Update app.py

app.py

@@ -3,6 +3,7 @@ import io
 import wave
 import struct
 import logging
+import hashlib
 from flask import Flask, render_template, request, send_file, jsonify
 from werkzeug.utils import secure_filename
 from pydub import AudioSegment
@@ -14,7 +15,7 @@ logging.basicConfig(level=logging.DEBUG)
 logger = app.logger
 
 UPLOAD_FOLDER = '/tmp'
-ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif', 'mp3', 'wav', 'txt', 'pdf'}
+ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif', 'mp3', 'wav', 'txt', 'pdf', 'docx'}
 
 app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
 
@@ -69,14 +70,24 @@ def process_file():
     logger.error("Invalid file type")
     return jsonify({'error': 'Invalid file type'}), 400
 
+def calculate_checksum(data):
+    return hashlib.md5(data).hexdigest()
+
 def encrypt_file(filepath, factor):
     logger.info(f"Encrypting file: {filepath}")
     with open(filepath, 'rb') as file:
         data = file.read()
     
+    # Calculate checksum
+    checksum = calculate_checksum(data)
+    
     # Convert data to audio samples
     samples = [ord(byte) * factor for byte in data.decode('latin-1')]
     
+    # Add checksum to the beginning of the samples
+    checksum_samples = [ord(byte) for byte in checksum.encode('ascii')]
+    samples = checksum_samples + samples
+    
     # Create WAV file
     wav_output = io.BytesIO()
     with wave.open(wav_output, 'wb') as wav_file:
@@ -109,15 +120,22 @@ def decrypt_file(filepath, factor):
     
     samples = struct.unpack(f'{len(frames)//2}h', frames)
     
-    # Use a list comprehension with error handling
+    # Extract checksum
+    checksum = ''.join(chr(sample // factor) for sample in samples[:32])
+    samples = samples[32:]
+    
+    # Decrypt data
     decrypted_bytes = bytearray()
     for sample in samples:
         byte_value = int(round(sample / factor))
-        # Ensure the byte value is within the valid range (0-255)
         decrypted_bytes.append(max(0, min(byte_value, 255)))
     
     decrypted_data = bytes(decrypted_bytes)
     
+    # Verify checksum
+    if calculate_checksum(decrypted_data) != checksum:
+        raise ValueError("Checksum verification failed. The file may be corrupted or the wrong factor was used.")
+    
     # Detect original file type
     mime = magic.Magic(mime=True)
     file_type = mime.from_buffer(decrypted_data)