Resolve "Use a HMAC instead of a hash"

3 files changed, 47 insertions, 21 deletions

diff --git a/matweb/frontend.py b/matweb/frontend.py
index 93432b4..2e25467 100644
--- a/matweb/frontend.py
+++ b/matweb/frontend.py
@@ -18,8 +18,8 @@ def info():
     )
 
 
-@routes.route('/download/<string:key>/<string:filename>')
-def download_file(key: str, filename: str):
+@routes.route('/download/<string:key>/<string:secret>/<string:filename>')
+def download_file(key: str, secret: str, filename: str):
     if filename != secure_filename(filename):
         return redirect(url_for('routes.upload_file'))
 
@@ -28,7 +28,7 @@ def download_file(key: str, filename: str):
 
     if not os.path.exists(complete_path):
         return redirect(url_for('routes.upload_file'))
-    if hmac.compare_digest(utils.hash_file(complete_path), key) is False:
+    if hmac.compare_digest(utils.hash_file(complete_path, secret), key) is False:
         return redirect(url_for('routes.upload_file'))
 
     @after_this_request
@@ -67,10 +67,14 @@ def upload_file():
             flash('Unable to clean %s' % mime)
             return redirect(url_for('routes.upload_file'))
 
-        key, meta_after, output_filename = utils.cleanup(parser, filepath, current_app.config['UPLOAD_FOLDER'])
+        key, secret, meta_after, output_filename = utils.cleanup(parser, filepath, current_app.config['UPLOAD_FOLDER'])
 
         return render_template(
-            'download.html', mimetypes=mime_types, meta=meta, filename=output_filename, meta_after=meta_after, key=key
+            'download.html',
+            mimetypes=mime_types,
+            meta=meta,
+            download_uri=url_for('routes.download_file', key=key, secret=secret, filename=output_filename),
+            meta_after=meta_after,
         )
 
     max_file_size = int(current_app.config['MAX_CONTENT_LENGTH'] / 1024 / 1024)
diff --git a/matweb/rest_api.py b/matweb/rest_api.py
index 60d834f..4098050 100644
--- a/matweb/rest_api.py
+++ b/matweb/rest_api.py
@@ -42,14 +42,15 @@ class APIUpload(Resource):
         if not parser.remove_all():
             abort(500, message='Unable to clean %s' % mime)
 
-        key, meta_after, output_filename = utils.cleanup(parser, filepath, self.upload_folder)
+        key, secret, meta_after, output_filename = utils.cleanup(parser, filepath, self.upload_folder)
         return utils.return_file_created_response(
             output_filename,
             mime,
             key,
+            secret,
             meta,
             meta_after,
-            urljoin(request.host_url, '%s/%s/%s/%s' % ('api', 'download', key, output_filename))
+            urljoin(request.host_url, '%s/%s/%s/%s/%s' % ('api', 'download', key, secret, output_filename))
         )
 
 
@@ -58,8 +59,8 @@ class APIDownload(Resource):
     def __init__(self, **kwargs):
         self.upload_folder = kwargs['upload_folder']
 
-    def get(self, key: str, filename: str):
-        complete_path, filepath = utils.is_valid_api_download_file(filename, key, self.upload_folder)
+    def get(self, key: str, secret: str, filename: str):
+        complete_path, filepath = utils.is_valid_api_download_file(filename, key, secret, self.upload_folder)
         # Make sure the file is NOT deleted on HEAD requests
         if request.method == 'GET':
             file_removal_scheduler.run_file_removal_job(self.upload_folder)
@@ -87,6 +88,7 @@ class APIBulkDownloadCreator(Resource):
                 'type': 'dict',
                 'schema': {
                     'key': {'type': 'string', 'required': True},
+                    'secret': {'type': 'string', 'required': True},
                     'file_name': {'type': 'string', 'required': True}
                 }
             }
@@ -108,6 +110,7 @@ class APIBulkDownloadCreator(Resource):
                 complete_path, file_path = utils.is_valid_api_download_file(
                     file_candidate['file_name'],
                     file_candidate['key'],
+                    file_candidate['secret'],
                     self.upload_folder
                 )
                 try:
