1"""Authentication: admin sessions (signed cookie) and API keys (hashed)."""23from __future__ import annotations45import hashlib6import secrets7import time8from dataclasses import dataclass910from argon2 import PasswordHasher11from argon2.exceptions import VerifyMismatchError12from fastapi import Depends, Request13from itsdangerous import BadSignature, SignatureExpired, URLSafeTimedSerializer1415from .config import Settings, get_settings16from .db import Database17from .errors import AuthError, Forbidden, RateLimited1819SESSION_COOKIE = "llm_session"20KEY_PREFIX = "llm_live_"21ph = PasswordHasher()222324def hash_password(pw: str) -> str:25 return ph.hash(pw)262728def verify_password(hash_: str, pw: str) -> bool:29 try:30 return ph.verify(hash_, pw)31 except VerifyMismatchError:32 return False33 except Exception:34 return False353637def hash_key(raw: str) -> str:38 return hashlib.sha256(raw.encode()).hexdigest()394041def generate_key() -> str:42 return KEY_PREFIX + secrets.token_urlsafe(32).replace("-", "").replace("_", "")[:40]434445@dataclass46class Principal:47 kind: str # session | api_key48 id: int | None49 name: str50 scopes: set[str]5152 def has(self, scope: str) -> bool:53 return "admin" in self.scopes or scope in self.scopes545556class RateLimiter:57 """Small in-memory sliding window limiter keyed by string."""5859 def __init__(self, limit: int, window: float):60 self.limit = limit61 self.window = window62 self.hits: dict[str, list[float]] = {}6364 def check(self, key: str) -> bool:65 now = time.time()66 arr = [t for t in self.hits.get(key, []) if now - t < self.window]67 if len(arr) >= self.limit:68 self.hits[key] = arr69 return False70 arr.append(now)71 self.hits[key] = arr72 return True737475login_limiter = RateLimiter(limit=8, window=60)76api_limiter = RateLimiter(limit=600, window=60)777879class Auth:80 def __init__(self, db: Database, settings: Settings):81 self.db = db82 self.settings = settings83 self.serializer = URLSafeTimedSerializer(settings.secret, salt="llm-api-session")8485 # ---- users ----------------------------------------------------------------86 async def user_count(self) -> int:87 return int(await self.db.scalar("SELECT COUNT(*) FROM users") or 0)8889 async def create_user(self, email: str, password: str) -> int:90 if len(password) < 10:91 raise ValueError("password must be at least 10 characters")92 return await self.db.execute("INSERT INTO users(email, password_hash, role, created_at) VALUES(?,?,?,?)",93 (email.strip().lower(), hash_password(password), "admin", time.time()))9495 async def authenticate(self, email: str, password: str) -> dict | None:96 row = await self.db.fetchone("SELECT * FROM users WHERE email=?", (email.strip().lower(),))97 if not row or not verify_password(row["password_hash"], password):98 return None99 await self.db.execute("UPDATE users SET last_login_at=? WHERE id=?", (time.time(), row["id"]))100 return row101102 async def change_password(self, user_id: int, new_password: str) -> None:103 if len(new_password) < 10:104 raise ValueError("password must be at least 10 characters")105 await self.db.execute("UPDATE users SET password_hash=? WHERE id=?", (hash_password(new_password), user_id))106107 # ---- sessions ---------------------------------------------------------------108 def make_session(self, user_id: int, email: str) -> str:109 return self.serializer.dumps({"uid": user_id, "email": email, "n": secrets.token_hex(8)})110111 def read_session(self, token: str) -> dict | None:112 try:113 return self.serializer.loads(token, max_age=self.settings.session_hours * 3600)114 except (BadSignature, SignatureExpired):115 return None116117 # ---- api keys ---------------------------------------------------------------118 async def create_key(self, name: str, scopes: list[str]) -> tuple[str, dict]:119 raw = generate_key()120 row_id = await self.db.execute(121 "INSERT INTO api_keys(name, prefix, key_hash, scopes, created_at) VALUES(?,?,?,?,?)",122 (name, raw[:16], hash_key(raw), ",".join(sorted(set(scopes))), time.time()))123 row = await self.db.fetchone("SELECT * FROM api_keys WHERE id=?", (row_id,))124 return raw, row # type: ignore[return-value]125126 async def lookup_key(self, raw: str) -> dict | None:127 row = await self.db.fetchone("SELECT * FROM api_keys WHERE key_hash=? AND revoked_at IS NULL", (hash_key(raw),))128 return row129130 async def touch_key(self, key_id: int) -> None:131 await self.db.execute("UPDATE api_keys SET last_used_at=?, request_count=request_count+1 WHERE id=?",132 (time.time(), key_id))133134135# ---------------------------------------------------------------------------136# FastAPI dependencies137# ---------------------------------------------------------------------------138139140def get_auth(request: Request) -> Auth:141 return request.app.state.auth142143144def _client_ip(request: Request) -> str:145 fwd = request.headers.get("x-forwarded-for")146 if fwd:147 return fwd.split(",")[0].strip()148 return request.client.host if request.client else "?"149150151async def principal_from_request(request: Request, auth: Auth) -> Principal | None:152 header = request.headers.get("authorization") or ""153 if header.lower().startswith("bearer "):154 raw = header[7:].strip()155 if raw:156 row = await auth.lookup_key(raw)157 if not row:158 return None159 await auth.touch_key(row["id"])160 return Principal("api_key", row["id"], row["name"], set(row["scopes"].split(",")))161 api_key_header = request.headers.get("x-api-key")162 if api_key_header:163 row = await auth.lookup_key(api_key_header.strip())164 if not row:165 return None166 await auth.touch_key(row["id"])167 return Principal("api_key", row["id"], row["name"], set(row["scopes"].split(",")))168 tok = request.cookies.get(SESSION_COOKIE)169 if tok:170 data = auth.read_session(tok)171 if data:172 return Principal("session", data["uid"], data["email"], {"admin"})173 return None174175176async def require_inference(request: Request, auth: Auth = Depends(get_auth)) -> Principal:177 """Bearer API key (scope inference/admin) or an admin session (playground)."""178 if not api_limiter.check(_client_ip(request)):179 raise RateLimited("Too many requests.")180 p = await principal_from_request(request, auth)181 if not p:182 if (request.headers.get("authorization") or request.headers.get("x-api-key")):183 raise AuthError("Incorrect API key provided.")184 raise AuthError("Missing API key. Pass it as 'Authorization: Bearer llm_live_...'.")185 if not p.has("inference"):186 raise Forbidden("This API key does not have the 'inference' scope.")187 request.state.principal = p188 return p189190191async def require_admin(request: Request, auth: Auth = Depends(get_auth)) -> Principal:192 """Admin session (dashboard) or API key with the admin scope. Mutations from a session need the CSRF header."""193 p = await principal_from_request(request, auth)194 if not p:195 raise AuthError("Authentication required.", code="UNAUTHENTICATED")196 if not p.has("admin"):197 raise Forbidden("Admin scope required.")198 if p.kind == "session" and request.method not in ("GET", "HEAD", "OPTIONS"):199 if request.headers.get("x-llm-csrf") != "1":200 raise Forbidden("Missing CSRF header.", code="CSRF")201 origin = request.headers.get("origin")202 if origin:203 host = request.headers.get("host", "")204 settings = get_settings()205 allowed = {settings.public_url.rstrip("/"), f"http://{host}", f"https://{host}", f"http://127.0.0.1:{settings.port}",206 f"http://localhost:{settings.port}", "http://localhost:3000", "http://127.0.0.1:3000"}207 if origin.rstrip("/") not in allowed:208 raise Forbidden("Origin not allowed.", code="CSRF")209 request.state.principal = p210 return p211