Update main.py

main.py

@@ -1,11 +1,40 @@
-from fastapi import FastAPI, UploadFile, File, HTTPException
+from fastapi import FastAPI, UploadFile, File, HTTPException, Depends, status
+from fastapi.security import HTTPBasic, HTTPBasicCredentials
 from fastapi.responses import FileResponse
 from typing import List, Dict
 import os
 import shutil
 import uuid
+from collections import defaultdict
+import secrets
+from config import DOCS_USERNAME, DOCS_PASSWORD
 
-app = FastAPI()
+# Initialize security
+security = HTTPBasic()
+
+def verify_credentials(credentials: HTTPBasicCredentials = Depends(security)):
+    """Verify HTTP Basic Auth credentials"""
+    is_username_correct = secrets.compare_digest(credentials.username, DOCS_USERNAME)
+    is_password_correct = secrets.compare_digest(credentials.password, DOCS_PASSWORD)
+    
+    if not (is_username_correct and is_password_correct):
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid credentials",
+            headers={"WWW-Authenticate": "Basic"},
+        )
+    return credentials
+
+# Initialize FastAPI with auth for docs
+app = FastAPI(
+    title="File Sharing API",
+    description="API for file sharing service",
+    version="1.0.0",
+    docs_url="/docs",
+    redoc_url="/redoc",
+    dependencies=[Depends(verify_credentials)],
+    openapi_url="/openapi.json"
+)
 
 # Create uploads directory if it doesn't exist
 UPLOAD_DIR = "uploads"
@@ -17,6 +46,37 @@ file_codes: Dict[str, dict] = {}
 # Reverse mapping: {filename: {"code": code, "user_id": user_id}}
 filename_codes: Dict[str, dict] = {}
 
+# Add this class after other class definitions
+class UserStats:
+    def __init__(self):
+        self.upload_counts = defaultdict(int)
+        self.download_counts = defaultdict(int)
+        self.total_bytes_uploaded = defaultdict(int)
+        self.total_bytes_downloaded = defaultdict(int)
+        self.last_activity = defaultdict(str)
+
+    def log_upload(self, user_id: str, file_size: int, filename: str):
+        self.upload_counts[user_id] += 1
+        self.total_bytes_uploaded[user_id] += file_size
+        self.last_activity[user_id] = f"Uploaded: {filename}"
+
+    def log_download(self, user_id: str, file_size: int, filename: str):
+        self.download_counts[user_id] += 1
+        self.total_bytes_downloaded[user_id] += file_size
+        self.last_activity[user_id] = f"Downloaded: {filename}"
+
+    def get_user_stats(self, user_id: str) -> dict:
+        return {
+            "uploads": self.upload_counts[user_id],
+            "downloads": self.download_counts[user_id],
+            "bytes_uploaded": self.total_bytes_uploaded[user_id],
+            "bytes_downloaded": self.total_bytes_downloaded[user_id],
+            "last_activity": self.last_activity.get(user_id, "No activity")
+        }
+
+# Initialize stats
+user_stats = UserStats()
+
 @app.post("/upload/")
 async def upload_file(file: UploadFile = File(...), user_id: str = None):
     try:
@@ -32,6 +92,10 @@ async def upload_file(file: UploadFile = File(...), user_id: str = None):
         file_codes[unique_code] = {"filename": file.filename, "user_id": user_id}
         filename_codes[file.filename] = {"code": unique_code, "user_id": user_id}
         
+        # Log the upload
+        file_size = os.path.getsize(file_path)
+        user_stats.log_upload(user_id, file_size, file.filename)
+        
         return {
             "filename": file.filename,
             "access_code": unique_code,
@@ -144,6 +208,23 @@ async def direct_download(access_code: str):
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
+# Add new endpoint for stats
+@app.get("/stats/{user_id}")
+async def get_user_stats(user_id: str):
+    try:
+        stats = user_stats.get_user_stats(user_id)
+        return stats
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/log_download")
+async def log_download(user_id: str, file_size: int, filename: str):
+    try:
+        user_stats.log_download(user_id, file_size, filename)
+        return {"message": "Download logged successfully"}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)