evo-sync/web/auth/session.py

from fastapi import HTTPException
from fastapi.responses import RedirectResponse
from sqlalchemy.orm import Session
from starlette.requests import Request

from web.models.user import User, UserRoleEnum, UserStatusEnum


def get_session_user_id(request: Request) -> int | None:
    return request.session.get("user_id")


def get_current_user(request: Request, db: Session) -> User:
    user_id = get_session_user_id(request)
    if not user_id:
        raise HTTPException(status_code=307, headers={"Location": "/login"})
    user = db.get(User, user_id)
    if not user or user.status == UserStatusEnum.suspended:
        request.session.clear()
        raise HTTPException(status_code=307, headers={"Location": "/login"})
    return user


def get_viewed_user(request: Request, db: Session) -> tuple[User, User]:
    """Return (real_user, viewed_user).

    Admins/system users can view another user's data by having
    `viewed_user_id` set in the session (via /admin/users/{id}/view-as).
    For regular users, both values are the same.
    """
    real_user = get_current_user(request, db)
    is_admin = real_user.role in (UserRoleEnum.admin, UserRoleEnum.system)
    viewed_id = request.session.get("viewed_user_id") if is_admin else None
    if viewed_id:
        viewed = db.get(User, viewed_id)
        if viewed:
            return real_user, viewed
    return real_user, real_user


def login_redirect() -> RedirectResponse:
    return RedirectResponse("/login", status_code=303)