stackchain-dashboard/src/session_store.py

"""Durable active-session registry used to revoke signed operator sessions."""

import hashlib
import secrets
import sqlite3
from dataclasses import dataclass
from pathlib import Path
from typing import Callable


class SessionStoreError(RuntimeError):
    """Raised when session state cannot be read or changed safely."""


@dataclass(frozen=True)
class ActiveDevice:
    management_id: str
    device_label: str
    created_at: int
    expires_at: int
    current: bool


class SessionStore:
    def __init__(
        self,
        path: str | Path,
        *,
        clock: Callable[[], float],
        lock_timeout_seconds: float = 0.1,
    ) -> None:
        self.path = Path(path)
        self.clock = clock
        self.lock_timeout_seconds = lock_timeout_seconds

    @staticmethod
    def _digest(session_id: str) -> str:
        return hashlib.sha256(session_id.encode()).hexdigest()

    def _connect(self, *, initialize: bool = False) -> sqlite3.Connection:
        try:
            if initialize:
                self.path.parent.mkdir(parents=True, exist_ok=True)
                connection = sqlite3.connect(
                    self.path, timeout=self.lock_timeout_seconds
                )
            else:
                connection = sqlite3.connect(
                    f"{self.path.resolve().as_uri()}?mode=rw",
                    timeout=self.lock_timeout_seconds,
                    uri=True,
                )
            if initialize:
                connection.execute(
                    """
                    CREATE TABLE IF NOT EXISTS active_sessions (
                        session_hash TEXT PRIMARY KEY,
                        expires_at INTEGER NOT NULL,
                        management_id TEXT,
                        device_label TEXT,
                        created_at INTEGER
                    )
                    """
                )
                columns = {
                    row[1] for row in connection.execute("PRAGMA table_info(active_sessions)")
                }
                additions = {
                    "management_id": "TEXT",
                    "device_label": "TEXT",
                    "created_at": "INTEGER",
                }
                for name, column_type in additions.items():
                    if name not in columns:
                        connection.execute(
                            f"ALTER TABLE active_sessions ADD COLUMN {name} {column_type}"
                        )
                connection.execute(
                    "UPDATE active_sessions SET management_id = lower(hex(randomblob(16))) "
                    "WHERE management_id IS NULL"
                )
                connection.execute(
                    "UPDATE active_sessions SET device_label = 'Existing device' "
                    "WHERE device_label IS NULL"
                )
                connection.execute(
                    "UPDATE active_sessions SET created_at = ? WHERE created_at IS NULL",
                    (int(self.clock()),),
                )
                connection.execute(
                    "CREATE UNIQUE INDEX IF NOT EXISTS active_sessions_management_id "
                    "ON active_sessions(management_id)"
                )
            return connection
        except (OSError, sqlite3.Error) as exc:
            raise SessionStoreError("Session registry is temporarily unavailable") from exc

    def activate(
        self, session_id: str, expires_at: int, *, device_label: str = "This device"
    ) -> None:
        label = " ".join(str(device_label).split())[:64] or "This device"
        now = int(self.clock())
        try:
            with self._connect(initialize=True) as connection:
                connection.execute(
                    "DELETE FROM active_sessions WHERE expires_at <= ?", (now,)
                )
                connection.execute(
                    "INSERT INTO active_sessions("
                    "session_hash, expires_at, management_id, device_label, created_at"
                    ") VALUES (?, ?, ?, ?, ?)",
                    (
                        self._digest(session_id),
                        expires_at,
                        secrets.token_urlsafe(18),
                        label,
                        now,
                    ),
                )
        except (OSError, sqlite3.Error) as exc:
            raise SessionStoreError("Session registry is temporarily unavailable") from exc

    def is_active(self, session_id: str, expires_at: int) -> bool:
        now = int(self.clock())
        try:
            with self._connect() as connection:
                row = connection.execute(
                    "SELECT expires_at FROM active_sessions WHERE session_hash = ?",
                    (self._digest(session_id),),
                ).fetchone()
        except (OSError, sqlite3.Error) as exc:
            raise SessionStoreError("Session registry is temporarily unavailable") from exc
        return row is not None and row[0] == expires_at and expires_at > now

    def revoke(self, session_id: str) -> None:
        try:
            with self._connect() as connection:
                connection.execute(
                    "DELETE FROM active_sessions WHERE session_hash = ?",
                    (self._digest(session_id),),
                )
        except (OSError, sqlite3.Error) as exc:
            raise SessionStoreError("Session registry is temporarily unavailable") from exc

    def list_active(self, current_session_id: str) -> list[ActiveDevice]:
        now = int(self.clock())
        current_hash = self._digest(current_session_id)
        try:
            with self._connect() as connection:
                rows = connection.execute(
                    "SELECT management_id, device_label, created_at, expires_at, session_hash "
                    "FROM active_sessions WHERE expires_at > ? "
                    "ORDER BY expires_at DESC, created_at DESC",
                    (now,),
                ).fetchall()
        except (OSError, sqlite3.Error) as exc:
            raise SessionStoreError("Session registry is temporarily unavailable") from exc
        return [
            ActiveDevice(
                management_id=row[0],
                device_label=row[1],
                created_at=row[2],
                expires_at=row[3],
                current=secrets.compare_digest(row[4], current_hash),
            )
            for row in rows
        ]

    def revoke_managed(self, management_id: str) -> bool:
        try:
            with self._connect() as connection:
                cursor = connection.execute(
                    "DELETE FROM active_sessions WHERE management_id = ?", (management_id,)
                )
                return cursor.rowcount == 1
        except (OSError, sqlite3.Error) as exc:
            raise SessionStoreError("Session registry is temporarily unavailable") from exc

    def revoke_all(self) -> None:
        try:
            with self._connect() as connection:
                connection.execute("DELETE FROM active_sessions")
        except (OSError, sqlite3.Error) as exc:
            raise SessionStoreError("Session registry is temporarily unavailable") from exc