import hashlib import hmac as hmac_mod import re import time import uuid from decimal import Decimal from typing import Optional from fastapi import APIRouter, Depends, HTTPException, Request, Response from fastapi.responses import JSONResponse from pydantic import BaseModel from sqlalchemy import func, select from sqlalchemy.exc import IntegrityError from sqlalchemy.orm import Session from app.dependencies.admin_auth import require_admin_jwt from app.models.database import ( Conversation, CreditEvent, Tenant, TenantAuthProvider, User, get_db, ) from app.services.google_auth import ( decode_full_jwt, decode_partial_jwt, issue_full_jwt, issue_partial_jwt, lookup_tenant_for_google, register_tenant_for_google, verify_google_token, ) from app.models.smart_models import KnowledgeBase from app.services.totp_service import ( consume_backup_code, encode_backup_codes, encrypt_secret, generate_backup_codes, generate_totp_secret, hash_backup_code, make_provisioning_uri, verify_totp_code, ) from app.utils.api_key import generate_api_key from app.utils.config import get_config from app.utils.registration_rate_limit import enforce_registration_rate_limit config = get_config() router = APIRouter(prefix="/admin/auth", tags=["admin-auth"]) _TOTP_RE = re.compile(r"^\d{6}$") _BACKUP_RE = re.compile(r"^[0-9a-f]{16}$") class GoogleAuthRequest(BaseModel): id_token: str class GoogleAuthResponse(BaseModel): totp_enabled: bool partial_token: Optional[str] = None registered: bool = False # B12 Layer 1: True when the Google identity is unknown and no tenant was # created — the consumer must ask the user to confirm before registering. new_account: bool = False class TOTPSetupResponse(BaseModel): provisioning_uri: str backup_codes: list[str] # plaintext — shown once, not retrievable again class TOTPVerifyRequest(BaseModel): code: str class TelegramAuthRequest(BaseModel): id: int first_name: str last_name: Optional[str] = None username: Optional[str] = None photo_url: Optional[str] = None auth_date: int hash: str phone_number: Optional[str] = None def _verify_telegram_hmac(data: dict, bot_token: str) -> bool: """Verify Telegram Login Widget data using HMAC-SHA256.""" received_hash = data.get("hash", "") check_pairs = sorted( (k, str(v)) for k, v in data.items() if k != "hash" and v is not None ) data_check_string = "\n".join(f"{k}={v}" for k, v in check_pairs) secret_key = hashlib.sha256(bot_token.encode()).digest() expected_hash = hmac_mod.new( secret_key, data_check_string.encode(), hashlib.sha256 ).hexdigest() return hmac_mod.compare_digest(expected_hash, received_hash) def _set_access_cookie(response: Response, tenant_id: uuid.UUID) -> None: """Issue a full JWT and set it as an HttpOnly cookie on the response.""" access_token = issue_full_jwt(tenant_id, config.jwt_secret_key) response.set_cookie( key="access_token", value=access_token, httponly=True, secure=config.env.name == "production", samesite="strict", path="/api", max_age=86400, ) def _resolve_enrollment_tenant_id(request: Request) -> tuple[uuid.UUID, bool]: """Resolve the tenant for a TOTP enrollment request via dual auth (B11). Two callers reach the enrollment endpoints: - Login-flow gate: a partial token (Authorization: Bearer, scope=totp_pending), issued by google_auth when 2FA is already enabled. - Settings enrollment: an already-logged-in tenant holding a full admin JWT in the HttpOnly access_token cookie (scope=admin) — JS cannot read it, so it can never be sent as a Bearer header. Returns (tenant_id, from_full_cookie). Partial token is tried first; on its absence/invalidity we fall back to the full cookie. Raises 401 if neither authenticates. Scope checks inside decode_* keep a full token from passing as a partial one and vice versa. """ auth_header = request.headers.get("Authorization", "") if auth_header.startswith("Bearer "): token = auth_header.removeprefix("Bearer ") try: tenant_id_str = decode_partial_jwt(token, config.jwt_secret_key) return uuid.UUID(tenant_id_str), False except (ValueError, AttributeError): pass # fall through to the full-cookie path cookie_token = request.cookies.get("access_token") if cookie_token: try: tenant_id_str = decode_full_jwt(cookie_token, config.jwt_secret_key) return uuid.UUID(tenant_id_str), True except (ValueError, AttributeError): pass raise HTTPException(status_code=401, detail="Unauthorized") @router.post("/telegram") async def telegram_auth( body: TelegramAuthRequest, request: Request, db: Session = Depends(get_db) ) -> JSONResponse: # 1. Guard: bot token must be configured if not config.telegram_bot_token: raise HTTPException(status_code=401, detail="Unauthorized") # 2. Verify HMAC data = body.model_dump(exclude_none=True) if not _verify_telegram_hmac(data, config.telegram_bot_token): raise HTTPException(status_code=401, detail="Unauthorized") # 3. Check auth_date freshness (24 hours) if time.time() - body.auth_date > 86400: raise HTTPException(status_code=401, detail="Unauthorized") telegram_user_id = str(body.id) phone = body.phone_number # 4. Direct match via auth_providers — telegram_id is the only lookup key. # Phone matching was removed with the pre-provisioning model (B8): it # carried a phone-recycling risk and contradicted self-registration. ap_row = db.scalars( select(TenantAuthProvider).where( TenantAuthProvider.provider == "telegram", TenantAuthProvider.provider_user_id == telegram_user_id, ) ).first() registered = False if ap_row: tenant = db.get(Tenant, ap_row.tenant_id) if not tenant or not tenant.is_active: raise HTTPException(status_code=401, detail="Unauthorized") else: # 5. B8 self-registration: unknown Telegram identity creates a tenant. # Telegram provides no email, so there is no cross-provider signal # to collide on — collision detection reduces to the direct match. client_ip = request.client.host if request.client else "unknown" enforce_registration_rate_limit(client_ip) name = body.first_name + (f" {body.last_name}" if body.last_name else "") _, hashed_key = generate_api_key() tenant = Tenant( name=name, api_key=hashed_key, is_active=True, billing_mode="credits", balance_usd=Decimal("0"), phone=phone, ) db.add(tenant) db.flush() db.add(TenantAuthProvider( tenant_id=tenant.id, provider="telegram", provider_user_id=telegram_user_id, )) try: db.commit() registered = True except IntegrityError: # B14: a concurrent request registered the same Telegram identity # between our lookup and insert. The (provider, provider_user_id) # unique constraint rejected our duplicate — recover by logging into # the tenant the winner created instead of creating a second one. db.rollback() ap_row = db.scalars( select(TenantAuthProvider).where( TenantAuthProvider.provider == "telegram", TenantAuthProvider.provider_user_id == telegram_user_id, ) ).first() if ap_row is None: # Not a lost race — a genuine integrity failure. Don't swallow it. raise tenant = db.get(Tenant, ap_row.tenant_id) if not tenant or not tenant.is_active: raise HTTPException(status_code=401, detail="Unauthorized") # 6. Store phone on tenant if provided and not yet set if phone and not tenant.phone: tenant.phone = phone db.commit() # 7. Issue full JWT via cookie (set cookie on the returned JSONResponse, not the injected Response) result = JSONResponse(content={"ok": True, "registered": registered}) _set_access_cookie(result, tenant.id) return result def _check_google_collision(email: str, db: Session) -> None: """Raise 409 if a non-Google account already owns this email (B12 guard). A Telegram-only tenant whose owner_email matches must not be silently bound to a Google identity — surface the collision so the consumer routes the user to link instead. """ pre_tenant = db.scalars(select(Tenant).where(Tenant.owner_email == email)).first() if pre_tenant and pre_tenant.google_id is None: existing_ap = db.scalars( select(TenantAuthProvider).where( TenantAuthProvider.tenant_id == pre_tenant.id, TenantAuthProvider.provider != "google", ) ).first() if existing_ap: raise HTTPException( status_code=409, detail={"collision": True, "existing_provider": existing_ap.provider}, ) def _finalize_google_login( tenant: Tenant, google_id: str, registered: bool, response: Response, db: Session ) -> GoogleAuthResponse: """Register the Google auth-provider row on first login, then issue the cookie (no 2FA) or a partial token (2FA enabled). Shared by login + register.""" has_google_ap = db.scalars( select(TenantAuthProvider).where( TenantAuthProvider.tenant_id == tenant.id, TenantAuthProvider.provider == "google", ) ).first() if not has_google_ap: db.add(TenantAuthProvider( tenant_id=tenant.id, provider="google", provider_user_id=google_id, )) db.commit() if not tenant.totp_enabled: _set_access_cookie(response, tenant.id) return GoogleAuthResponse(totp_enabled=False, registered=registered) partial_token = issue_partial_jwt(tenant.id, config.jwt_secret_key) return GoogleAuthResponse( totp_enabled=True, partial_token=partial_token, registered=registered ) @router.post("/google", response_model=GoogleAuthResponse) async def google_auth( body: GoogleAuthRequest, request: Request, response: Response, db: Session = Depends(get_db), ) -> GoogleAuthResponse: try: google_payload = verify_google_token(body.id_token, config.google_client_id) except ValueError: raise HTTPException(status_code=401, detail="Unauthorized") email = google_payload["email"] google_id = google_payload["google_id"] _check_google_collision(email, db) tenant = lookup_tenant_for_google(email=email, google_id=google_id, db=db) # B12 Layer 1: unknown Google identity is NOT auto-registered. Signal a new # account (no cookie, no tenant created) so the consumer can require an # explicit "create workspace" confirmation before calling /google/register. if tenant is None: return GoogleAuthResponse(totp_enabled=False, registered=False, new_account=True) return _finalize_google_login(tenant, google_id, False, response, db) @router.post("/google/register", response_model=GoogleAuthResponse) async def google_register( body: GoogleAuthRequest, request: Request, response: Response, db: Session = Depends(get_db), ) -> GoogleAuthResponse: """Confirmed registration for a new Google identity (B12 Layer 1). Reached only after the user explicitly confirms "create a new workspace" in response to google_auth's new_account signal. Idempotent: if a tenant already exists for the identity (e.g. confirm clicked twice), it logs in rather than creating a duplicate.""" try: google_payload = verify_google_token(body.id_token, config.google_client_id) except ValueError: raise HTTPException(status_code=401, detail="Unauthorized") email = google_payload["email"] google_id = google_payload["google_id"] profile_name = google_payload.get("name") _check_google_collision(email, db) tenant = lookup_tenant_for_google(email=email, google_id=google_id, db=db) registered = False if tenant is None: client_ip = request.client.host if request.client else "unknown" try: tenant = register_tenant_for_google( email=email, google_id=google_id, name=profile_name, db=db, client_ip=client_ip ) registered = True except IntegrityError: # B13: a concurrent request registered this identity between our # lookup and insert (TOCTOU race). owner_email/google_id are unique, # so the DB rejected our duplicate — recover by logging into the # tenant the winner created instead of surfacing a 500. db.rollback() tenant = lookup_tenant_for_google(email=email, google_id=google_id, db=db) if tenant is None: # Not a lost race — a genuine integrity failure. Don't swallow it. raise return _finalize_google_login(tenant, google_id, registered, response, db) @router.post("/link/telegram") async def link_telegram( body: TelegramAuthRequest, tenant_id: uuid.UUID = Depends(require_admin_jwt), db: Session = Depends(get_db), ) -> JSONResponse: if not config.telegram_bot_token: raise HTTPException(status_code=401, detail="Unauthorized") data = body.model_dump(exclude_none=True) if not _verify_telegram_hmac(data, config.telegram_bot_token): raise HTTPException(status_code=401, detail="Unauthorized") if time.time() - body.auth_date > 86400: raise HTTPException(status_code=401, detail="Unauthorized") telegram_user_id = str(body.id) # Already linked to this tenant if db.scalars( select(TenantAuthProvider).where( TenantAuthProvider.tenant_id == tenant_id, TenantAuthProvider.provider == "telegram", ) ).first(): raise HTTPException(status_code=409, detail={"already_linked": True}) # Telegram ID already used by another tenant if db.scalars( select(TenantAuthProvider).where( TenantAuthProvider.provider == "telegram", TenantAuthProvider.provider_user_id == telegram_user_id, TenantAuthProvider.tenant_id != tenant_id, ) ).first(): raise HTTPException(status_code=409, detail={"already_linked": True, "used_by_other": True}) db.add(TenantAuthProvider( tenant_id=tenant_id, provider="telegram", provider_user_id=telegram_user_id, )) tenant = db.get(Tenant, tenant_id) if body.phone_number and tenant and not tenant.phone: tenant.phone = body.phone_number db.commit() return JSONResponse(content={"linked": True}) @router.post("/link/google") async def link_google( body: GoogleAuthRequest, tenant_id: uuid.UUID = Depends(require_admin_jwt), db: Session = Depends(get_db), ) -> JSONResponse: try: google_payload = verify_google_token(body.id_token, config.google_client_id) except ValueError: raise HTTPException(status_code=401, detail="Unauthorized") google_id = google_payload["google_id"] email = google_payload["email"] # Already linked to this tenant if db.scalars( select(TenantAuthProvider).where( TenantAuthProvider.tenant_id == tenant_id, TenantAuthProvider.provider == "google", ) ).first(): raise HTTPException(status_code=409, detail={"already_linked": True}) # Google ID already used by another tenant if db.scalars( select(TenantAuthProvider).where( TenantAuthProvider.provider == "google", TenantAuthProvider.provider_user_id == google_id, TenantAuthProvider.tenant_id != tenant_id, ) ).first(): raise HTTPException(status_code=409, detail={"already_linked": True, "used_by_other": True}) db.add(TenantAuthProvider( tenant_id=tenant_id, provider="google", provider_user_id=google_id, )) tenant = db.get(Tenant, tenant_id) if tenant: if not tenant.owner_email: tenant.owner_email = email if not tenant.google_id: tenant.google_id = google_id db.commit() return JSONResponse(content={"linked": True}) @router.post("/totp/setup", response_model=TOTPSetupResponse) @router.post("/totp/enroll", response_model=TOTPSetupResponse) async def totp_setup(request: Request, db: Session = Depends(get_db)) -> TOTPSetupResponse: tenant_id, _ = _resolve_enrollment_tenant_id(request) tenant = db.get(Tenant, tenant_id) if not tenant: raise HTTPException(status_code=401, detail="Unauthorized") if tenant.totp_enabled: raise HTTPException(status_code=403, detail="TOTP already configured") raw_secret = generate_totp_secret() plain_codes = generate_backup_codes() tenant.totp_secret = encrypt_secret(raw_secret, config.fernet_key) tenant.backup_codes = encode_backup_codes([hash_backup_code(c) for c in plain_codes]) tenant.totp_enabled = False db.commit() return TOTPSetupResponse( provisioning_uri=make_provisioning_uri(raw_secret, tenant.owner_email or ""), backup_codes=plain_codes, ) @router.post("/totp/verify") async def totp_verify(body: TOTPVerifyRequest, request: Request, db: Session = Depends(get_db)) -> JSONResponse: tenant_id, from_full_cookie = _resolve_enrollment_tenant_id(request) tenant = db.get(Tenant, tenant_id) if not tenant or not tenant.totp_secret: raise HTTPException(status_code=401, detail="Unauthorized") code = body.code.strip().lower() verified = False if _TOTP_RE.match(code): verified = verify_totp_code(tenant.totp_secret, code, config.fernet_key) elif _BACKUP_RE.match(code) and tenant.backup_codes: updated = consume_backup_code(tenant.backup_codes, code) if updated is not None: tenant.backup_codes = encode_backup_codes(updated) verified = True if not verified: raise HTTPException(status_code=401, detail="Unauthorized") tenant.totp_enabled = True db.commit() # Settings enrollment (full-cookie path): the tenant already holds a valid # admin cookie — just confirm enrollment, do not re-issue it. if from_full_cookie: return JSONResponse(content={"ok": True}) # Login-flow gate (partial-token path): issue the full admin cookie now. access_token = issue_full_jwt(tenant.id, config.jwt_secret_key) response = JSONResponse(content={"ok": True}) # SameSite=Strict provides CSRF protection without a CSRF token response.set_cookie( key="access_token", value=access_token, httponly=True, secure=config.env.name == "production", samesite="strict", path="/api", max_age=86400, ) return response @router.delete("/totp") async def totp_disable( body: TOTPVerifyRequest, tenant_id: uuid.UUID = Depends(require_admin_jwt), db: Session = Depends(get_db), ) -> JSONResponse: tenant = db.get(Tenant, tenant_id) if not tenant: raise HTTPException(status_code=401, detail="Unauthorized") if not tenant.totp_enabled or not tenant.totp_secret: raise HTTPException(status_code=400, detail="TOTP not configured") code = body.code.strip().lower() verified = False if _TOTP_RE.match(code): verified = verify_totp_code(tenant.totp_secret, code, config.fernet_key) elif _BACKUP_RE.match(code) and tenant.backup_codes: updated = consume_backup_code(tenant.backup_codes, code) if updated is not None: tenant.backup_codes = encode_backup_codes(updated) verified = True if not verified: raise HTTPException(status_code=401, detail="Unauthorized") tenant.totp_enabled = False tenant.totp_secret = None tenant.backup_codes = None db.commit() return JSONResponse(content={"totp_enabled": False}) @router.delete("/workspace") async def delete_workspace( response: Response, tenant_id: uuid.UUID = Depends(require_admin_jwt), db: Session = Depends(get_db), ) -> JSONResponse: """B12 Layer 3: self-service deletion of an EMPTY workspace. Lets a user stranded with an accidental orphan account clear it themselves so they can link/reuse the identity. Only an empty workspace qualifies: zero balance, no KB entries, no conversations. A workspace with real data is refused (409) — bulk deletion of real content is out of scope here. The guard also refuses when chat users or credit-event history exist: those rows carry NOT-NULL tenant_id FKs with no cascade, so deleting past them would raise a database FK error. Refusing keeps deletion safe and matches the "truly empty" intent (a never-used orphan has none of these). """ tenant = db.get(Tenant, tenant_id) if not tenant: raise HTTPException(status_code=401, detail="Unauthorized") def _count(model) -> int: return db.scalar( select(func.count()).select_from(model).where(model.tenant_id == tenant_id) ) kb_count = _count(KnowledgeBase) conv_count = _count(Conversation) user_count = _count(User) credit_event_count = _count(CreditEvent) if ( tenant.balance_usd != 0 or kb_count > 0 or conv_count > 0 or user_count > 0 or credit_event_count > 0 ): raise HTTPException( status_code=409, detail={ "not_empty": True, "balance_usd": str(tenant.balance_usd), "kb_entries": kb_count, "conversations": conv_count, "users": user_count, "credit_events": credit_event_count, }, ) db.delete(tenant) # auth_providers cascade via delete-orphan db.commit() # Clear the now-stale session cookie for the deleted workspace. result = JSONResponse(content={"deleted": True}) result.delete_cookie(key="access_token", path="/api") return result @router.post("/logout") async def logout(response: Response) -> dict: response.delete_cookie(key="access_token", path="/api") return {"ok": True}