1 files changed, 55 insertions, 34 deletions

diff --git a/test/test.py b/test/test.py
index 02216ac..2d09662 100644
--- a/test/test.py
+++ b/test/test.py
@@ -6,12 +6,13 @@ import io
 import os
 
 from unittest.mock import patch
+from flask_testing import TestCase
 
 import main
 
 
-class Mat2WebTestCase(unittest.TestCase):
-    def setUp(self):
+class Mat2WebTestCase(TestCase):
+    def create_app(self):
         os.environ.setdefault('MAT2_ALLOW_ORIGIN_WHITELIST', 'origin1.gnu origin2.gnu')
         self.upload_folder = tempfile.mkdtemp()
         app = main.create_app(
@@ -20,45 +21,45 @@ class Mat2WebTestCase(unittest.TestCase):
                 'UPLOAD_FOLDER': self.upload_folder
             }
         )
-        self.app = app.test_client()
+        return app
 
     def tearDown(self):
         shutil.rmtree(self.upload_folder)
 
     def test_get_root(self):
-        rv = self.app.get('/')
+        rv = self.client.get('/')
         self.assertIn(b'mat2-web', rv.data)
 
     def test_check_mimetypes(self):
-        rv = self.app.get('/')
+        rv = self.client.get('/')
         self.assertIn(b'.torrent', rv.data)
         self.assertIn(b'.ods', rv.data)
 
     def test_get_download_dangerous_file(self):
-        rv = self.app.get('/download/1337/\..\filename')
+        rv = self.client.get('/download/1337/aabb/\..\filename')
         self.assertEqual(rv.status_code, 302)
 
     def test_get_download_without_key_file(self):
-        rv = self.app.get('/download/non_existant')
+        rv = self.client.get('/download/non_existant')
         self.assertEqual(rv.status_code, 404)
 
     def test_get_download_nonexistant_file(self):
-        rv = self.app.get('/download/1337/non_existant')
+        rv = self.client.get('/download/1337/aabb/non_existant')
         self.assertEqual(rv.status_code, 302)
 
     def test_get_upload_without_file(self):
-        rv = self.app.post('/')
+        rv = self.client.post('/')
         self.assertEqual(rv.status_code, 302)
 
     def test_get_upload_empty_file(self):
