Fix file extensions by preserving original filenames

app.py

@@ -1,5 +1,6 @@
 import asyncio
 import os
+import re
 import time
 import uuid
 import shutil
@@ -52,17 +53,7 @@ async def startup_event():
     asyncio.create_task(cleanup_old_files())
 
 async def download_file(task_id: str, url: str):
-    file_path = os.path.join(DATA_DIR, f"{task_id}.bin")
-    tasks[task_id] = {
-        "task_id": task_id,
-        "url": url,
-        "status": "downloading",
-        "total_size": 0,
-        "downloaded": 0,
-        "speed": 0.0,
-        "file_path": file_path,
-        "timestamp": time.time()
-    }
+    file_path = tasks[task_id]["file_path"]
     
     try:
         async with aiohttp.ClientSession() as session:
@@ -71,6 +62,20 @@ async def download_file(task_id: str, url: str):
                 total_size = int(response.headers.get('Content-Length', 0))
                 tasks[task_id]["total_size"] = total_size
                 
+                # Check for Content-Disposition header for filename
+                cd = response.headers.get('Content-Disposition')
+                if cd and 'filename=' in cd:
+                    # Extract filename from header
+                    fname = re.findall('filename="([^"]+)"', cd)
+                    if not fname:
+                        fname = re.findall('filename=([^;]+)', cd)
+                    if fname:
+                        new_filename = fname[0]
+                        new_file_path = os.path.join(DATA_DIR, f"{task_id}_{new_filename}")
+                        tasks[task_id]["file_path"] = new_file_path
+                        tasks[task_id]["original_filename"] = new_filename
+                        file_path = new_file_path
+
                 downloaded = 0
                 start_time = time.time()
                 last_time = start_time
@@ -100,6 +105,26 @@ async def download_file(task_id: str, url: str):
 @app.post("/start_download")
 async def start_download(req: DownloadRequest, background_tasks: BackgroundTasks):
     task_id = str(uuid.uuid4())
+    
+    # Try to extract filename from URL
+    filename = req.url.split("/")[-1]
+    if "?" in filename:
+        filename = filename.split("?")[0]
+    if not filename or "." not in filename:
+        filename = "downloaded_file.bin"
+        
+    tasks[task_id] = {
+        "task_id": task_id,
+        "url": req.url,
+        "status": "downloading",
+        "total_size": 0,
+        "downloaded": 0,
+        "speed": 0.0,
+        "file_path": os.path.join(DATA_DIR, f"{task_id}_{filename}"),
+        "original_filename": filename,
+        "timestamp": time.time()
+    }
+    
     background_tasks.add_task(download_file, task_id, req.url)
     return {"task_id": task_id}
 
@@ -113,7 +138,11 @@ async def get_status(task_id: str):
 async def download(task_id: str):
     if task_id not in tasks or tasks[task_id]["status"] != "completed":
         return {"error": "File not ready"}
-    return FileResponse(tasks[task_id]["file_path"], filename=f"downloaded_{task_id}.bin")
+    
+    file_path = tasks[task_id]["file_path"]
+    filename = tasks[task_id].get("original_filename", f"downloaded_{task_id}.bin")
+    
+    return FileResponse(file_path, filename=filename)
 
 @app.get("/stream/{task_id}")
 async def stream(task_id: str, request: Request):
@@ -182,7 +211,8 @@ async def get_history():
 @app.post("/upload")
 async def upload_file(file: UploadFile = File(...)):
     task_id = str(uuid.uuid4())
-    file_path = os.path.join(DATA_DIR, f"{task_id}.bin")
+    original_filename = file.filename if file.filename else "uploaded_file.bin"
+    file_path = os.path.join(DATA_DIR, f"{task_id}_{original_filename}")
     
     with open(file_path, "wb") as buffer:
         shutil.copyfileobj(file.file, buffer)
@@ -197,6 +227,7 @@ async def upload_file(file: UploadFile = File(...)):
         "downloaded": file_size,
         "speed": 0.0,
         "file_path": file_path,
+        "original_filename": original_filename,
         "timestamp": time.time()
     }
     return {"task_id": task_id}