Verify download filenames in URL

3 files changed, 9 insertions, 6 deletions

diff --git a/encrypt.py b/encrypt.py
index 80cedf3..e813d36 100644
--- a/encrypt.py
+++ b/encrypt.py
@@ -14,6 +14,7 @@ if __name__ == '__main__':
         print(f'{infile} is not a file or directory, exiting.')
         os.exit(2)
 
-    file_id, token = encrypt_file(args.infile)
-    print(f'/{file_id}/{token}/{os.path.basename(args.infile)}')
+    download_filename = os.path.basename(args.infile)
+    file_id, token = encrypt_file(args.infile, download_filename)
+    print(f'/{file_id}/{token}/{download_filename}')
     
diff --git a/filecrypt.py b/filecrypt.py
index 85d9eea..da0db9b 100644
--- a/filecrypt.py
+++ b/filecrypt.py
@@ -10,13 +10,14 @@ FILE_ID_LENGTH = 22
 TOKEN_LENGTH = 22
 HEADER_LENGTH = 56
 
-def encrypt_file(filename_in, chunk_size=1000000//16):
+def encrypt_file(filename_in, download_filename, chunk_size=1000000//16):
     file_id = secrets.token_urlsafe(16)
     auth_secret = secrets.token_bytes(16)
     key = secrets.token_bytes(16)
     data_nonce = secrets.token_bytes(8)
 
     token_cipher = AES.new(auth_secret, AES.MODE_GCM)
+    token_cipher.update(download_filename.encode())
     ciphertext, token_tag = token_cipher.encrypt_and_digest(key)
     token = base64.b64encode(ciphertext).rstrip(b'=').decode()
 
@@ -41,7 +42,7 @@ def encrypt_file(filename_in, chunk_size=1000000//16):
 def payload_size(path):
     return os.stat(path).st_size - HEADER_LENGTH
 
-def decrypt_generator(filename, token, seek=0, end=None, chunk_size=1000000//16):
+def decrypt_generator(filename, download_filename, token, seek=0, end=None, chunk_size=1000000//16):
     with open(filename, 'rb') as fin:
         token_nonce = fin.read(16)
         token_tag = fin.read(16)
@@ -51,6 +52,7 @@ def decrypt_generator(filename, token, seek=0, end=None, chunk_size=1000000//16)
 
     ciphertext = base64.b64decode(token + '='*((3-len(token)%3)%3))
     token_cipher = AES.new(auth_secret, AES.MODE_GCM, nonce=token_nonce)
+    token_cipher.update(download_filename.encode())
     key = token_cipher.decrypt_and_verify(ciphertext, token_tag)
 
     def generator():
diff --git a/server.py b/server.py
index b084298..694b1ac 100644
--- a/server.py
+++ b/server.py
@@ -26,7 +26,7 @@ def download(file_id, token, filename):
     range_header = re.match('^bytes=([0-9]+)-([0-9]*)$', request.headers.get('Range', ''))
     if not range_header:
         try:
-            generator = filecrypt.decrypt_generator(path, token)
+            generator = filecrypt.decrypt_generator(path, filename, token)
         except ValueError: # MAC check failed
             abort(403) # forbidden
 
@@ -40,7 +40,7 @@ def download(file_id, token, filename):
             abort(416) # range not satisfiable
 
         try:
-            generator = filecrypt.decrypt_generator(path, token, seek=range_start, end=range_end)
+            generator = filecrypt.decrypt_generator(path, filename, token, seek=range_start, end=range_end)
         except ValueError: # MAC check failed
             abort(403) # forbidden
         response = Response(generator, status=206, mimetype='application/octet-stream')