@@ -124,13 +127,17 @@ class APIBulkDownloadCreator(Resource):
         parser, mime = utils.get_file_parser(zip_path)
         if not parser.remove_all():
             abort(500, message='Unable to clean %s' % mime)
-        key, meta_after, output_filename = utils.cleanup(parser, zip_path, self.upload_folder)
+        key, secret, meta_after, output_filename = utils.cleanup(parser, zip_path, self.upload_folder)
         return {
                    'output_filename': output_filename,
                    'mime': mime,
                    'key': key,
+                   'secret': secret,
                    'meta_after': meta_after,
-                   'download_link': urljoin(request.host_url, '%s/%s/%s/%s' % ('api', 'download', key, output_filename))
+                   'download_link': urljoin(
+                       request.host_url,
+                       '%s/%s/%s/%s/%s' % ('api', 'download', key, secret, output_filename)
+                   )
                }, 201
 
 
diff --git a/matweb/utils.py b/matweb/utils.py
index 8dfff45..ec9b99c 100644
--- a/matweb/utils.py
+++ b/matweb/utils.py
@@ -12,15 +12,21 @@ def get_allow_origin_header_value():
     return os.environ.get('MAT2_ALLOW_ORIGIN_WHITELIST', '*').split(" ")
 
 
-def hash_file(filepath: str) -> str:
-    sha256 = hashlib.sha256()
+def hash_file(filepath: str, secret: str) -> str:
+    """
+    The goal of the hmac is to ONLY make the hashes unpredictable
+    :param filepath: Path of the file
+    :param secret: a server side generated secret
+    :return: digest, secret
+    """
+    mac = hmac.new(secret.encode(), None, hashlib.sha256)
     with open(filepath, 'rb') as f:
         while True:
             data = f.read(65536)  # read the file by chunk of 64k
             if not data:
                 break
-            sha256.update(data)
-    return sha256.hexdigest()
+            mac.update(data)
+    return mac.hexdigest()
 
 
 def check_upload_folder(upload_folder):
@@ -28,11 +34,20 @@ def check_upload_folder(upload_folder):
         os.mkdir(upload_folder)
 
 
-def return_file_created_response(output_filename, mime, key, meta, meta_after, download_link):
+def return_file_created_response(
+        output_filename: str,
+        mime: str,
+        key: str,
+        secret: str,
+        meta: list,
+        meta_after: list,
+        download_link: str
+) -> dict:
     return {
         'output_filename': output_filename,
         'mime': mime,
         'key': key,
+        'secret': secret,
         'meta': meta,
         'meta_after': meta_after,
         'download_link': download_link
@@ -65,9 +80,9 @@ def cleanup(parser, filepath, upload_folder):
     parser, _ = parser_factory.get_parser(parser.output_filename)
     meta_after = parser.get_meta()
     os.remove(filepath)
-
-    key = hash_file(os.path.join(upload_folder, output_filename))
-    return key, meta_after, output_filename
+    secret = os.urandom(32).hex()
+    key = hash_file(os.path.join(upload_folder, output_filename), secret)
+    return key, secret, meta_after, output_filename
 
 
 def get_file_paths(filename, upload_folder):
@@ -77,7 +92,7 @@ def get_file_paths(filename, upload_folder):
     return complete_path, filepath
 
 
-def is_valid_api_download_file(filename, key, upload_folder):
+def is_valid_api_download_file(filename: str, key: str, secret: str, upload_folder: str) -> [str, str]:
     if filename != secure_filename(filename):
         abort(400, message='Insecure filename')
 
@@ -86,6 +101,6 @@ def is_valid_api_download_file(filename, key, upload_folder):
     if not os.path.exists(complete_path):
         abort(404, message='File not found')
 
-    if hmac.compare_digest(hash_file(complete_path), key) is False:
+    if hmac.compare_digest(hash_file(complete_path, secret), key) is False:
         abort(400, message='The file hash does not match')
     return complete_path, filepath