5 files changed, 46 insertions, 5 deletions

diff --git a/matweb/frontend.py b/matweb/frontend.py
index 2e25467..8295f4e 100644
--- a/matweb/frontend.py
+++ b/matweb/frontend.py
@@ -53,8 +53,12 @@ def upload_file():
         if not uploaded_file.filename:
             flash('No selected file')
             return redirect(request.url)
+        try:
+            filename, filepath = utils.save_file(uploaded_file, current_app.config['UPLOAD_FOLDER'])
+        except ValueError:
+            flash('Invalid Filename')
+            return redirect(request.url)
 
-        filename, filepath = utils.save_file(uploaded_file, current_app.config['UPLOAD_FOLDER'])
         parser, mime = utils.get_file_parser(filepath)
 
         if parser is None:
diff --git a/matweb/rest_api.py b/matweb/rest_api.py
index 4098050..a07d2d2 100644
--- a/matweb/rest_api.py
+++ b/matweb/rest_api.py
@@ -28,11 +28,15 @@ class APIUpload(Resource):
         args = req_parser.parse_args()
         try:
             file_data = base64.b64decode(args['file'])
-        except binascii.Error as err:
-            abort(400, message='Failed decoding file: ' + str(err))
+        except (binascii.Error, ValueError):
+            abort(400, message='Failed decoding file')
 
         file = FileStorage(stream=io.BytesIO(file_data), filename=args['file_name'])
-        filename, filepath = utils.save_file(file, self.upload_folder)
+        try:
+            filename, filepath = utils.save_file(file, self.upload_folder)
+        except ValueError:
+            abort(400, message='Invalid Filename')
+
         parser, mime = utils.get_file_parser(filepath)
 
         if parser is None:
diff --git a/matweb/utils.py b/matweb/utils.py
index ec9b99c..20c213d 100644
--- a/matweb/utils.py
+++ b/matweb/utils.py
@@ -65,6 +65,8 @@ def get_supported_extensions():
 
 def save_file(file, upload_folder):
     filename = secure_filename(file.filename)
+    if not filename:
+        raise ValueError('Invalid Filename')
     filepath = os.path.join(upload_folder, filename)
     file.save(os.path.join(filepath))
     return filename, filepath
diff --git a/test/test.py b/test/test.py
index 2d09662..7431881 100644
--- a/test/test.py
+++ b/test/test.py
@@ -179,6 +179,18 @@ class Mat2WebTestCase(TestCase):
         self.assertIn(b'.mp2', rv.data)
         self.assertEqual(rv.status_code, 200)
 
+    def test_get_upload_naughty_input(self):
+        rv = self.client.post(
+            '/',
+            data=dict(
+                file=(io.BytesIO(b"a"), '﷽'),
+            ),
+            follow_redirects=True
+        )
+        self.assertEqual(rv.status_code, 200)
+        self.assertIn(b'Invalid Filename', rv.data)
+
+
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/test/test_api.py b/test/test_api.py
index 4925d9e..af736af 100644
--- a/test/test_api.py
+++ b/test/test_api.py
@@ -70,7 +70,7 @@ class Mat2APITestCase(unittest.TestCase):
 
         self.assertEqual(request.status_code, 400)
         error = request.get_json()['message']
-        self.assertEqual(error, 'Failed decoding file: Incorrect padding')
+        self.assertEqual(error, 'Failed decoding file')
 
     def test_api_not_supported(self):
         request = self.app.post('/api/upload',
@@ -400,6 +400,25 @@ class Mat2APITestCase(unittest.TestCase):
         request = app.get(download_link)
         self.assertEqual(code, request.status_code)
 
+    def test_upload_naughty_input(self):
+        request = self.app.post('/api/upload',
+                           data='{"file_name": "\\\\", '
+                                '"file": "\\\\"}',
+                           headers={'content-type': 'application/json'}
+                           )
+        error_message = request.get_json()['message']
+        self.assertEqual(400, request.status_code)
+        self.assertEqual("Invalid Filename", error_message)
+
+        request = self.app.post('/api/upload',
+                                data='{"file_name": "﷽", '
+                                     '"file": "﷽"}',
+                                headers={'content-type': 'application/json'}
+                                )
+        error_message = request.get_json()['message']
+        self.assertEqual(400, request.status_code)
+        self.assertEqual("Failed decoding file", error_message)
+
 
 if __name__ == '__main__':
     unittest.main()