import shutil import os import asyncio from fastapi import ( APIRouter, UploadFile, File, HTTPException, Header, status, Depends, WebSocket, WebSocketDisconnect ) # Import Auth System from app.routers.auth.system import get_current_user # Router Configuration router = APIRouter( prefix="/drive", tags=["Cloud Drive"] ) # Settings UPLOAD_DIR = "uploaded_files" MAX_FILE_SIZE = 1 * 1024 * 1024 * 1024 # 1GB Limit # 🔐 Auth Toggle AUTH_ENABLED = False # <<< set False to disable authentication # Auth Wrapper (returns dependency or None) def auth_dependency(): return Depends(get_current_user) if AUTH_ENABLED else None # Ensure the upload directory exists os.makedirs(UPLOAD_DIR, exist_ok=True) # ========================================== # 🛠️ Helper: Safe Username Extraction # ========================================== def get_safe_username(user_obj): if not user_obj: return "anonymous" if hasattr(user_obj, "username"): return user_obj.username return user_obj.get("username", "anonymous") # ========================================== # 📤 Standard File Upload # ========================================== @router.post("/upload/") async def upload_file( file: UploadFile = File(...), content_length: int = Header(None), current_user: object = auth_dependency() ): if content_length and content_length > MAX_FILE_SIZE: raise HTTPException( status_code=status.HTTP_413_REQUEST_ENTITY_TOO_LARGE, detail="File is too large. Maximum limit is 1GB." ) username = get_safe_username(current_user) file_path = os.path.join(UPLOAD_DIR, file.filename) try: with open(file_path, "wb") as buffer: shutil.copyfileobj(file.file, buffer) except Exception as e: raise HTTPException(status_code=500, detail=f"Internal Server Error: {str(e)}") finally: file.file.close() return { "user": username, "filename": file.filename, "path": file_path, "message": "File uploaded successfully" } # ========================================== # 🚀 Remote Upload (Aria2) # ========================================== @router.post("/remote-upload/") async def remote_url_upload( url: str, custom_filename: str = None, current_user: object = auth_dependency() ): username = get_safe_username(current_user) command = ["aria2c", "-x", "6", "-s", "6", "-d", UPLOAD_DIR, url] if custom_filename: command.extend(["-o", custom_filename]) try: process = await asyncio.create_subprocess_exec( *command, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE ) stdout, stderr = await process.communicate() if process.returncode != 0: raise HTTPException(status_code=500, detail=f"Aria2 Error: {stderr.decode()}") except Exception as e: raise HTTPException(status_code=500, detail=f"System Error: {str(e)}") return { "user": username, "source": url, "filename": custom_filename or "Auto-detected", "message": "Download completed via Aria2" } # ========================================== # 📡 WebSocket Remote Upload (No Auth Applied) # ========================================== @router.websocket("/ws/remote-upload/") async def websocket_remote_upload(websocket: WebSocket): await websocket.accept() try: data = await websocket.receive_json() target_url = data.get("url") custom_filename = data.get("filename") if not target_url: await websocket.send_text("❌ Error: No URL provided.") await websocket.close() return command = ["aria2c", "-x", "6", "-s", "6", "-d", UPLOAD_DIR, target_url] if custom_filename: command.extend(["-o", custom_filename]) await websocket.send_text(f"🚀 Starting download: {target_url}") process = await asyncio.create_subprocess_exec( *command, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE ) while True: line = await process.stdout.readline() if not line: if process.returncode is not None: break await asyncio.sleep(0.1) continue decoded_line = line.decode().strip() if decoded_line: await websocket.send_text(decoded_line) await process.wait() if process.returncode == 0: await websocket.send_text("✅ Download Complete.") else: await websocket.send_text("❌ Download Failed.") await websocket.close() except WebSocketDisconnect: print("WS: Client disconnected") except Exception as e: try: await websocket.send_text(f"🔥 Critical Error: {str(e)}") await websocket.close() except: pass