"""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 def assert_mfa_phone(account: Account) -> None: try: 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