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 store_refresh_token from app.core.admin_database import get_admin_db from app.core.config import settings from app.core.security import create_access_token, create_refresh_token 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, ) refresh_token_str = create_refresh_token({"sub": str(user.id)}) # Persist the refresh-token JTI so the first /auth/refresh call doesn't # reject this token as "revoked" (the rotation logic requires a row to # mark as used). _sign_in_or_register already committed; this needs a # second commit. await store_refresh_token(db, refresh_token_str, user.id) await db.commit() return OAuthCallbackResponse( access_token=create_access_token({"sub": str(user.id)}), refresh_token=refresh_token_str, is_new_user=is_new, ) @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, ) refresh_token_str = create_refresh_token({"sub": str(user.id)}) # Persist the refresh-token JTI so the first /auth/refresh call doesn't # reject this token as "revoked" (the rotation logic requires a row to # mark as used). _sign_in_or_register already committed; this needs a # second commit. await store_refresh_token(db, refresh_token_str, user.id) await db.commit() return OAuthCallbackResponse( access_token=create_access_token({"sub": str(user.id)}), refresh_token=refresh_token_str, is_new_user=is_new, )