DataOps-platform/app/core/system/identity_sessions.py

"""Server-side session ledger with opaque rotating refresh credentials."""

from __future__ import annotations

import hashlib
import secrets
from collections.abc import Callable
from dataclasses import dataclass
from datetime import UTC, datetime, timedelta
from typing import Any

from app.core.common.identifiers import new_governance_uid
from app.core.system.enterprise_identity import IdentityPolicyError
from app.core.system.tokens import issue_access_token


@dataclass(frozen=True)
class SessionCredentials:
    session_uid: str
    access_token: str
    refresh_token: str


class SessionManager:
    def __init__(self, repository: Any, *, secret: str, clock: Callable[[], datetime] | None = None,
                 max_sessions: int = 5, refresh_lifetime: timedelta = timedelta(hours=8)) -> None:
        self.repository = repository
        self.secret = secret
        self.clock = clock or (lambda: datetime.now(UTC))
        self.max_sessions = max_sessions
        self.refresh_lifetime = refresh_lifetime

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

    def _access(self, record: dict[str, Any]) -> str:
        return issue_access_token(user_id=record["user_uid"], roles=list(record["roles"]), secret=self.secret,
                                  now=self.clock(), session_uid=record["uid"], token_version=record["token_version"],
                                  identity_source=record["identity_source"])

    def create(self, *, subject: str, user_uid: str, roles: list[str], identity_source: str,
               token_version: int = 1, provider_uid: str | None = None,
               emergency_request_uid: str | None = None,
               emergency_expires_at: datetime | None = None, commit: bool = True) -> SessionCredentials:
        if identity_source not in {"oidc", "local", "emergency"}:
            raise IdentityPolicyError("unsupported identity source")
        if identity_source == "oidc" and not provider_uid:
            raise IdentityPolicyError("OIDC session requires provider binding")
        if identity_source == "emergency" and (not emergency_request_uid or not emergency_expires_at):
            raise IdentityPolicyError("emergency session requires request and expiry binding")
        now = self.clock()
        if identity_source == "emergency" and not self.repository.emergency_session_active(
                str(emergency_request_uid), user_uid, now):
            raise IdentityPolicyError("emergency account is no longer an active local administrator")
        refresh = secrets.token_urlsafe(48)
        record = {"uid": new_governance_uid(), "family_uid": new_governance_uid(), "subject": subject,
                  "user_uid": user_uid, "roles": sorted(set(roles)), "identity_source": identity_source,
                  "provider_uid": provider_uid, "emergency_request_uid": emergency_request_uid,
                  "token_version": token_version, "refresh_hash": self._digest(refresh), "status": "active",
                  "created_at": now, "last_seen_at": now,
                  "expires_at": min(now + self.refresh_lifetime, emergency_expires_at) if emergency_expires_at else now + self.refresh_lifetime,
                  "rotated_from_uid": None, "risk_reason": None}
        self.repository.create_session(record, max_sessions=self.max_sessions, commit=commit)
        return SessionCredentials(record["uid"], self._access(record), refresh)

    def refresh(self, refresh_token: str) -> SessionCredentials:
        record = self.repository.get_session_by_refresh(self._digest(refresh_token))
        if not record:
            raise IdentityPolicyError("unknown refresh credential")
        if record["identity_source"] == "emergency" and not self.repository.emergency_session_active(
                record["emergency_request_uid"], record["user_uid"], self.clock()):
            self.repository.revoke_family(record["family_uid"], reason="emergency_window_closed")
            raise IdentityPolicyError("emergency access window is not active")
        if record["status"] != "active":
            self.repository.revoke_family(record["family_uid"], reason="refresh_reuse")
            raise IdentityPolicyError("refresh reuse detected")
        if self.clock() >= record["expires_at"]:
            self.repository.update_session(record["uid"], status="revoked", revoke_reason="expired")
            raise IdentityPolicyError("refresh credential expired")
        refresh = secrets.token_urlsafe(48)
        next_expiry = self.clock() + self.refresh_lifetime
        if record["identity_source"] == "emergency":
            next_expiry = min(record["expires_at"], next_expiry)
        next_record = {**record, "uid": new_governance_uid(), "refresh_hash": self._digest(refresh),
                       "status": "active", "created_at": self.clock(), "last_seen_at": self.clock(),
                       "expires_at": next_expiry, "rotated_from_uid": record["uid"]}
        outcome = self.repository.rotate_refresh(self._digest(refresh_token), next_record)
        if outcome == "unknown":
            raise IdentityPolicyError("unknown refresh credential")
        if outcome == "reuse":
            raise IdentityPolicyError("refresh reuse detected")
        return SessionCredentials(next_record["uid"], self._access(next_record), refresh)

    def revoke(self, session_uid: str, *, reason: str = "logout") -> None:
        if self.repository.get_session(session_uid):
            self.repository.update_session(session_uid, status="revoked", revoke_reason=reason)

    def risk_revoke(self, session_uid: str, *, reason: str) -> None:
        self.repository.update_session(session_uid, status="revoked", revoke_reason="risk", risk_reason=reason)

    def revoke_subject(self, provider_uid: str | None, subject: str, *, reason: str) -> None:
        self.repository.revoke_subject_sessions(provider_uid, subject, reason=reason)