feat: launch PK testing application

src/pk_lab_server/security.py

@@ -0,0 +1,86 @@
+import base64
+import binascii
+import hashlib
+import hmac
+import json
+import secrets
+import time
+from dataclasses import dataclass
+
+from fastapi import Depends, HTTPException
+from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
+from sqlalchemy.orm import Session
+
+from .config import get_settings
+from .db import get_db
+from .models import User
+
+
+bearer = HTTPBearer(auto_error=False)
+
+
+def _b64(data: bytes) -> str:
+    return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii")
+
+
+def _unb64(value: str) -> bytes:
+    return base64.urlsafe_b64decode(value + "=" * (-len(value) % 4))
+
+
+def issue_token(user_id: str, app_key: str) -> str:
+    settings = get_settings()
+    payload = {
+        "sub": user_id,
+        "app": app_key,
+        "iat": int(time.time()),
+        "exp": int(time.time()) + settings.token_ttl_seconds,
+        "nonce": secrets.token_hex(8),
+    }
+    encoded = _b64(json.dumps(payload, separators=(",", ":"), sort_keys=True).encode())
+    signature = hmac.new(settings.token_secret.encode(), encoded.encode(), hashlib.sha256).digest()
+    return encoded + "." + _b64(signature)
+
+
+@dataclass(frozen=True)
+class TokenClaims:
+    user_id: str
+    app_key: str
+
+
+def parse_token(token: str) -> TokenClaims:
+    settings = get_settings()
+    try:
+        encoded, supplied = token.split(".", 1)
+        expected = hmac.new(settings.token_secret.encode(), encoded.encode(), hashlib.sha256).digest()
+        if not hmac.compare_digest(expected, _unb64(supplied)):
+            raise ValueError("signature")
+        payload = json.loads(_unb64(encoded))
+        if int(payload["exp"]) < int(time.time()):
+            raise ValueError("expired")
+        return TokenClaims(user_id=str(payload["sub"]), app_key=str(payload["app"]))
+    except (
+        KeyError,
+        ValueError,
+        TypeError,
+        UnicodeDecodeError,
+        binascii.Error,
+        json.JSONDecodeError,
+    ) as exc:
+        raise HTTPException(status_code=401, detail="invalid or expired token") from exc
+
+
+def current_claims(
+    credentials: HTTPAuthorizationCredentials | None = Depends(bearer),
+) -> TokenClaims:
+    if credentials is None or credentials.scheme.lower() != "bearer":
+        raise HTTPException(status_code=401, detail="authentication required")
+    return parse_token(credentials.credentials)
+
+
+def current_user(
+    claims: TokenClaims = Depends(current_claims), db: Session = Depends(get_db)
+) -> User:
+    user = db.get(User, claims.user_id)
+    if user is None:
+        raise HTTPException(status_code=401, detail="user no longer exists")
+    return user