Update app.py

app.py

@@ -5,6 +5,25 @@ import bcrypt
 from fastapi import FastAPI, HTTPException, Depends, Response, Request
 from fastapi.middleware.cors import CORSMiddleware
 from dotenv import load_dotenv
+from slowapi import Limiter
+from slowapi.util import get_remote_address
+from slowapi.errors import RateLimitExceeded
+from starlette.responses import JSONResponse
+
+# Initialize FastAPI app
+app = FastAPI()
+
+# Initialize rate limiter
+limiter = Limiter(key_func=get_remote_address)
+app.state.limiter = limiter
+
+# Custom error handler for too many requests
+@app.exception_handler(RateLimitExceeded)
+async def ratelimit_handler(request: Request, exc: RateLimitExceeded):
+    return JSONResponse(
+        {"detail": "Too many login attempts. Try again later."},
+        status_code=429
+    )
 
 # Load environment variables
 load_dotenv()
@@ -13,21 +32,16 @@ load_dotenv()
 SECRET_KEY = os.getenv("SECRET_KEY")
 TOKEN_EXPIRATION_MINUTES = int(os.getenv("TOKEN_EXPIRATION_MINUTES", 30))
 REFRESH_TOKEN_EXPIRATION_DAYS = int(os.getenv("REFRESH_TOKEN_EXPIRATION_DAYS", 7))
-ALLOWED_ORIGIN = os.getenv("ALLOWED_ORIGIN")
+ALLOWED_ORIGIN = os.getenv("ALLOWED_ORIGIN", "*").split(',')  # Ensure it's a list
 
-
-
-# Load environment variables
+# Load dummy user API key
 hashed_password = os.getenv("DUMMY_USER_KEY")
-
 if hashed_password:
     # Ensure it's stored as a hashed password (not plain text)
     hashed_password = bcrypt.hashpw(hashed_password.encode(), bcrypt.gensalt()).decode()
 
 # Fake database of API keys (hashed)
-API_KEYS_DB = {
-    "user1": hashed_password
-}
+API_KEYS_DB = {"user1": hashed_password} if hashed_password else {}
 
 def verify_api_key(api_key: str) -> bool:
     """Check if the provided API key is valid."""
@@ -36,53 +50,44 @@ def verify_api_key(api_key: str) -> bool:
             return True
     return False
 
-app = FastAPI()
-
 # Configure CORS for security (allow only trusted frontend)
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=[ALLOWED_ORIGIN],  
+    allow_origins=ALLOWED_ORIGIN,  
     allow_credentials=True,
     allow_methods=["GET", "POST"],
     allow_headers=["Authorization", "Content-Type"],
 )
 
-def create_jwt_token(expiration_minutes: int):
-    """Generate JWT Token."""
+def create_jwt_token(user_id: str, expiration_minutes: int):
+    """Generate JWT Token with user ID."""
     expiration = datetime.datetime.utcnow() + datetime.timedelta(minutes=expiration_minutes)
-    payload = {"exp": expiration}
+    payload = {"sub": user_id, "exp": expiration}
     return jwt.encode(payload, SECRET_KEY, algorithm="HS256")
 
 def verify_jwt_token(token: str):
     """Verify and decode JWT token."""
     try:
-        jwt.decode(token, SECRET_KEY, algorithms=["HS256"])
-        return True
+        return jwt.decode(token, SECRET_KEY, algorithms=["HS256"])
     except jwt.ExpiredSignatureError:
         raise HTTPException(status_code=401, detail="Token expired")
     except jwt.InvalidTokenError:
         raise HTTPException(status_code=401, detail="Invalid token")
 
-def verify_api_key(api_key: str) -> bool:
-    """Check if the provided API key is valid."""
-    for hashed_key in API_KEYS_DB.values():
-        if bcrypt.checkpw(api_key.encode(), hashed_key.encode()):
-            return True
-    return False
-
 @app.get("/sample-data")
 def get_sample_data():
     """Public endpoint: No token required"""
     return {"message": "This is public sample data."}
 
 @app.post("/login")
-def login(response: Response, api_key: str):
+@limiter.limit("5/minute")  # Allow max 5 login attempts per minute
+def login(request: Request, response: Response, api_key: str):
     """User must provide a valid API key to obtain JWT tokens."""
     if not verify_api_key(api_key):
         raise HTTPException(status_code=403, detail="Invalid API key")
 
-    access_token = create_jwt_token(TOKEN_EXPIRATION_MINUTES)
-    refresh_token = create_jwt_token(REFRESH_TOKEN_EXPIRATION_DAYS * 24 * 60)
+    access_token = create_jwt_token("user1", TOKEN_EXPIRATION_MINUTES)
+    refresh_token = create_jwt_token("user1", REFRESH_TOKEN_EXPIRATION_DAYS * 24 * 60)
 
     # Secure HTTP-only cookies (prevent XSS)
     response.set_cookie(
@@ -119,7 +124,7 @@ def refresh_token(request: Request, response: Response):
         raise HTTPException(status_code=401, detail="Invalid refresh token")
 
     # Issue new access token
-    new_access_token = create_jwt_token(TOKEN_EXPIRATION_MINUTES)
+    new_access_token = create_jwt_token("user1", TOKEN_EXPIRATION_MINUTES)
     response.set_cookie(
         key="access_token",
         value=new_access_token,