import secrets import string from datetime import datetime, timezone from typing import Annotated from fastapi import APIRouter, Depends, HTTPException from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from app.api.endpoints.auth import _mint_session_tokens from app.core.admin_database import get_admin_db from app.core.config import settings from app.models.account import Account from app.models.account_invite import AccountInvite from app.models.oauth_identity import OAuthIdentity from app.models.user import User from app.schemas.oauth import OAuthCallbackPayload, OAuthCallbackResponse from app.services.billing import BillingService from app.services.oauth_providers import ( google_exchange_code, microsoft_exchange_code, OAuthProfile, ) router = APIRouter(prefix="/auth", tags=["auth-oauth"]) def _generate_display_code(length: int = 8) -> str: """Match the helper used by /auth/register — A-Z + 0-9, length 8.""" alphabet = string.ascii_uppercase + string.digits return "".join(secrets.choice(alphabet) for _ in range(length)) async def _sign_in_or_register( db: AsyncSession, provider: str, profile: OAuthProfile, *, account_invite_code: str | None = None, invited_email: str | None = None, ) -> tuple[User, bool]: """Returns (user, is_new_user). Idempotent on (provider, provider_subject). When ``account_invite_code`` is supplied (from the /accept-invite flow), a brand-new user is created inside the invited account instead of getting a personal account + Pro trial. Mismatch between the OAuth profile email and ``invited_email`` raises ``invite_email_mismatch`` per the spec contract that mirrors the email+password register path. """ identity = ( await db.execute( select(OAuthIdentity).where( OAuthIdentity.provider == provider, OAuthIdentity.provider_subject == profile.provider_subject, ) ) ).scalar_one_or_none() if identity: user = ( await db.execute(select(User).where(User.id == identity.user_id)) ).scalar_one() return user, False user = ( await db.execute(select(User).where(User.email == profile.email)) ).scalar_one_or_none() is_new_user = user is None # If the user arrived via an invite link but already has a ResolutionFlow # account (e.g., previously signed up with email+password), silently # linking the OAuth identity to that existing account would bypass the # invite — they'd stay in their personal account and the invite would # never be consumed. Fail loud instead so they can sign in and accept the # invite from the dashboard. The "invited user wants to transfer accounts" # case is a v2 concern. if account_invite_code and not is_new_user: raise HTTPException( status_code=400, detail={ "error": "email_already_registered_use_login", "message": ( "An account already exists for this email. Please sign in " "instead, then accept the invite from your dashboard." ), }, ) invite_record: AccountInvite | None = None if is_new_user and account_invite_code: # SELECT FOR UPDATE so two concurrent OAuth callbacks can't both # consume the same invite code. invite_record = ( await db.execute( select(AccountInvite) .where(AccountInvite.code == account_invite_code) .with_for_update() ) ).scalar_one_or_none() if invite_record is None or not invite_record.is_valid: raise HTTPException( status_code=400, detail={"error": "invite_invalid_or_expired_or_revoked"}, ) # Verify the OAuth profile email matches what was invited. We compare # against the invite row directly (source of truth), but also accept # the client-supplied invited_email as a defensive equality check. if invite_record.email.lower() != profile.email.lower(): raise HTTPException( status_code=400, detail={"error": "invite_email_mismatch"}, ) if invited_email and invited_email.lower() != invite_record.email.lower(): raise HTTPException( status_code=400, detail={"error": "invite_email_mismatch"}, ) if is_new_user: if invite_record is not None: # Join the invited account directly — no personal account, no # trial creation. user = User( email=profile.email, name=profile.name, password_hash=None, account_id=invite_record.account_id, account_role=invite_record.role, role="engineer", email_verified_at=datetime.now(timezone.utc), ) db.add(user) await db.flush() invite_record.accepted_by_id = user.id invite_record.used_at = datetime.now(timezone.utc) await db.flush() else: account = Account( name=f"{profile.name}'s Account", display_code=_generate_display_code(), ) db.add(account) await db.flush() user = User( email=profile.email, name=profile.name, password_hash=None, account_id=account.id, account_role="owner", role="engineer", email_verified_at=datetime.now(timezone.utc), ) db.add(user) await db.flush() account.owner_id = user.id await db.flush() # start_trial commits internally; flushed account/user above. await BillingService.start_trial(db, account.id) db.add( OAuthIdentity( user_id=user.id, provider=provider, provider_subject=profile.provider_subject, provider_email_at_link=profile.email, ) ) await db.commit() await db.refresh(user) return user, is_new_user @router.post("/google/callback", response_model=OAuthCallbackResponse) async def google_callback( payload: OAuthCallbackPayload, db: Annotated[AsyncSession, Depends(get_admin_db)], ) -> OAuthCallbackResponse: if not settings.GOOGLE_CLIENT_ID: raise HTTPException(status_code=503, detail="Google sign-in not configured") redirect_uri = f"{settings.OAUTH_REDIRECT_BASE}/auth/google/callback" profile = await google_exchange_code(payload.code, redirect_uri) user, is_new = await _sign_in_or_register( db, "google", profile, account_invite_code=payload.account_invite_code, invited_email=payload.invited_email, ) token = await _mint_session_tokens(user, db) await db.commit() return OAuthCallbackResponse( access_token=token.access_token, refresh_token=token.refresh_token, is_new_user=is_new, idle_expires_at=token.idle_expires_at, absolute_expires_at=token.absolute_expires_at, ) @router.post("/microsoft/callback", response_model=OAuthCallbackResponse) async def microsoft_callback( payload: OAuthCallbackPayload, db: Annotated[AsyncSession, Depends(get_admin_db)], ) -> OAuthCallbackResponse: if not settings.MS_CLIENT_ID: raise HTTPException(status_code=503, detail="Microsoft sign-in not configured") redirect_uri = f"{settings.OAUTH_REDIRECT_BASE}/auth/microsoft/callback" profile = await microsoft_exchange_code(payload.code, redirect_uri) user, is_new = await _sign_in_or_register( db, "microsoft", profile, account_invite_code=payload.account_invite_code, invited_email=payload.invited_email, ) token = await _mint_session_tokens(user, db) await db.commit() return OAuthCallbackResponse( access_token=token.access_token, refresh_token=token.refresh_token, is_new_user=is_new, idle_expires_at=token.idle_expires_at, absolute_expires_at=token.absolute_expires_at, )