import secrets import string from datetime import datetime, timezone from typing import Annotated from fastapi import APIRouter, Depends, HTTPException, status 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: # Seat enforcement: re-check at OAuth accept time (race-condition guard). if invite_record.role in ("engineer", "l1_tech"): from app.core.subscriptions import get_account_subscription from app.services.seat_enforcement import check_seat_available sub = await get_account_subscription(invite_record.account_id, db) if sub is not None: acct_result = await db.execute( select(Account).where(Account.id == invite_record.account_id) ) acct = acct_result.scalar_one() seat_result = await check_seat_available(acct, sub, invite_record.role, db) if not seat_result.available: raise HTTPException( status_code=status.HTTP_402_PAYMENT_REQUIRED, detail={ "code": "seat_limit_exceeded", "role": seat_result.role, "current": seat_result.current, "limit": seat_result.limit, "upgrade_url": "/account/billing", }, ) # 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, )