import shutil import os import asyncio from fastapi import ( APIRouter, UploadFile, File, HTTPException, Header, status, Depends, WebSocket, WebSocketDisconnect ) from fastapi.responses import HTMLResponse # Import Auth System (Assuming this exists in your project structure) # from app.routers.auth.system import get_current_user # Mocking auth for standalone functionality if import fails try: from app.routers.auth.system import get_current_user except ImportError: def get_current_user(): return "anonymous" # 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 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 if isinstance(user_obj, dict): return user_obj.get("username", "anonymous") return str(user_obj) # ========================================== # 📤 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" } # ========================================== # 📡 WebSocket Remote Upload # ========================================== @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 # Check if aria2c is installed if shutil.which("aria2c") is None: await websocket.send_text("❌ Error: 'aria2c' is not installed on the server.") await websocket.close() return command = ["aria2c", "-x", "8", "-s", "8", "-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 ) # Stream stdout line by line 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: stderr_data = await process.stderr.read() await websocket.send_text(f"❌ Download Failed: {stderr_data.decode()}") 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 # ========================================== # 🎨 Modern UI Endpoint # ========================================== # Path becomes: /drive/storage/ui @router.get("/storage/ui", response_class=HTMLResponse) async def drive_ui(): html_content = """ Cloud Drive | by Sam
📂

Click to upload or drag & drop

Max file size: 1GB

""" return html_content