src/auth/session.py

"""HMAC-signed session cookies — no database, no session store.

The cookie payload is a base64 JSON blob {uid, username, exp} followed by an
HMAC-SHA256 signature over that blob using SESSION_SECRET. Tampering with the
payload invalidates the signature; expiry is checked server-side.

This is a stateless session: the cookie IS the session. Revocation means
rotating SESSION_SECRET (which invalidates every outstanding session).
"""
from __future__ import annotations   # defer annotation evaluation (PEP 563)

import base64   # url-safe base64 encoding of the payload + signature
import hashlib   # sha256 for the HMAC
import hmac   # constant-time signature comparison (prevents timing attacks)
import json   # serialize the session payload to JSON
import os   # read SESSION_SECRET env var
import time   # unix timestamps for expiry math

COOKIE_NAME = "dr_session"   # the browser cookie name
COOKIE_MAX_AGE = 7 * 24 * 60 * 60   # 7 days in seconds (cookie lifetime)

# The signing secret. MUST be set in production (HF Space secret). If unset,
# sessions are disabled and only Bearer-key auth works.
_SECRET = os.environ.get("SESSION_SECRET", "").encode("utf-8")   # bytes for HMAC


def _sign(payload_b64: str) -> str:   # HMAC-SHA256 over the base64 payload
    """Return a url-safe base64 HMAC signature over the payload string."""
    mac = hmac.new(_SECRET, payload_b64.encode("utf-8"), hashlib.sha256).digest()   # compute HMAC
    return base64.urlsafe_b64encode(mac).decode("ascii")   # encode signature to ascii string


def create_session(uid: str, username: str, reveal_key: str = "") -> str:   # build a signed session cookie value
    """Create a signed session cookie value for a Google user.

    reveal_key: an optional one-time plaintext API key to show the user at
    signup. Carried inside the HttpOnly signed cookie so it can't be forged
    or stolen by XSS. Cleared by /auth/clear-reveal after first display.
    """
    if not _SECRET:   # no secret configured
        raise RuntimeError("SESSION_SECRET is not set")   # refuse to issue sessions
    payload = {   # the session contents
        "uid": str(uid),   # provider user id (string for JSON safety)
        "username": username,   # display name or email (display handle)
        "exp": int(time.time()) + COOKIE_MAX_AGE,   # absolute expiry timestamp
    }
    if reveal_key:   # carry the one-time key reveal inside the signed cookie
        payload["reveal_key"] = reveal_key   # /me returns it once, then clear-reveal drops it
    payload_b64 = base64.urlsafe_b64encode(   # encode the JSON payload
        json.dumps(payload, separators=(",", ":")).encode("utf-8")   # compact JSON (no spaces)
    ).decode("ascii")   # to ascii string
    sig = _sign(payload_b64)   # sign the encoded payload
    return f"{payload_b64}.{sig}"   # payload.signature cookie value


def verify_session(cookie_value: str) -> dict | None:   # validate + decode a session cookie
    """Return the session payload if valid, else None.

    Checks: signature matches (constant-time), not expired, well-formed.
    Returns None on ANY failure so callers can treat it as 'not logged in'.
    """
    if not _SECRET or not cookie_value:   # no secret OR empty cookie
        return None   # cannot verify
    parts = cookie_value.split(".", 1)   # split payload.signature
    if len(parts) != 2:   # malformed cookie
        return None   # reject
    payload_b64, sig = parts   # unpack
    expected = _sign(payload_b64)   # recompute the expected signature
    if not hmac.compare_digest(sig, expected):   # constant-time compare (timing-safe)
        return None   # signature mismatch -> tampered or forged
    try:   # decode the payload
        raw = base64.urlsafe_b64decode(payload_b64.encode("ascii"))   # base64 decode
        payload = json.loads(raw)   # parse JSON
    except Exception:   # corrupt payload
        return None   # reject
    if not isinstance(payload, dict):   # not a JSON object
        return None   # reject
    if int(payload.get("exp", 0)) < int(time.time()):   # expired?
        return None   # session expired
    return payload   # valid session -> return {uid, username, exp}


def cookie_kwargs() -> dict:   # the cookie attributes we always set
    """Return the kwargs for setting the session cookie (security flags)."""
    return {   # starlette set_cookie kwargs
        "key": COOKIE_NAME,   # cookie name
        "httponly": True,   # JS cannot read it (XSS protection)
        "secure": os.environ.get("SESSION_COOKIE_SECURE", "1") != "0",   # HTTPS-only in prod (set 0 for local http)
        "samesite": "lax",   # allows top-level navigation (OAuth callback) but blocks CSRF
        "max_age": COOKIE_MAX_AGE,   # 7 days
        "path": "/",   # cookie visible on all paths
    }