Ten31-Portal/backend/ten31portal/auth.py

"""Authentication, session management, and role enforcement."""

from typing import Annotated

from argon2 import PasswordHasher
from argon2.exceptions import VerifyMismatchError
from fastapi import Depends, HTTPException, Request
from sqlmodel import Session, select, col

from ten31portal.database import get_session
from ten31portal.models import (
    CapitalAccountStatement, EntityAccess, User, UserRole, EXTERNAL_ROLES,
)

ph = PasswordHasher()


def hash_password(password: str) -> str:
    return ph.hash(password)


def verify_password(password: str, password_hash: str) -> bool:
    try:
        return ph.verify(password_hash, password)
    except VerifyMismatchError:
        return False


# While an account is on the shared default password it may reach only these endpoints, so the
# forced password change cannot be skipped by calling the API directly (the browser modal is not
# a security control). Everything else is blocked until the password is rotated.
_PASSWORD_CHANGE_ALLOWED = frozenset({
    "/api/auth/me",
    "/api/auth/change-password",
    "/api/auth/logout",
})


def get_current_user(request: Request, session: Session = Depends(get_session)) -> User:
    """FastAPI dependency: extract user from session cookie."""
    user_id = request.session.get("user_id")
    if user_id is None:
        raise HTTPException(status_code=401, detail="Not authenticated")
    user = session.get(User, user_id)
    if user is None or not user.is_active:
        raise HTTPException(status_code=401, detail="Not authenticated")
    if user.must_change_password and request.url.path not in _PASSWORD_CHANGE_ALLOWED:
        raise HTTPException(status_code=403, detail="password_change_required")
    return user


def require_role(*roles: UserRole):
    """Return a dependency that enforces one of the given roles."""
    def checker(user: User = Depends(get_current_user)) -> User:
        if user.role not in roles:
            raise HTTPException(status_code=403, detail="Insufficient permissions")
        return user
    return checker


def require_internal(user: User = Depends(get_current_user)) -> User:
    """Block external (entity-scoped) accounts from internal staff endpoints."""
    if user.role in EXTERNAL_ROLES:
        raise HTTPException(status_code=403, detail="Insufficient permissions")
    return user


def household_user_ids(user: User, session: Session) -> list[int]:
    """All account ids that share this user's login.

    An investor who invests under several legal names has one "primary" account (the login)
    and one or more secondary accounts linked to it via ``primary_account_id``. Signing in as
    the primary should surface every linked name's entities, statements, and documents. For a
    standalone account this is just ``[user.id]``.
    """
    root_id = user.primary_account_id or user.id
    linked = session.exec(
        select(User.id).where(User.primary_account_id == root_id)
    ).all()
    return list({root_id, user.id, *linked})


def accessible_entity_ids(user: User, session: Session) -> set[int] | None:
    """Entity ids an external account may view. None means unrestricted (internal staff)."""
    if user.role not in EXTERNAL_ROLES:
        return None
    rows = session.exec(
        select(EntityAccess.entity_id).where(
            col(EntityAccess.user_id).in_(household_user_ids(user, session))
        )
    ).all()
    return set(rows)


def can_access_entity(user: User, entity_id: int, session: Session) -> bool:
    allowed = accessible_entity_ids(user, session)
    return allowed is None or entity_id in allowed


# Convenience aliases
require_user = get_current_user
require_writer = require_role(
    UserRole.fund_admin, UserRole.cfo, UserRole.approver, UserRole.operations
)
require_approver = require_role(UserRole.approver)  # final sign-off — Managing Partners only
require_audit_reader = require_role(UserRole.approver, UserRole.cfo, UserRole.operations)
# Account administration (create/manage users, documents, capital accounts)
require_internal_admin = require_role(UserRole.approver, UserRole.cfo, UserRole.operations)

# Internal admins plus the external Administrator (fund_administrator). Every endpoint
# using this gate must also call check_administrator_scope for the entity it touches —
# the role alone says nothing about WHICH entities an Administrator may manage.
require_admin = require_role(
    UserRole.approver, UserRole.cfo, UserRole.operations, UserRole.fund_administrator
)
# Entity-record writers: internal writers plus the external Administrator (scope-checked).
require_entity_writer = require_role(
    UserRole.fund_admin, UserRole.cfo, UserRole.approver, UserRole.operations,
    UserRole.fund_administrator,
)


def require_internal_or_administrator(user: User = Depends(get_current_user)) -> User:
    """Read gate for admin screens: any internal role, or an external Administrator
    (managing or view-only).

    Investors are blocked; external calls must still be scope-checked per entity.
    """
    if user.role == UserRole.investor:
        raise HTTPException(status_code=403, detail="Insufficient permissions")
    return user


def check_administrator_scope(user: User, entity_id: int, session: Session) -> None:
    """403 when an external account touches an entity outside their grants.

    Internal roles pass through untouched — their reach is decided by the route's gate.
    """
    if user.role in EXTERNAL_ROLES and not can_access_entity(user, entity_id, session):
        raise HTTPException(status_code=403, detail="No access to this entity")


def visible_investor_ids(user: User, session: Session) -> set[int] | None:
    """Investor ids an external account may see or manage. None means unrestricted (internal).

    An investor is in scope when they already have an access grant or a capital-account
    statement in one of the caller's entities. This is the set an external Administrator is
    allowed to touch; it is derived from rows that predate the current request, so a write
    that would CREATE such a row for a new investor must be scope-checked BEFORE it runs.
    """
    scope = accessible_entity_ids(user, session)
    if scope is None:
        return None
    if not scope:
        return set()
    ids = set(session.exec(
        select(EntityAccess.user_id).where(col(EntityAccess.entity_id).in_(scope))
    ).all())
    ids |= set(session.exec(
        select(CapitalAccountStatement.investor_user_id).where(
            col(CapitalAccountStatement.entity_id).in_(scope)
        )
    ).all())
    return ids


def check_target_investor_scope(admin: User, target_user_id: int, session: Session) -> None:
    """403 unless the caller is internal, or the target is an investor already inside the
    caller's scope. Blocks an external Administrator from pulling an out-of-scope investor
    into their funds (which would otherwise widen their scope and unlock further actions)."""
    allowed = visible_investor_ids(admin, session)
    if allowed is None:
        return
    if target_user_id not in allowed:
        raise HTTPException(status_code=403, detail="No access to this account")