Spaces:
Sleeping
Sleeping
File size: 5,092 Bytes
55eb784 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 |
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
|