Update app.py

app.py

@@ -1,5 +1,6 @@
-from fastapi import FastAPI, UploadFile, File, HTTPException
+from fastapi import FastAPI, UploadFile, File, HTTPException, Depends
 from fastapi.responses import FileResponse
+from fastapi.security import OAuth2PasswordBearer
 from pathlib import Path
 import psutil
 import shutil
@@ -11,8 +12,11 @@ app = FastAPI()
 UPLOAD_DIRECTORY = Path("uploads")
 UPLOAD_DIRECTORY.mkdir(parents=True, exist_ok=True)
 
-# Hugging Face Spaces has a 50GB limit
-TOTAL_SPACE_GB = 50
+# OAuth2 scheme for token-based authorization
+oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
+
+# Hugging Face secret token from environment variables
+VALID_TOKEN = os.environ.get("SECRET_TOKEN")
 
 # Helper function to calculate the size of a directory
 def get_directory_size(directory: Path) -> int:
@@ -23,16 +27,22 @@ def get_directory_size(directory: Path) -> int:
             total_size += os.path.getsize(fp)
     return total_size
 
+# Dependency to verify the token
+def get_current_token(token: str = Depends(oauth2_scheme)):
+    if token != VALID_TOKEN:
+        raise HTTPException(status_code=401, detail="Invalid token")
+    return token
+
 # Health check endpoint
 @app.get("/health")
-def health_check():
+def health_check(token: str = Depends(get_current_token)):
     return {
         "status": "healthy"
     }
 
 # System metrics endpoint: CPU, RAM, and disk usage for uploads folder (in GB and percentage of 50GB limit)
 @app.get("/metrics")
-def get_metrics():
+def get_metrics(token: str = Depends(get_current_token)):
     # CPU percentage (rounded to the nearest whole number)
     cpu_percent = round(psutil.cpu_percent(interval=1))
 
@@ -44,7 +54,7 @@ def get_metrics():
     # Disk stats for uploads directory
     uploads_size_bytes = get_directory_size(UPLOAD_DIRECTORY)
     uploads_size_gb = uploads_size_bytes / (1024 ** 3)
-    uploads_percent = (uploads_size_gb / TOTAL_SPACE_GB) * 100
+    uploads_percent = (uploads_size_gb / 50) * 100  # Assuming 50GB total space limit
 
     return {
         "cpu_percent": cpu_percent,  # Rounded CPU percentage
@@ -56,13 +66,13 @@ def get_metrics():
         "disk": {
             "uploads_folder_size_gb": round(uploads_size_gb, 2),
             "uploads_usage_percent_of_50gb": round(uploads_percent, 2),
-            "total_space_gb": TOTAL_SPACE_GB
+            "total_space_gb": 50
         }
     }
 
 # File upload endpoint
 @app.post("/uploadfile/")
-async def upload_file(file: UploadFile = File(...)):
+async def upload_file(file: UploadFile = File(...), token: str = Depends(get_current_token)):
     file_location = UPLOAD_DIRECTORY / file.filename
 
     with file_location.open("wb") as buffer:
@@ -72,7 +82,7 @@ async def upload_file(file: UploadFile = File(...)):
 
 # File download endpoint
 @app.get("/downloadfile/{filename}")
-def download_file(filename: str):
+def download_file(filename: str, token: str = Depends(get_current_token)):
     file_location = UPLOAD_DIRECTORY / filename
 
     if not file_location.exists():
@@ -82,6 +92,6 @@ def download_file(filename: str):
 
 # Show current files endpoint
 @app.get("/files/")
-def list_files():
+def list_files(token: str = Depends(get_current_token)):
     files = [f for f in os.listdir(UPLOAD_DIRECTORY) if os.path.isfile(UPLOAD_DIRECTORY / f)]
     return {"files": files}