PacketStation/server/packetstation/auth/sessions.py

"""Server-side sessions. The cookie holds a random token; the row holds HMAC-SHA256(session key, token)."""

from __future__ import annotations

import hashlib
import hmac
import secrets
import uuid
from dataclasses import dataclass
from typing import Any

from psycopg.rows import dict_row
from psycopg.types.json import Jsonb

from .. import events
from ..settings import Settings
from . import secsettings

SESSION_COOKIE = "ps_session"
DEVICE_COOKIE = "ps_device"


@dataclass(frozen=True)
class Actor:
    """Who is acting in a session: what step-up checks, signatures and audit rows need."""

    session_id: uuid.UUID
    user_id: uuid.UUID
    sam_account: str
    display_name: str
    kind: str  # password | pin
    roles: tuple[str, ...]

    @property
    def role_snapshot(self) -> str:
        return ",".join(self.roles)


def new_token() -> str:
    return secrets.token_urlsafe(32)


def token_hash(key: bytes, token: str) -> bytes:
    return hmac.new(key, token.encode("utf-8"), hashlib.sha256).digest()


def create(
    conn: Any,
    settings: Settings,
    key: bytes,
    *,
    user_id: Any,
    kind: str,
    roles: list[str],
    groups: list[str],
    device_id: Any = None,
    ip: str | None = None,
    user_agent: str | None = None,
) -> tuple[str, dict[str, Any]]:
    """Insert a session row; returns (cookie token, row). The token itself is never stored."""
    token = new_token()
    settings = secsettings.overlay(conn, settings)
    hours = settings.session_pin_hours if kind == "pin" else settings.session_password_hours
    with conn.cursor(row_factory=dict_row) as cur:
        row = cur.execute(
            "INSERT INTO sessions (user_id, token_hash, kind, device_id, expires_at, last_seen_at, ip, user_agent,"
            " roles_snapshot, roles_checked_at)"
            " VALUES (%s, %s, %s::session_kind, %s, now() + make_interval(hours => %s::int), now(), %s::inet, %s, %s,"
            " now()) RETURNING *",
            (
                user_id,
                token_hash(key, token),
                kind,
                device_id,
                hours,
                ip,
                (user_agent or "")[:400] or None,
                Jsonb({"roles": roles, "groups": groups}),
            ),
        ).fetchone()
    assert row is not None
    return token, row


def load(conn: Any, key: bytes, token: str) -> dict[str, Any] | None:
    """The session row for a cookie token plus `db_now`, or None. Revoked and expired rows are returned too."""
    with conn.cursor(row_factory=dict_row) as cur:
        row: dict[str, Any] | None = cur.execute(
            "SELECT s.*, s.kind::text AS kind, now() AS db_now FROM sessions s WHERE s.token_hash = %s",
            (token_hash(key, token),),
        ).fetchone()
    return row


def revoke(conn: Any, session_id: Any, reason: str, *, notify: bool = True) -> bool:
    """Revoke one session. Publishes a `session` event so that browser signs out. False when already revoked."""
    row = conn.execute(
        "UPDATE sessions SET revoked_at = now(), revoked_reason = %s WHERE id = %s AND revoked_at IS NULL RETURNING id",
        (reason, session_id),
    ).fetchone()
    if row is None:
        return False
    if notify:
        events.publish(conn, {"event": "session", "session_id": str(row[0]), "action": "revoked"})
    return True


def revoke_where(conn: Any, where: str, params: tuple[Any, ...], reason: str) -> list[uuid.UUID]:
    """Revoke every active session matching a WHERE fragment; publishes one `session` event each."""
    rows = conn.execute(
        f"UPDATE sessions SET revoked_at = now(), revoked_reason = %s WHERE revoked_at IS NULL AND ({where})"  # noqa: S608
        " RETURNING id",
        (reason, *params),
    ).fetchall()
    ids = [r[0] for r in rows]
    for sid in ids:
        events.publish(conn, {"event": "session", "session_id": str(sid), "action": "revoked"})
    return ids


def revoke_user(conn: Any, user_id: Any, reason: str, *, except_id: Any = None) -> list[uuid.UUID]:
    if except_id is None:
        return revoke_where(conn, "user_id = %s", (user_id,), reason)
    return revoke_where(conn, "user_id = %s AND id <> %s", (user_id, except_id), reason)