feat(security): implement IP rate limiting and secure session cookies on Hub

src/Server.py

@@ -10,6 +10,7 @@ import asyncio
 import secrets
 import hashlib
 import hmac
+import time
 import warnings
 import threading
 from datetime import datetime, timezone, timedelta
@@ -88,6 +89,54 @@ class ChangePasswordRequest(BaseModel):
     new_password: str
 
 
+# --- Rate Limiting & Brute-Force Mitigation ---
+
+_RATE_LIMIT_LOCK = threading.Lock()
+_LOGIN_FAILURES: Dict[str, List[float]] = {}
+_ENROLL_FAILURES: Dict[str, List[float]] = {}
+
+
+def get_client_ip(request: Optional[Request] = None) -> str:
+    """Extracts client IP address respecting reverse-proxy headers."""
+    if not request:
+        return "127.0.0.1"
+    forwarded = request.headers.get("x-forwarded-for")
+    if forwarded:
+        return forwarded.split(",")[0].strip()
+    real_ip = request.headers.get("x-real-ip")
+    if real_ip:
+        return real_ip.strip()
+    if request.client:
+        return request.client.host
+    return "127.0.0.1"
+
+
+def check_ip_rate_limit(ip: str, store: Dict[str, List[float]], max_attempts: int = 10, window_seconds: int = 60) -> bool:
+    """Checks whether the client IP has exceeded failed attempts within the sliding time window."""
+    now = datetime.now(timezone.utc).timestamp()
+    with _RATE_LIMIT_LOCK:
+        history = store.get(ip, [])
+        history = [ts for ts in history if now - ts < window_seconds]
+        store[ip] = history
+        return len(history) < max_attempts
+
+
+def record_failed_ip_attempt(ip: str, store: Dict[str, List[float]]):
+    """Records a failed attempt timestamp for the client IP."""
+    now = datetime.now(timezone.utc).timestamp()
+    with _RATE_LIMIT_LOCK:
+        if ip not in store:
+            store[ip] = []
+        store[ip].append(now)
+
+
+def reset_failed_ip_attempts(ip: str, store: Dict[str, List[float]]):
+    """Resets failed attempt records for client IP upon successful authentication."""
+    with _RATE_LIMIT_LOCK:
+        if ip in store:
+            del store[ip]
+
+
 # --- Authentication & Password Hashing Utilities ---
 
 def hash_password(password: str, salt: Optional[str] = None) -> Tuple[str, str]:
@@ -1068,6 +1117,10 @@ def login_page(request: Request):
 @app.post("/login")
 async def process_login(request: Request):
     """Processes login credentials and issues session cookie."""
+    client_ip = get_client_ip(request)
+    if not check_ip_rate_limit(client_ip, _LOGIN_FAILURES, max_attempts=10, window_seconds=60):
+        raise HTTPException(status_code=429, detail="Too many failed login attempts. Please try again later.")
+
     username, password = None, None
     content_type = request.headers.get("content-type", "")
     if "application/json" in content_type:
@@ -1087,14 +1140,19 @@ async def process_login(request: Request):
 
     db_path = SERVER_STATE.get("config", {}).get("db_path", DEFAULT_DB_FILE)
     if not verify_user_credentials(username, password, db_path=db_path):
+        record_failed_ip_attempt(client_ip, _LOGIN_FAILURES)
+        await asyncio.sleep(0.5)
         raise HTTPException(status_code=401, detail="Invalid username or password")
 
+    reset_failed_ip_attempts(client_ip, _LOGIN_FAILURES)
     token = generate_session_token(username)
+    is_https = request.url.scheme == "https" or request.headers.get("x-forwarded-proto") == "https"
     response = JSONResponse(content={"status": "success", "username": username, "token": token})
     response.set_cookie(
         key=SESSION_COOKIE_NAME,
         value=token,
         httponly=True,
+        secure=is_https,
         max_age=86400,
         samesite="lax"
     )
@@ -1122,12 +1180,16 @@ async def api_change_password(req: ChangePasswordRequest, current_user: str = De
 
 
 @app.post("/api/client/enroll")
-def enroll_client(req: ClientEnrollRequest):
+def enroll_client(req: ClientEnrollRequest, request: Request = None):
     """
     Enrolls an edge client by validating the enrollment secret,
     checking license seat limits, issuing a signed client certificate + key,
     and recording the client in the SQLite accounting database.
     """
+    client_ip = get_client_ip(request)
+    if not check_ip_rate_limit(client_ip, _ENROLL_FAILURES, max_attempts=15, window_seconds=60):
+        raise HTTPException(status_code=429, detail="Too many failed enrollment attempts. Please try again later.")
+
     db_path = SERVER_STATE.get("config", {}).get("db_path", DEFAULT_DB_FILE)
     conn = sqlite3.connect(db_path)
     c = conn.cursor()
@@ -1142,8 +1204,12 @@ def enroll_client(req: ClientEnrollRequest):
     expected_secret, max_seats = row
     if not secrets.compare_digest(str(req.enrollment_secret), str(expected_secret)):
         conn.close()
+        record_failed_ip_attempt(client_ip, _ENROLL_FAILURES)
+        time.sleep(0.5)
         raise HTTPException(status_code=403, detail="Invalid enrollment secret")
 
+    reset_failed_ip_attempts(client_ip, _ENROLL_FAILURES)
+
     ca_cert = SERVER_STATE.get("ca_cert")
     ca_key = SERVER_STATE.get("ca_key")
     ca_cert_pem = SERVER_STATE.get("ca_cert_pem")