Update app.py

app.py

@@ -1,19 +1,16 @@
 from fastapi import FastAPI, File, UploadFile, HTTPException, Response
-from pydantic import BaseModel
-from typing import Optional
 import uuid
-import os
-import schedule
+import mimetypes
 import time
 import threading
-import mimetypes
+import schedule
 
 app = FastAPI()
 
-# In-memory storage (NOT SUITABLE FOR PRODUCTION)
-# Consider Redis or an embedded NoSQL DB for production
+# **In-Memory Storage (FOR DEMO ONLY)**
+# Replace with Redis, NoSQL DB, or Cloud Storage for Production
 data_storage = {}
-expiration_timeout = 15 * 60  # 15 minutes in seconds
+expiration_timeout = 5 * 60  # 5 minutes in seconds
 
 def clean_expired_data():
     global data_storage
@@ -21,8 +18,8 @@ def clean_expired_data():
     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 15 minutes
-schedule.every(15).minutes.do(clean_expired_data) 
+# Schedule data cleaning every minute
+schedule.every(1).minutes.do(clean_expired_data) 
 
 def run_schedule():
     while True:
@@ -34,62 +31,63 @@ thread = threading.Thread(target=run_schedule)
 thread.daemon = True  # So that it dies when main thread dies
 thread.start()
 
-class TextUpload(BaseModel):
-    text: str
-    file_name: Optional[str] = "uploaded_text.txt"
+@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")
-async def upload_file(file: UploadFile = File(...)):
-    global data_storage
+@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}")
+    
     upload_id = str(uuid.uuid4())
-    data_storage[upload_id] = {
-        'data_type': "file",
-        'file_name': file.filename,
-        'file_data': await file.read(),
-        'uploaded_at': time.time()
-    }
-    # Optionally save file to filesystem (commented out for brevity)
-    # with open(f"uploads/{file.filename}", "wb") as f:
-    #     f.write(await file.read())
-    return {"upload_id": upload_id, "message": "File uploaded successfully"}
+    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")
 
-@app.post("/upload/text")
-async def upload_text(text_upload: TextUpload):
+async def _store_file(upload_id, file_name, file_data, mime_type, custom_link=None):
     global data_storage
-    upload_id = str(uuid.uuid4())
     data_storage[upload_id] = {
-        'data_type': "text",
-        'file_name': text_upload.file_name,
-        'file_data': text_upload.text.encode('utf-8'),
-        'uploaded_at': time.time()
+        '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
     }
-    return {"upload_id": upload_id, "message": "Text uploaded successfully"}
 
-@app.get("/download/{upload_id}")
-async def download_file(upload_id: str):
+async def _handle_download(upload_id, file_name=None, file_data=None, mime_type=None):
     global data_storage
     if upload_id not in data_storage:
-        raise HTTPException(status_code=404, detail="Upload not found or has expired")
-    
-    upload_data = data_storage[upload_id]
-    file_name = upload_data['file_name']
-    
-    # Guess the MIME type based on the file extension
-    mime_type, _ = mimetypes.guess_type(file_name)
-    if mime_type is None:  # Default to octet-stream if type is unknown
-        mime_type = "application/octet-stream"
+        raise HTTPException(status_code=404, detail="File not found or has expired")
     
-    # For text uploads, ensure the MIME type is set correctly
-    if upload_data['data_type'] == "text":
-        mime_type = "text/plain"
-        # Optionally, force a `.txt` extension for text files if no extension is provided
-        if '.' not in file_name:
-            file_name += '.txt'
+    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"
     
-    return Response(content=upload_data['file_data'], media_type=mime_type, 
+    return Response(content=file_data, media_type=mime_type, 
                     headers={"Content-Disposition": f"attachment; filename={file_name}"})
 
 if __name__ == "__main__":
-    scheduled_function()
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=8083)