admin / Synapse-Cortex

import uuid

from fastapi import APIRouter, Depends, Form, Request
from fastapi.responses import HTMLResponse, RedirectResponse
from fastapi.templating import Jinja2Templates
from sqlalchemy.orm import Session

from .. import mfa, models
from ..audit import log_audit
from ..database import get_db
from ..paths import TEMPLATES_DIR
from ..security import verify_password

router = APIRouter()
templates = Jinja2Templates(directory=str(TEMPLATES_DIR))


def _establish_session(request: Request, db: Session, user: models.User) -> RedirectResponse:
    """Grant the full session and log the successful login. If the tenant admin
    has flagged this account as needing MFA but it isn't enrolled yet, send
    them straight to enrollment instead of the dashboard."""
    request.session["user_id"] = str(user.id)
    log_audit(db, tenant_id=user.tenant_id, user_id=user.id, action="login_success")
    db.commit()
    if user.mfa_required and not user.mfa_enabled:
        return RedirectResponse("/settings/mfa/enroll", status_code=303)
    return RedirectResponse("/", status_code=303)


@router.get("/login", response_class=HTMLResponse)
def login_form(request: Request):
    return templates.TemplateResponse("login.html", {"request": request, "error": None})


@router.post("/login")
def login_submit(
    request: Request,
    db: Session = Depends(get_db),
    email: str = Form(...),
    password: str = Form(...),
):
    user = db.query(models.User).filter(models.User.email == email.strip().lower()).first()
    if not user or not user.is_active or not verify_password(password, user.hashed_password):
        # Only log the attempt when it maps to a real (tenant-scoped) account -
        # unknown emails have no tenant to attribute the row to.
        if user:
            log_audit(db, tenant_id=user.tenant_id, user_id=user.id, action="login_failed")
            db.commit()
        return templates.TemplateResponse(
            "login.html",
            {"request": request, "error": "Invalid credentials."},
            status_code=401,
        )

    if user.mfa_enabled:
        # Password verified but the session isn't established yet - a second
        # factor is still required before we trust this login.
        request.session["pending_mfa_user_id"] = str(user.id)
        return templates.TemplateResponse("mfa_challenge.html", {"request": request, "error": None})

    return _establish_session(request, db, user)


@router.post("/login/mfa")
def login_mfa_submit(request: Request, db: Session = Depends(get_db), code: str = Form(...)):
    pending_id = request.session.get("pending_mfa_user_id")
    if not pending_id:
        return RedirectResponse("/login", status_code=303)

    user = db.get(models.User, uuid.UUID(pending_id))
    if not user or not user.mfa_enabled or not mfa.verify_code(user.mfa_secret, code):
        return templates.TemplateResponse(
            "mfa_challenge.html",
            {"request": request, "error": "Invalid or expired code."},
            status_code=401,
        )

    request.session.pop("pending_mfa_user_id", None)
    return _establish_session(request, db, user)


@router.get("/logout")
def logout(request: Request):
    request.session.clear()
    return RedirectResponse("/login", status_code=303)