Update app.py

app.py

@@ -5,6 +5,8 @@ import struct
 import logging
 from flask import Flask, render_template, request, send_file, jsonify
 from werkzeug.utils import secure_filename
+from pydub import AudioSegment
+import magic
 
 app = Flask(__name__)
 
@@ -40,7 +42,7 @@ def process_file():
         logger.info(f"File saved: {filepath}")
 
         factor = int(request.form['factor'])
-        is_decrypting = filename.lower().endswith('.wav')
+        is_decrypting = filename.lower().endswith('.mp3')
         logger.info(f"Processing file: {'decrypting' if is_decrypting else 'encrypting'} with factor {factor}")
 
         try:
@@ -76,27 +78,45 @@ def encrypt_file(filepath, factor):
     samples = [ord(byte) * factor for byte in data.decode('latin-1')]
     
     # Create WAV file
-    output = io.BytesIO()
-    with wave.open(output, 'wb') as wav_output:
-        wav_output.setnchannels(1)
-        wav_output.setsampwidth(2)
-        wav_output.setframerate(44100)
-        wav_output.writeframes(struct.pack(f'{len(samples)}h', *samples))
+    wav_output = io.BytesIO()
+    with wave.open(wav_output, 'wb') as wav_file:
+        wav_file.setnchannels(1)
+        wav_file.setsampwidth(2)
+        wav_file.setframerate(44100)
+        wav_file.writeframes(struct.pack(f'{len(samples)}h', *samples))
     
-    logger.info(f"Encryption complete. Output size: {output.getbuffer().nbytes} bytes")
-    return output.getvalue(), 'encrypted.wav', 'audio/wav'
+    # Convert WAV to MP3
+    wav_output.seek(0)
+    audio = AudioSegment.from_wav(wav_output)
+    mp3_output = io.BytesIO()
+    audio.export(mp3_output, format="mp3")
+    
+    logger.info(f"Encryption complete. Output size: {mp3_output.getbuffer().nbytes} bytes")
+    return mp3_output.getvalue(), 'encrypted.mp3', 'audio/mpeg'
 
 def decrypt_file(filepath, factor):
     logger.info(f"Decrypting file: {filepath}")
-    with wave.open(filepath, 'rb') as wav_file:
+    
+    # Convert MP3 to WAV
+    audio = AudioSegment.from_mp3(filepath)
+    wav_output = io.BytesIO()
+    audio.export(wav_output, format="wav")
+    wav_output.seek(0)
+    
+    with wave.open(wav_output, 'rb') as wav_file:
         params = wav_file.getparams()
         frames = wav_file.readframes(wav_file.getnframes())
     
     samples = struct.unpack(f'{len(frames)//2}h', frames)
     decrypted_data = ''.join(chr(int(sample / factor)) for sample in samples).encode('latin-1')
     
+    # Detect original file type
+    mime = magic.Magic(mime=True)
+    file_type = mime.from_buffer(decrypted_data)
+    file_extension = file_type.split('/')[-1]
+    
     logger.info(f"Decryption complete. Output size: {len(decrypted_data)} bytes")
-    return decrypted_data, 'decrypted_file', 'application/octet-stream'
+    return decrypted_data, f'decrypted_file.{file_extension}', file_type
 
 if __name__ == '__main__':
     app.run(debug=True)