agentic.email/backend/app/routers/auth.py

import re
from datetime import date, datetime, timedelta, timezone
from uuid import UUID

import jwt
from fastapi import APIRouter, Depends, HTTPException, status
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from pwdlib import PasswordHash
from pydantic import EmailStr, TypeAdapter, ValidationError
from sqlalchemy import or_, select
from sqlalchemy.exc import IntegrityError
from sqlalchemy.ext.asyncio import AsyncSession

from app.config import get_settings
from app.database import get_postgres_session
from app.models import Account
from app.schemas import AccountRegister, AccountRead, AuthResponse, LoginRequest
from app.services.domain_provisioning import (
    get_domain_by_name,
    provision_mail_domain,
)

router = APIRouter(prefix="/auth", tags=["authentication"])
password_hash = PasswordHash.recommended()
dummy_password_hash = password_hash.hash("Agmail timing protection value")
bearer_scheme = HTTPBearer(auto_error=False)
settings = get_settings()
email_adapter = TypeAdapter(EmailStr)
domain_pattern = re.compile(
    r"^(?=.{1,253}$)(?:[a-z0-9](?:[a-z0-9-]{0,61}[a-z0-9])?\.)+"
    r"[a-z]{2,63}$"
)


def normalize_phone(value: str) -> str:
    prefix = "+" if value.strip().startswith("+") else ""
    return prefix + "".join(character for character in value if character.isdigit())


def validate_password(value: str) -> None:
    if not (
        re.search(r"[a-z]", value)
        and re.search(r"[A-Z]", value)
        and re.search(r"\d", value)
    ):
        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
            detail="Password must include uppercase, lowercase, and a number.",
        )


def build_email(payload: AccountRegister) -> tuple[str, str | None]:
    domain = "getall.help"
    company_domain = None

    if payload.account_type == "company":
        if not payload.company_domain:
            raise HTTPException(
                status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
                detail="Company domain is required for company accounts.",
            )
        domain = payload.company_domain.strip().lower().removeprefix("@")
        if not domain_pattern.fullmatch(domain):
            raise HTTPException(
                status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
                detail="Enter a valid company domain, such as example.com.",
            )
        company_domain = domain

    email = f"{payload.email_name.lower()}@{domain}"
    try:
        return str(email_adapter.validate_python(email)), company_domain
    except ValidationError as error:
        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
            detail="The resulting email address is invalid.",
        ) from error


def create_access_token(account: Account) -> str:
    expires_at = datetime.now(timezone.utc) + timedelta(
        minutes=settings.access_token_minutes
    )
    return jwt.encode(
        {"sub": str(account.id), "email": account.email, "exp": expires_at},
        settings.jwt_secret,
        algorithm="HS256",
    )


async def get_current_account(
    credentials: HTTPAuthorizationCredentials | None = Depends(bearer_scheme),
    session: AsyncSession = Depends(get_postgres_session),
) -> Account:
    unauthorized = HTTPException(
        status_code=status.HTTP_401_UNAUTHORIZED,
        detail="Your session is invalid or has expired.",
        headers={"WWW-Authenticate": "Bearer"},
    )
    if credentials is None or credentials.scheme.lower() != "bearer":
        raise unauthorized

    try:
        claims = jwt.decode(
            credentials.credentials,
            settings.jwt_secret,
            algorithms=["HS256"],
            options={"require": ["sub", "exp"]},
        )
        account_id = UUID(claims["sub"])
    except (jwt.PyJWTError, KeyError, TypeError, ValueError):
        raise unauthorized

    account = await session.get(Account, account_id)
    if account is None:
        raise unauthorized
    return account


@router.post(
    "/register",
    response_model=AuthResponse,
    status_code=status.HTTP_201_CREATED,
)
async def register(
    payload: AccountRegister,
    session: AsyncSession = Depends(get_postgres_session),
) -> AuthResponse:
    if payload.date_of_birth >= date.today():
        raise HTTPException(
            status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
            detail="Date of birth must be in the past.",
        )
    validate_password(payload.password)
    email, company_domain = build_email(payload)
    phone_number = normalize_phone(payload.phone_number)

    existing = await session.scalar(
        select(Account).where(
            or_(Account.email == email, Account.phone_number == phone_number)
        )
    )
    if existing:
        detail = (
            "That email address is already registered."
            if existing.email == email
            else "That phone number is already registered."
        )
        raise HTTPException(status_code=status.HTTP_409_CONFLICT, detail=detail)

    account = Account(
        account_type=payload.account_type,
        email=email,
        first_name=payload.first_name.strip(),
        last_name=payload.last_name.strip(),
        date_of_birth=payload.date_of_birth,
        phone_number=phone_number,
        company_domain=company_domain,
        password_hash=password_hash.hash(payload.password),
    )
    session.add(account)
    try:
        await session.flush()
        if company_domain:
            existing_domain = await get_domain_by_name(session, company_domain)
            if existing_domain is None:
                await provision_mail_domain(
                    session,
                    domain_name=company_domain,
                    domain_type="company",
                    account_id=account.id,
                )
        await session.commit()
    except IntegrityError as error:
        await session.rollback()
        raise HTTPException(
            status_code=status.HTTP_409_CONFLICT,
            detail="An account with those details already exists.",
        ) from error
    await session.refresh(account)
    return AuthResponse(
        access_token=create_access_token(account),
        account=AccountRead.model_validate(account),
    )


@router.post("/login", response_model=AuthResponse)
async def login(
    payload: LoginRequest,
    session: AsyncSession = Depends(get_postgres_session),
) -> AuthResponse:
    account = await session.scalar(
        select(Account).where(Account.email == str(payload.email).lower())
    )
    password_is_valid = password_hash.verify(
        payload.password,
        account.password_hash if account is not None else dummy_password_hash,
    )
    if account is None or not password_is_valid:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Incorrect email or password.",
        )
    return AuthResponse(
        access_token=create_access_token(account),
        account=AccountRead.model_validate(account),
    )


@router.get("/me", response_model=AccountRead)
async def get_me(account: Account = Depends(get_current_account)) -> Account:
    return account