agentic.email/backend/app/services/login_mfa.py

"""Login-time MFA challenges and OTP verification."""

from __future__ import annotations

import re
import secrets
from datetime import datetime, timedelta, timezone

from fastapi import HTTPException, status
from sqlalchemy.ext.asyncio import AsyncSession

from app.models import Account
from app.services.org_security import get_or_create_security_settings
from app.services.urmi_otp import (
    INDIA_COUNTRY_CODE,
    OTP_TTL_MINUTES,
    _hash_otp,
    parse_india_phone,
    send_phone_otp,
)

CHALLENGE_TTL_MINUTES = 10
MAX_VERIFY_ATTEMPTS = 5


def phone_hint(phone_number: str) -> str:
    digits = re.sub(r"\D", "", phone_number)
    if len(digits) >= 4:
        return f"•••• {digits[-4:]}"
    return "your phone"


async def account_requires_mfa(session: AsyncSession, account: Account) -> bool:
    if account.mfa_enabled:
        return True
    if not account.company_domain:
        return False
    org = await get_or_create_security_settings(session, account.company_domain)
    if org.require_mfa_all_users:
        return True
    if org.require_mfa_admins and account.org_role in {"admin", "owner"}:
        return True
    return False


def mfa_phone_parts(account: Account) -> tuple[str, str]:
    try:
        return parse_india_phone(account.phone_number)
    except ValueError as error:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=str(error),
        ) from error


async def create_login_challenge(
    database,
    account: Account,
) -> tuple[str, str]:
    country_code, local_phone = mfa_phone_parts(account)
    challenge_token = secrets.token_urlsafe(32)
    otp, _ = await send_phone_otp(
        database,
        country_code,
        local_phone,
        purpose="login_mfa",
        purpose_key=challenge_token,
    )
    now = datetime.now(timezone.utc)
    await database.login_mfa_challenges.update_many(
        {"account_id": str(account.id), "consumed": False},
        {"$set": {"consumed": True}},
    )
    await database.login_mfa_challenges.insert_one(
        {
            "challenge_token": challenge_token,
            "account_id": str(account.id),
            "phone_key": f"{country_code}{local_phone}",
            "otp_hash": _hash_otp(otp),
            "attempts": 0,
            "consumed": False,
            "sent_at": now,
            "created_at": now,
            "expires_at": now + timedelta(minutes=CHALLENGE_TTL_MINUTES),
        }
    )
    return challenge_token, phone_hint(account.phone_number)


async def resend_login_challenge_otp(database, challenge_token: str) -> str:
    record = await _active_challenge(database, challenge_token)
    country_code, local_phone = record["phone_key"][:2], record["phone_key"][2:]
    otp, _ = await send_phone_otp(
        database,
        country_code,
        local_phone,
        purpose="login_mfa",
        purpose_key=challenge_token,
    )
    now = datetime.now(timezone.utc)
    await database.login_mfa_challenges.update_one(
        {"_id": record["_id"]},
        {
            "$set": {
                "otp_hash": _hash_otp(otp),
                "attempts": 0,
                "sent_at": now,
                "expires_at": now + timedelta(minutes=CHALLENGE_TTL_MINUTES),
            }
        },
    )
    return phone_hint(f"+{record['phone_key']}")


async def verify_login_challenge(
    database,
    challenge_token: str,
    account: Account,
    otp: str,
) -> None:
    record = await _active_challenge(database, challenge_token)
    if record["account_id"] != str(account.id):
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Verification session is invalid. Sign in again.",
        )
    code = otp.strip()
    if not re.fullmatch(r"\d{6}", code):
        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
            detail="Enter the 6-digit verification code.",
        )
    if record.get("attempts", 0) >= MAX_VERIFY_ATTEMPTS:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="Too many incorrect attempts. Sign in again to receive a new code.",
        )
    if _hash_otp(code) != record["otp_hash"]:
        await database.login_mfa_challenges.update_one(
            {"_id": record["_id"]},
            {"$inc": {"attempts": 1}},
        )
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="Incorrect verification code.",
        )
    await database.login_mfa_challenges.update_one(
        {"_id": record["_id"]},
        {"$set": {"consumed": True, "verified_at": datetime.now(timezone.utc)}},
    )


async def _active_challenge(database, challenge_token: str) -> dict:
    now = datetime.now(timezone.utc)
    record = await database.login_mfa_challenges.find_one(
        {"challenge_token": challenge_token, "consumed": False},
    )
    if record is None or record.get("expires_at") <= now:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="Verification session expired. Sign in again.",
        )
    return record