govoplan-core/src/govoplan_core/security/sessions.py

from __future__ import annotations

from dataclasses import dataclass
from datetime import timedelta

from sqlalchemy.orm import Session

from govoplan_core.db.models import (
    Account,
    AuthSession,
    Group,
    GroupRoleAssignment,
    Role,
    SystemRoleAssignment,
    Tenant,
    User,
    UserGroupMembership,
    UserRoleAssignment,
)
from govoplan_core.access.auth.tokens import generate_secret, hash_secret, verify_secret
from govoplan_core.security.permissions import expand_scopes
from govoplan_core.security.time import ensure_aware_utc, utc_now

SESSION_RANDOM_BYTES = 32
DEFAULT_SESSION_HOURS = 12


@dataclass(slots=True)
class CreatedSession:
    model: AuthSession
    token: str
    csrf_token: str


def generate_session_token() -> str:
    return generate_secret("ms_", random_bytes=SESSION_RANDOM_BYTES)


def hash_session_token(token: str) -> str:
    return hash_secret(token)


def generate_csrf_token() -> str:
    return generate_secret("", random_bytes=SESSION_RANDOM_BYTES)


def hash_csrf_token(token: str) -> str:
    return hash_secret(token)


def verify_session_token(token: str, expected_hash: str) -> bool:
    return verify_secret(token, expected_hash)


def verify_auth_session_csrf(auth_session: AuthSession, csrf_token: str | None) -> bool:
    if not csrf_token or not auth_session.csrf_token_hash:
        return False
    return verify_secret(csrf_token, auth_session.csrf_token_hash)


def create_auth_session(
    session: Session,
    *,
    user: User,
    hours: int = DEFAULT_SESSION_HOURS,
    user_agent: str | None = None,
    ip_address: str | None = None,
) -> CreatedSession:
    if not user.account_id:
        raise ValueError("Tenant membership has no global account.")
    token = generate_session_token()
    csrf_token = generate_csrf_token()
    now = utc_now()
    model = AuthSession(
        tenant_id=user.tenant_id,
        user_id=user.id,
        account_id=user.account_id,
        token_hash=hash_session_token(token),
        csrf_token_hash=hash_csrf_token(csrf_token),
        expires_at=now + timedelta(hours=hours),
        last_seen_at=now,
        user_agent=user_agent,
        ip_address=ip_address,
    )
    user.last_login_at = now
    if user.account:
        user.account.last_login_at = now
        session.add(user.account)
    session.add(model)
    session.add(user)
    session.flush()
    return CreatedSession(model=model, token=token, csrf_token=csrf_token)


def authenticate_session_token(session: Session, token: str) -> AuthSession | None:
    token_hash = hash_session_token(token)
    model = session.query(AuthSession).filter(AuthSession.token_hash == token_hash, AuthSession.revoked_at.is_(None)).one_or_none()
    if not model:
        return None
    now = utc_now()
    expires_at = ensure_aware_utc(model.expires_at)
    if expires_at is None or expires_at < now:
        return None
    model.last_seen_at = now
    session.add(model)
    return model


def revoke_auth_session(session: Session, token: str) -> bool:
    model = authenticate_session_token(session, token)
    if not model:
        return False
    model.revoked_at = utc_now()
    session.add(model)
    return True


def switch_auth_session_tenant(session: Session, auth_session: AuthSession, tenant_id: str) -> User:
    membership = (
        session.query(User)
        .join(Tenant, Tenant.id == User.tenant_id)
        .filter(
            User.account_id == auth_session.account_id,
            User.tenant_id == tenant_id,
            User.is_active.is_(True),
            Tenant.is_active.is_(True),
        )
        .one_or_none()
    )
    if membership is None:
        raise LookupError("The account does not have an active membership in this tenant.")
    auth_session.tenant_id = membership.tenant_id
    auth_session.user_id = membership.id
    auth_session.last_seen_at = utc_now()
    session.add(auth_session)
    session.flush()
    return membership


def collect_direct_user_roles(session: Session, user: User) -> list[Role]:
    return (
        session.query(Role)
        .join(UserRoleAssignment, UserRoleAssignment.role_id == Role.id)
        .filter(
            UserRoleAssignment.user_id == user.id,
            UserRoleAssignment.tenant_id == user.tenant_id,
            Role.tenant_id == user.tenant_id,
        )
        .order_by(Role.name.asc())
        .all()
    )


def collect_user_roles(session: Session, user: User) -> list[Role]:
    roles_by_id: dict[str, Role] = {role.id: role for role in collect_direct_user_roles(session, user)}

    group_roles = (
        session.query(Role)
        .join(GroupRoleAssignment, GroupRoleAssignment.role_id == Role.id)
        .join(UserGroupMembership, UserGroupMembership.group_id == GroupRoleAssignment.group_id)
        .join(Group, Group.id == UserGroupMembership.group_id)
        .filter(
            UserGroupMembership.user_id == user.id,
            UserGroupMembership.tenant_id == user.tenant_id,
            Group.is_active.is_(True),
            Role.tenant_id == user.tenant_id,
        )
        .all()
    )
    for role in group_roles:
        roles_by_id[role.id] = role
    return list(roles_by_id.values())


def collect_system_roles(session: Session, account: Account) -> list[Role]:
    return (
        session.query(Role)
        .join(SystemRoleAssignment, SystemRoleAssignment.role_id == Role.id)
        .filter(SystemRoleAssignment.account_id == account.id, Role.tenant_id.is_(None))
        .order_by(Role.name.asc())
        .all()
    )


def collect_user_groups(session: Session, user: User) -> list[Group]:
    return (
        session.query(Group)
        .join(UserGroupMembership, UserGroupMembership.group_id == Group.id)
        .filter(
            UserGroupMembership.user_id == user.id,
            UserGroupMembership.tenant_id == user.tenant_id,
            Group.is_active.is_(True),
        )
        .order_by(Group.name.asc())
        .all()
    )


def collect_user_scopes(session: Session, user: User, *, include_system: bool = True) -> list[str]:
    scopes: set[str] = set()
    for role in collect_user_roles(session, user):
        scopes.update(role.permissions or [])
    if include_system and user.account:
        for role in collect_system_roles(session, user.account):
            scopes.update(role.permissions or [])
    return expand_scopes(scopes)


def collect_tenant_memberships(session: Session, account: Account) -> list[tuple[User, Tenant]]:
    return (
        session.query(User, Tenant)
        .join(Tenant, Tenant.id == User.tenant_id)
        .filter(
            User.account_id == account.id,
            User.is_active.is_(True),
            Tenant.is_active.is_(True),
        )
        .order_by(Tenant.name.asc())
        .all()
    )