Update app.py

app.py

@@ -1,91 +1,59 @@
-from fastapi import FastAPI, File, UploadFile, HTTPException, Response
+from fastapi import FastAPI, File, UploadFile, HTTPException, Form
 import uuid
 import mimetypes
-import time
-import threading
-import schedule
+from datetime import datetime, timedelta
 
 app = FastAPI()
 
 # **In-Memory Storage (FOR DEMO ONLY)**
 # Replace with Redis, NoSQL DB, or Cloud Storage for Production
 data_storage = {}
-expiration_timeout = 5 * 60  # 5 minutes in seconds
 
-def clean_expired_data():
-    global data_storage
-    current_time = time.time()
-    data_storage = {key: value for key, value in data_storage.items() if current_time - value['uploaded_at'] < expiration_timeout}
-    print("Data storage cleaned.")
-
-# Schedule data cleaning every minute
-schedule.every(1).minutes.do(clean_expired_data) 
-
-def run_schedule():
-    while True:
-        schedule.run_pending()
-        time.sleep(1)
-
-# Start the scheduler in a separate thread
-thread = threading.Thread(target=run_schedule)
-thread.daemon = True  # So that it dies when main thread dies
-thread.start()
-
-@app.post("/upload/file/auto-link")
-async def upload_file_auto_link(file: UploadFile = File(...)):
-    upload_id = str(uuid.uuid4())
-    file_name = file.filename
-    mime_type = file.content_type
-    await _store_file(upload_id, file_name, await file.read(), mime_type)
-    return {"download_link": f"/download/auto/{upload_id}", "message": "File uploaded successfully"}
-
-@app.post("/upload/file/custom-link")
-async def upload_file_custom_link(link: str, file: UploadFile = File(...)):
-    # Ensure link doesn't collide and is URL-friendly (simplified for demo)
-    if link in [item['custom_link'] for item in data_storage.values() if 'custom_link' in item]:
-        link += f"_{str(uuid.uuid4())[:5]}"  # Append unique suffix
-        print(f"Link collision detected. New link: {link}")
-    
+@app.post("/upload/file")
+async def upload_file(file: UploadFile = File(...)):
     upload_id = str(uuid.uuid4())
-    file_name = file.filename
-    mime_type = file.content_type
-    await _store_file(upload_id, file_name, await file.read(), mime_type, custom_link=link)
-    return {"download_link": f"/download/custom/{link}", "message": "File uploaded successfully"}
-
-@app.get("/download/auto/{upload_id}")
-async def download_file_auto(upload_id: str):
-    return await _handle_download(upload_id)
-
-@app.get("/download/custom/{link}")
-async def download_file_custom(link: str):
-    for item in data_storage.values():
-        if 'custom_link' in item and item['custom_link'] == link:
-            return await _handle_download(item['upload_id'], item['file_name'], item['file_data'], item['mime_type'])
-    raise HTTPException(status_code=404, detail="File not found")
-
-async def _store_file(upload_id, file_name, file_data, mime_type, custom_link=None):
-    global data_storage
     data_storage[upload_id] = {
-        'upload_id': upload_id,
-        'file_name': file_name,
-        'file_data': file_data,
-        'ime_type': mime_type,
-        'uploaded_at': time.time(),
-        'custom_link': custom_link  # Optional
+        'file_name': file.filename,
+        'file_data': await file.read(),
+       'mime_type': file.content_type,
+        'upload_time': datetime.datetime.utcnow()
     }
+    return {"upload_id": upload_id, "message": "File uploaded successfully"}
 
-async def _handle_download(upload_id, file_name=None, file_data=None, mime_type=None):
-    global data_storage
+@app.post("/upload/custom")
+async def upload_custom(file: UploadFile = File(...), download_link: str = Form(...)):
+    if not download_link:
+        raise HTTPException(status_code=400, detail="Download link is required")
+    
+    # Validate the download_link to be URL-safe
+    if not download_link.isalnum() and not all(c in string.ascii_letters + string.digits + '-_' for c in download_link):
+        raise HTTPException(status_code=400, detail="Invalid download link format")
+    
+    # Assign the custom download link
+    data_storage[download_link] = {
+        'file_name': file.filename,
+        'file_data': await file.read(),
+       'mime_type': file.content_type,
+        'upload_time': datetime.datetime.utcnow()
+    }
+    return {"download_link": download_link, "message": "File uploaded successfully"}
+
+@app.get("/download/{upload_id}")
+async def download_file(upload_id: str):
+    # Remove expired files
+    expire_time = datetime.datetime.utcnow() - timedelta(minutes=5)
+    keys_to_remove = [key for key, value in data_storage.items() if value['upload_time'] < expire_time]
+    for key in keys_to_remove:
+        del data_storage[key]
+    
     if upload_id not in data_storage:
-        raise HTTPException(status_code=404, detail="File not found or has expired")
+        raise HTTPException(status_code=404, detail="Upload not found")
     
-    if not file_name or not file_data or not mime_type:
-        upload_data = data_storage[upload_id]
-        file_name = upload_data['file_name']
-        file_data = upload_data['file_data']
-        mime_type = upload_data['mime_type'] or mimetypes.guess_type(file_name)[0] or "application/octet-stream"
+    upload_data = data_storage[upload_id]
+    file_name = upload_data['file_name']
+    mime_type = upload_data['mime_type'] or mimetypes.guess_type(file_name)[0] or "application/octet-stream"
     
-    return Response(content=file_data, media_type=mime_type, 
+    return Response(content=upload_data['file_data'], media_type=mime_type, 
                     headers={"Content-Disposition": f"attachment; filename={file_name}"})
 
 if __name__ == "__main__":