-        rv = self.app.post('/',
+        rv = self.client.post('/',
                 data=dict(
                     file=(io.BytesIO(b""), 'test.pdf'),
                     ), follow_redirects=False)
         self.assertEqual(rv.status_code, 302)
 
     def test_get_upload_empty_file_redir(self):
-        rv = self.app.post('/',
+        rv = self.client.post('/',
                 data=dict(
                     file=(io.BytesIO(b""), 'test.pdf'),
                     ), follow_redirects=True)
@@ -67,7 +68,7 @@ class Mat2WebTestCase(unittest.TestCase):
         self.assertEqual(rv.status_code, 200)
 
     def test_get_upload_no_selected_file(self):
-        rv = self.app.post('/',
+        rv = self.client.post('/',
                            data=dict(
                                file=(io.BytesIO(b""), ''),
                            ), follow_redirects=True)
@@ -86,7 +87,7 @@ class Mat2WebTestCase(unittest.TestCase):
             'AAAAAAAAAAApIFnAAAAdGVzdC5qc29uVVQNAAfomo9d6JqPXeiaj111eAsAAQTpAwAABOkDAAB'
             'QSwUGAAAAAAIAAgC8AAAAwAAAAAAA'
         )
-        rv = self.app.post('/',
+        rv = self.client.post('/',
                            data=dict(
                                file=(io.BytesIO(zip_file_bytes), 'test.zip'),
                            ), follow_redirects=True)
@@ -94,7 +95,7 @@ class Mat2WebTestCase(unittest.TestCase):
         self.assertEqual(rv.status_code, 200)
 
     def test_get_upload_no_file_name(self):
-        rv = self.app.post('/',
+        rv = self.client.post('/',
                 data=dict(
                     file=(io.BytesIO(b"aaa")),
                     ), follow_redirects=True)
@@ -102,30 +103,51 @@ class Mat2WebTestCase(unittest.TestCase):
         self.assertEqual(rv.status_code, 200)
 
     def test_get_upload_harmless_file(self):
-        rv = self.app.post('/',
-                data=dict(
-                    file=(io.BytesIO(b"Some text"), 'test.txt'),
-                    ), follow_redirects=True)
-        self.assertIn(b'/download/4c2e9e6da31a64c70623619c449a040968cdbea85945bf384fa30ed2d5d24fa3/test.cleaned.txt', rv.data)
+        rv = self.client.post(
+            '/',
+            data=dict(
+                file=(io.BytesIO(b"Some text"), 'test.txt'),
+            ),
+            follow_redirects=True
+        )
+        download_uri = self.get_context_variable('download_uri')
+        self.assertIn('/test.cleaned.txt', download_uri)
         self.assertEqual(rv.status_code, 200)
         self.assertNotIn('Access-Control-Allow-Origin', rv.headers)
 
-        rv = self.app.get('/download/4c2e9e6da31a64c70623619c449a040968cdbea85945bf384fa30ed2d5d24fa3/test.cleaned.txt')
+        rv = self.client.get(download_uri)
         self.assertEqual(rv.status_code, 200)
 
-        rv = self.app.get('/download/4c2e9e6da31a64c70623619c449a040968cdbea85945bf384fa30ed2d5d24fa3/test.cleaned.txt')
+        rv = self.client.get(download_uri)
         self.assertEqual(rv.status_code, 302)
 
-    def test_upload_wrong_hash(self):
-        rv = self.app.post('/',
-                           data=dict(
-                               file=(io.BytesIO(b"Some text"), 'test.txt'),
-                           ), follow_redirects=True)
-        self.assertIn(b'/download/4c2e9e6da31a64c70623619c449a040968cdbea85945bf384fa30ed2d5d24fa3/test.cleaned.txt',
-                      rv.data)
+    def test_upload_wrong_hash_or_secret(self):
+        rv = self.client.post(
+            '/',
+            data=dict(
+                file=(io.BytesIO(b"Some text"), 'test.txt'),
+            ),
+            follow_redirects=True
+        )
+
+        download_uri = self.get_context_variable('download_uri')
+
+        self.assertIn('/test.cleaned.txt', download_uri)
+        self.assertIn('/download', download_uri)
         self.assertEqual(rv.status_code, 200)
 
-        rv = self.app.get('/download/70623619c449a040968cdbea85945bf384fa30ed2d5d24fa3/test.cleaned.txt')
+        uri_parts = download_uri.split("/")
+        self.assertEqual(len(uri_parts[2]), len(uri_parts[3]))
+        self.assertEqual(64, len(uri_parts[2]))
+
+        key_uri_parts = uri_parts
+        key_uri_parts[2] = '70623619c'
+        rv = self.client.get("/".join(key_uri_parts))
+        self.assertEqual(rv.status_code, 302)
+
+        key_uri_parts = uri_parts
+        key_uri_parts[3] = '70623619c'
+        rv = self.client.get("/".join(key_uri_parts))
         self.assertEqual(rv.status_code, 302)
 
     @patch('matweb.file_removal_scheduler.random.randint')
@@ -140,19 +162,18 @@ class Mat2WebTestCase(unittest.TestCase):
         )
         app = app.test_client()
 
-        request = self.app.post('/',
+        request = self.client.post('/',
                            data=dict(
                                file=(io.BytesIO(b"Some text"), 'test.txt'),
                            ), follow_redirects=True)
         self.assertEqual(request.status_code, 200)
-        request = app.get(
-            b'/download/4c2e9e6da31a64c70623619c449a040968cdbea85945bf384fa30ed2d5d24fa3/test.cleaned.txt'
-        )
+
+        request = app.get(self.get_context_variable('download_uri'))
         self.assertEqual(302, request.status_code)
         os.environ['MAT2_MAX_FILE_AGE_FOR_REMOVAL'] = '9999'
 
     def test_info_page(self):
-        rv = self.app.get('/info')
+        rv = self.client.get('/info')
         self.assertIn(b'What are metadata?', rv.data)
         self.assertIn(b'.asc', rv.data)
         self.assertIn(b'.mp2', rv.data)