Update main.py

main.py

@@ -12,13 +12,13 @@ UPLOAD_DIR = "uploads"
 if not os.path.exists(UPLOAD_DIR):
     os.makedirs(UPLOAD_DIR)
 
-# Store file mappings: {unique_code: filename}
-file_codes: Dict[str, str] = {}
-# Reverse mapping for easy lookup: {filename: unique_code}
-filename_codes: Dict[str, str] = {}
+# Store file mappings with user info: {unique_code: {"filename": filename, "user_id": user_id}}
+file_codes: Dict[str, dict] = {}
+# Reverse mapping: {filename: {"code": code, "user_id": user_id}}
+filename_codes: Dict[str, dict] = {}
 
 @app.post("/upload/")
-async def upload_file(file: UploadFile = File(...)):
+async def upload_file(file: UploadFile = File(...), user_id: str = None):
     try:
         # Generate unique code
         unique_code = str(uuid.uuid4())[:8]
@@ -28,9 +28,9 @@ async def upload_file(file: UploadFile = File(...)):
         with open(file_path, "wb") as buffer:
             shutil.copyfileobj(file.file, buffer)
         
-        # Store the mapping
-        file_codes[unique_code] = file.filename
-        filename_codes[file.filename] = unique_code
+        # Store the mapping with user info
+        file_codes[unique_code] = {"filename": file.filename, "user_id": user_id}
+        filename_codes[file.filename] = {"code": unique_code, "user_id": user_id}
         
         return {
             "filename": file.filename,
@@ -53,8 +53,8 @@ async def upload_multiple_files(files: List[UploadFile] = File(...)):
                 shutil.copyfileobj(file.file, buffer)
             
             # Store the mapping
-            file_codes[unique_code] = file.filename
-            filename_codes[file.filename] = unique_code
+            file_codes[unique_code] = {"filename": file.filename, "user_id": None}
+            filename_codes[file.filename] = {"code": unique_code, "user_id": None}
             
             uploaded_files.append({
                 "filename": file.filename,
@@ -64,16 +64,18 @@ async def upload_multiple_files(files: List[UploadFile] = File(...)):
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
-@app.get("/files/")
-async def list_files():
+@app.get("/files/{user_id}")
+async def list_files(user_id: str):
+    """List files for a specific user"""
     try:
         files = []
         for filename in os.listdir(UPLOAD_DIR):
-            code = filename_codes.get(filename, "no_code")
-            files.append({
-                "filename": filename,
-                "access_code": code
-            })
+            file_info = filename_codes.get(filename, {})
+            if file_info.get("user_id") == user_id:
+                files.append({
+                    "filename": filename,
+                    "access_code": file_info["code"]
+                })
         return {"files": files}
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
@@ -84,7 +86,7 @@ async def download_file(access_code: str):
         if access_code not in file_codes:
             raise HTTPException(status_code=404, detail="Invalid access code")
         
-        filename = file_codes[access_code]
+        filename = file_codes[access_code]["filename"]
         file_path = os.path.join(UPLOAD_DIR, filename)
         
         if os.path.exists(file_path):
@@ -100,7 +102,7 @@ async def delete_file(access_code: str):
         if access_code not in file_codes:
             raise HTTPException(status_code=404, detail="Invalid access code")
         
-        filename = file_codes[access_code]
+        filename = file_codes[access_code]["filename"]
         file_path = os.path.join(UPLOAD_DIR, filename)
         
         if os.path.exists(file_path):
@@ -122,7 +124,7 @@ async def direct_download(access_code: str):
             # If not a valid access code, return 404 or redirect to home
             raise HTTPException(status_code=404, detail="Invalid access code")
         
-        filename = file_codes[access_code]
+        filename = file_codes[access_code]["filename"]
         file_path = os.path.join(UPLOAD_DIR, filename)
         
         if os.path.exists(file_path):