diff --git a/app/api/routes_auth.py b/app/api/routes_auth.py index b009d16..6abfbac 100644 --- a/app/api/routes_auth.py +++ b/app/api/routes_auth.py @@ -30,6 +30,7 @@ from app.core.auth import hash_password, mint_jwt, verify_password from app.core import password_reset from app.core import refresh_tokens +from app.core import totp from app.core.token_revocation import revoke as revoke_jti from app.db.tenants_repository import get_tenants_repository from app.db.users_repository import get_users_repository @@ -97,6 +98,42 @@ class AuthResponse(BaseModel): refresh_token: str principal: AuthPrincipal onboarding_required: bool = False + # Only populated on the response that completes 2FA enrollment - the + # one-time recovery codes, shown to the user once and never again. + recovery_codes: list[str] | None = None + + +class MfaChallengeResponse(BaseModel): + """Returned by /auth/signin when a second factor is required. No session + token is issued yet; the client exchanges ``mfa_token`` (plus a code, or an + enrollment) at /auth/2fa/verify for the real AuthResponse.""" + + mfa_required: bool = True + # False means the user must enrol first (call /auth/2fa/enroll), True means + # they already have an authenticator and should just enter a code. + enrolled: bool + mfa_token: str + + +class MfaEnrollRequest(BaseModel): + mfa_token: str = Field(min_length=1) + + +class MfaEnrollResponse(BaseModel): + secret: str + otpauth_uri: str + issuer: str + account: str + + +class MfaVerifyRequest(BaseModel): + mfa_token: str = Field(min_length=1) + code: str = Field(min_length=1) + + +class MfaStatusResponse(BaseModel): + enabled: bool + required: bool class RefreshRequest(BaseModel): @@ -147,6 +184,10 @@ def _safe_delete_tenant(repo, tenant_id: str) -> None: pass +def _totp_issuer(settings) -> str: + return str(getattr(settings, "totp_issuer", "") or "PetroBrain") + + def _to_principal_payload(record: dict) -> AuthPrincipal: return AuthPrincipal( user_id=record["id"], @@ -228,8 +269,8 @@ def _validate_password(plain: str) -> None: raise HTTPException(status_code=422, detail="password is too long (max 72 bytes)") -@router.post("/signup", response_model=AuthResponse, status_code=201) -async def signup(req: SignupRequest) -> AuthResponse: +@router.post("/signup", status_code=201) +async def signup(req: SignupRequest) -> AuthResponse | MfaChallengeResponse: settings = get_settings() if not settings.enable_self_signup: raise HTTPException(status_code=403, detail="self-serve signup is disabled") @@ -286,6 +327,30 @@ async def signup(req: SignupRequest) -> AuthResponse: _safe_delete_tenant(tenants, tenant_id) raise + # When 2FA is mandatory, a brand-new account does not get a session yet - + # it must enrol an authenticator first, same as signin. The user row exists + # (active) so the challenge can resolve it; the frontend handles the + # enrollment step identically for signup and signin. + if bool(getattr(settings, "require_2fa", False)): + from datetime import timedelta + challenge = totp.mint_challenge_token( + user_id=record["id"], + tenant_id=record["tenant_id"], + jwt_secret=settings.jwt_secret, + ttl=timedelta(minutes=int(getattr(settings, "mfa_challenge_ttl_minutes", 10))), + ) + audit_logger.write(AuditEvent( + event_type="auth_signup", + tenant_id=record["tenant_id"], + user_id=record["id"], + role=record["role"], + route="/auth/signup", + request={"email": record["email"]}, + response={"user_id": record["id"], "mfa_pending": True}, + metadata={"flow": "self_serve", "account_type": account_type}, + )) + return MfaChallengeResponse(enrolled=False, mfa_token=challenge) + token = _mint_for(record) refresh_token = _issue_refresh(record) audit_logger.write(AuditEvent( @@ -307,8 +372,8 @@ async def signup(req: SignupRequest) -> AuthResponse: ) -@router.post("/signin", response_model=AuthResponse) -async def signin(req: SigninRequest) -> AuthResponse: +@router.post("/signin") +async def signin(req: SigninRequest) -> AuthResponse | MfaChallengeResponse: settings = get_settings() users = get_users_repository() record = users.find_by_email_any_tenant(req.email) @@ -338,6 +403,31 @@ async def signin(req: SigninRequest) -> AuthResponse: raise invalid auth_lockout.record_success(tenant_id, req.email) + # Two-factor gate. The password is correct, but we do not issue a session + # yet if 2FA applies. A user who has already enrolled is ALWAYS challenged + # (even if the global flag is off - you can't un-enrol your way past it); + # when PB_REQUIRE_2FA is on, an unenrolled user is sent to enroll first. + enrolled = bool(record.get("totp_enabled")) + if enrolled or bool(getattr(settings, "require_2fa", False)): + from datetime import timedelta + challenge = totp.mint_challenge_token( + user_id=record["id"], + tenant_id=record["tenant_id"], + jwt_secret=settings.jwt_secret, + ttl=timedelta(minutes=int(getattr(settings, "mfa_challenge_ttl_minutes", 10))), + ) + audit_logger.write(AuditEvent( + event_type="auth_signin_mfa_challenge", + tenant_id=record["tenant_id"], + user_id=record["id"], + role=record["role"], + route="/auth/signin", + request={"email": record["email"]}, + response={"enrolled": enrolled}, + metadata={}, + )) + return MfaChallengeResponse(enrolled=enrolled, mfa_token=challenge) + try: users.touch_last_active(tenant_id=tenant_id, user_id=record["id"]) except Exception: @@ -369,6 +459,146 @@ async def signin(req: SigninRequest) -> AuthResponse: ) +def _load_active_user_from_challenge(req_token: str) -> dict: + """Resolve and validate the user behind an MFA challenge token, or 401.""" + settings = get_settings() + expired = HTTPException( + status_code=401, detail="this sign-in step expired, please sign in again" + ) + claims = totp.verify_challenge_token(req_token, jwt_secret=settings.jwt_secret) + if claims is None: + raise expired + record = get_users_repository().get( + tenant_id=str(claims["tenant_id"]), user_id=str(claims["user_id"]) + ) + if record is None or record.get("status") != "active": + raise expired + return record + + +@router.post("/2fa/enroll", response_model=MfaEnrollResponse) +async def enroll_2fa(req: MfaEnrollRequest) -> MfaEnrollResponse: + """Begin TOTP enrollment for the user behind a valid challenge token. + + Generates a fresh secret (stored as pending, not yet enabled) and returns + the otpauth:// URI for the authenticator app. The user then proves a code + at /auth/2fa/verify, which is what actually enables 2FA. + """ + settings = get_settings() + record = _load_active_user_from_challenge(req.mfa_token) + if record.get("totp_enabled"): + raise HTTPException( + status_code=409, + detail="two-factor is already set up; enter a code instead", + ) + secret = totp.generate_secret() + get_users_repository().set_totp_pending( + tenant_id=record["tenant_id"], user_id=record["id"], secret=secret + ) + return MfaEnrollResponse( + secret=secret, + otpauth_uri=totp.provisioning_uri( + secret, account_email=record["email"], issuer=_totp_issuer(settings) + ), + issuer=_totp_issuer(settings), + account=record["email"], + ) + + +@router.post("/2fa/verify", response_model=AuthResponse) +async def verify_2fa(req: MfaVerifyRequest) -> AuthResponse: + """Complete sign-in (or enrollment) by proving a 6-digit code or a recovery + code, exchanging the challenge token for a real access + refresh pair. + + Brute-force protection reuses the per-account lockout, keyed separately from + the password step so the two can't exhaust each other's budgets. + """ + record = _load_active_user_from_challenge(req.mfa_token) + tenant_id = record["tenant_id"] + user_id = record["id"] + users = get_users_repository() + + lock_key = f"{record['email']}:2fa" + if auth_lockout.is_locked(tenant_id, lock_key): + raise HTTPException( + status_code=429, detail="too many attempts, please wait and try again" + ) + invalid = HTTPException(status_code=401, detail="invalid or expired code") + + recovery_codes: list[str] | None = None + enrolling = not bool(record.get("totp_enabled")) + if enrolling: + secret = record.get("totp_secret") + if not secret: + # /auth/2fa/enroll has to run first to provision a secret. + raise HTTPException(status_code=400, detail="start two-factor setup first") + if not totp.verify_totp(secret, req.code): + auth_lockout.record_failure(tenant_id, lock_key) + raise invalid + recovery_codes = totp.generate_recovery_codes() + record = users.enable_totp( + tenant_id=tenant_id, + user_id=user_id, + recovery_code_hashes=totp.hash_recovery_codes(recovery_codes), + ) + else: + if totp.verify_totp(record.get("totp_secret"), req.code): + pass + else: + # Not a valid TOTP - try a one-time recovery code, which is consumed. + remaining = totp.consume_recovery_code( + req.code, list(record.get("totp_recovery_codes") or []) + ) + if remaining is None: + auth_lockout.record_failure(tenant_id, lock_key) + raise invalid + record = users.replace_recovery_codes( + tenant_id=tenant_id, user_id=user_id, recovery_code_hashes=remaining + ) + auth_lockout.record_success(tenant_id, lock_key) + + try: + users.touch_last_active(tenant_id=tenant_id, user_id=user_id) + except Exception: + pass + + token = _mint_for(record) + refresh_token = _issue_refresh(record) + audit_logger.write(AuditEvent( + event_type="auth_2fa_verified", + tenant_id=tenant_id, + user_id=user_id, + role=record["role"], + route="/auth/2fa/verify", + request={"email": record["email"]}, + response={"enrolled_now": enrolling}, + metadata={"method": "totp" if not enrolling else "enrollment"}, + )) + tenant = get_tenants_repository().get(tenant_id) or {} + onboarding_required = ( + (tenant.get("attributes") or {}).get("onboarding_status") != "completed" + and bool((tenant.get("attributes") or {}).get("created_by_signup")) + ) + return AuthResponse( + token=token, + refresh_token=refresh_token, + principal=_to_principal_payload(record), + onboarding_required=onboarding_required, + recovery_codes=recovery_codes, + ) + + +@router.get("/2fa/status", response_model=MfaStatusResponse) +async def mfa_status(who: Principal = Depends(get_principal)) -> MfaStatusResponse: + """Whether the signed-in user has 2FA enabled, and whether it's mandatory.""" + settings = get_settings() + record = get_users_repository().get(tenant_id=who.tenant_id, user_id=who.user_id) + return MfaStatusResponse( + enabled=bool(record and record.get("totp_enabled")), + required=bool(getattr(settings, "require_2fa", False)), + ) + + @router.post("/refresh", response_model=AuthResponse) async def refresh(req: RefreshRequest) -> AuthResponse: """Exchange a valid refresh token for a new access token + a new refresh diff --git a/app/config.py b/app/config.py index 5bd3a98..d63be19 100644 --- a/app/config.py +++ b/app/config.py @@ -232,6 +232,15 @@ class Settings(BaseSettings): auth_lockout_max_failures: int = 5 auth_lockout_window_minutes: int = 15 auth_lockout_minutes: int = 15 + # Two-factor (TOTP). When require_2fa is true every account must enrol in an + # authenticator app and pass a 6-digit code at sign-in. Roll out by deploying + # with this OFF (enrollment available, not enforced), enrolling, then turning + # it ON - flipping it on before anyone has enrolled forces all users through + # enrollment at next login. The challenge token bridges the password step and + # the code step and is intentionally short-lived. + require_2fa: bool = False + mfa_challenge_ttl_minutes: int = 10 + totp_issuer: str = "PetroBrain" # RAG embedding_model: str = "text-embedding-3-large" diff --git a/app/core/totp.py b/app/core/totp.py new file mode 100644 index 0000000..dd5dce3 --- /dev/null +++ b/app/core/totp.py @@ -0,0 +1,131 @@ +""" +Two-factor (TOTP) primitives: authenticator secrets, code verification, +one-time recovery codes, and the short-lived "MFA challenge" token that +bridges the password step and the code step of sign-in. + +Design notes: + * The challenge token is signed with a namespaced derivative of the JWT + secret (``::mfa``) so it can never be replayed as a real access + token - ``app.api.deps.get_principal`` verifies against the bare secret and + would reject this signature outright. It carries only enough to resume the + flow (tenant_id, user_id) and a tight ``exp``. + * Recovery codes are stored as bcrypt hashes, exactly like passwords; the + plaintext is shown to the user once at enrollment and never persisted. + * TOTP verification uses a +/-1 step window to tolerate clock skew. +""" +from __future__ import annotations + +import secrets +from datetime import datetime, timedelta, timezone +from uuid import uuid4 + +import bcrypt +import jwt +import pyotp + + +MFA_TOKEN_PURPOSE = "mfa_challenge" +RECOVERY_CODE_COUNT = 10 + + +def generate_secret() -> str: + """A fresh base32 TOTP secret for a new enrollment.""" + return pyotp.random_base32() + + +def provisioning_uri(secret: str, *, account_email: str, issuer: str) -> str: + """The ``otpauth://`` URI an authenticator app imports (we render it as a + QR on the client). ``issuer`` is the label the app shows, e.g. PetroBrain.""" + return pyotp.TOTP(secret).provisioning_uri(name=account_email, issuer_name=issuer) + + +def verify_totp(secret: str | None, code: str) -> bool: + """True if ``code`` is a currently-valid 6-digit TOTP for ``secret``. + + Tolerates one step of clock skew on either side. Non-digit / wrong-length + input is rejected before hitting pyotp so a junk code can't raise. + """ + if not secret: + return False + cleaned = (code or "").strip().replace(" ", "") + if len(cleaned) != 6 or not cleaned.isdigit(): + return False + try: + return pyotp.TOTP(secret).verify(cleaned, valid_window=1) + except Exception: # noqa: BLE001 - any pyotp failure is a non-match, not a 500 + return False + + +def generate_recovery_codes(count: int = RECOVERY_CODE_COUNT) -> list[str]: + """Human-friendly one-time backup codes, e.g. ``a1b2c3d4-e5f6g7h8``. + + Returned as plaintext for the one-time display; callers must hash them with + :func:`hash_recovery_codes` before storing. + """ + return [f"{secrets.token_hex(4)}-{secrets.token_hex(4)}" for _ in range(count)] + + +def hash_recovery_codes(codes: list[str]) -> list[str]: + return [ + bcrypt.hashpw(code.encode("utf-8"), bcrypt.gensalt()).decode("utf-8") + for code in codes + ] + + +def consume_recovery_code(code: str, hashed: list[str]) -> list[str] | None: + """If ``code`` matches one of the stored bcrypt hashes, return the remaining + hashes (that code burned); otherwise return None. Single-use by construction. + """ + cleaned = (code or "").strip().lower() + if not cleaned: + return None + encoded = cleaned.encode("utf-8") + for index, stored in enumerate(hashed): + try: + if bcrypt.checkpw(encoded, stored.encode("utf-8")): + return [h for i, h in enumerate(hashed) if i != index] + except ValueError: + continue + return None + + +def _mfa_secret(jwt_secret: str) -> str: + return f"{jwt_secret}::mfa" + + +def mint_challenge_token( + *, + user_id: str, + tenant_id: str, + jwt_secret: str, + ttl: timedelta, +) -> str: + """Short-lived token proving the password step passed, pending the code.""" + if not jwt_secret: + raise ValueError("jwt secret is required to mint a challenge token") + now = datetime.now(timezone.utc) + claims = { + "sub": user_id, + "user_id": user_id, + "tenant_id": tenant_id, + "purpose": MFA_TOKEN_PURPOSE, + "iat": now, + "exp": now + ttl, + "jti": str(uuid4()), + } + return jwt.encode(claims, _mfa_secret(jwt_secret), algorithm="HS256") + + +def verify_challenge_token(token: str, *, jwt_secret: str) -> dict | None: + """Return the claims for a valid, unexpired challenge token, else None.""" + if not token: + return None + try: + claims = jwt.decode(token, _mfa_secret(jwt_secret), algorithms=["HS256"]) + except jwt.PyJWTError: + return None + if claims.get("purpose") != MFA_TOKEN_PURPOSE: + return None + if not claims.get("user_id") or not claims.get("tenant_id"): + return None + return claims diff --git a/app/db/migrations/019_user_totp.sql b/app/db/migrations/019_user_totp.sql new file mode 100644 index 0000000..39ef198 --- /dev/null +++ b/app/db/migrations/019_user_totp.sql @@ -0,0 +1,14 @@ +-- Two-factor auth (TOTP) on the users table from migration 004. +-- +-- A user enrols an authenticator app: totp_secret holds their base32 secret, +-- totp_enabled flips true once they prove a code, totp_recovery_codes holds the +-- bcrypt hashes of their one-time backup codes (never the plaintext), and +-- totp_enrolled_utc records when enrollment completed. Existing rows default to +-- not-enrolled; whether they are *forced* to enrol is the PB_REQUIRE_2FA app +-- flag, not a column, so rollout can be toggled without a migration. + +ALTER TABLE users + ADD COLUMN IF NOT EXISTS totp_secret TEXT, + ADD COLUMN IF NOT EXISTS totp_enabled BOOLEAN NOT NULL DEFAULT false, + ADD COLUMN IF NOT EXISTS totp_recovery_codes JSONB NOT NULL DEFAULT '[]'::jsonb, + ADD COLUMN IF NOT EXISTS totp_enrolled_utc TIMESTAMPTZ; diff --git a/app/db/users_repository.py b/app/db/users_repository.py index 0886b67..35e3f7b 100644 --- a/app/db/users_repository.py +++ b/app/db/users_repository.py @@ -38,6 +38,12 @@ class UserRecord: updated_utc: str = "" password_hash: str | None = None password_set_utc: str | None = None + # Two-factor (TOTP). secret is base32; recovery codes are bcrypt hashes, + # never plaintext. enabled flips true only once a code is proven. + totp_secret: str | None = None + totp_enabled: bool = False + totp_recovery_codes: list[str] = field(default_factory=list) + totp_enrolled_utc: str | None = None def as_dict(self) -> dict[str, Any]: return asdict(self) @@ -192,6 +198,36 @@ def set_allowed_assets(self, *, tenant_id: str, user_id: str, allowed_assets: list[str]) -> dict[str, Any]: return self._update(tenant_id, user_id, allowed_assets=list(allowed_assets)) + def set_totp_pending(self, *, tenant_id: str, user_id: str, + secret: str) -> dict[str, Any]: + """Store a not-yet-confirmed TOTP secret (enrollment in progress).""" + return self._update(tenant_id, user_id, totp_secret=secret, totp_enabled=False) + + def enable_totp(self, *, tenant_id: str, user_id: str, + recovery_code_hashes: list[str]) -> dict[str, Any]: + """Confirm enrollment: flip enabled on and store hashed recovery codes.""" + return self._update( + tenant_id, user_id, + totp_enabled=True, + totp_recovery_codes=list(recovery_code_hashes), + totp_enrolled_utc=_now(), + ) + + def replace_recovery_codes(self, *, tenant_id: str, user_id: str, + recovery_code_hashes: list[str]) -> dict[str, Any]: + return self._update( + tenant_id, user_id, totp_recovery_codes=list(recovery_code_hashes) + ) + + def disable_totp(self, *, tenant_id: str, user_id: str) -> dict[str, Any]: + return self._update( + tenant_id, user_id, + totp_secret=None, + totp_enabled=False, + totp_recovery_codes=[], + totp_enrolled_utc=None, + ) + def _update(self, tenant_id: str, user_id: str, **changes: Any) -> dict[str, Any]: with self._lock: rows = self._read_all_locked() @@ -232,7 +268,8 @@ def _write_all_locked(self, rows: list[dict[str, Any]]) -> None: _USER_COLUMNS = ( "id, tenant_id, email, role, status, allowed_assets, " "invited_at_utc, last_active_utc, created_utc, updated_utc, " - "password_hash, password_set_utc" + "password_hash, password_set_utc, " + "totp_secret, totp_enabled, totp_recovery_codes, totp_enrolled_utc" ) @@ -396,6 +433,34 @@ def set_allowed_assets(self, *, tenant_id: str, user_id: str, allowed_assets: list[str]) -> dict[str, Any]: return self._update(tenant_id, user_id, allowed_assets=list(allowed_assets)) + def set_totp_pending(self, *, tenant_id: str, user_id: str, + secret: str) -> dict[str, Any]: + return self._update(tenant_id, user_id, totp_secret=secret, totp_enabled=False) + + def enable_totp(self, *, tenant_id: str, user_id: str, + recovery_code_hashes: list[str]) -> dict[str, Any]: + return self._update( + tenant_id, user_id, + totp_enabled=True, + totp_recovery_codes=list(recovery_code_hashes), + totp_enrolled_utc=_now(), + ) + + def replace_recovery_codes(self, *, tenant_id: str, user_id: str, + recovery_code_hashes: list[str]) -> dict[str, Any]: + return self._update( + tenant_id, user_id, totp_recovery_codes=list(recovery_code_hashes) + ) + + def disable_totp(self, *, tenant_id: str, user_id: str) -> dict[str, Any]: + return self._update( + tenant_id, user_id, + totp_secret=None, + totp_enabled=False, + totp_recovery_codes=[], + totp_enrolled_utc=None, + ) + def _update(self, tenant_id: str, user_id: str, **changes: Any) -> dict[str, Any]: from psycopg.types.json import Json @@ -403,7 +468,9 @@ def _update(self, tenant_id: str, user_id: str, **changes: Any) -> dict[str, Any params: list[Any] = [] for key, value in changes.items(): assignments.append(f"{key} = %s") - params.append(Json(value) if key == "allowed_assets" else value) + params.append( + Json(value) if key in ("allowed_assets", "totp_recovery_codes") else value + ) assignments.append("updated_utc = now()") params.extend([tenant_id, user_id]) sql = ( @@ -445,6 +512,7 @@ def _serialize_row(row: dict[str, Any]) -> dict[str, Any]: "created_utc", "updated_utc", "password_set_utc", + "totp_enrolled_utc", ): value = out.get(key) if value is not None and not isinstance(value, str): @@ -464,6 +532,10 @@ def _record_from_row(row: dict[str, Any]) -> UserRecord: updated_utc=data.get("updated_utc", ""), password_hash=data.get("password_hash"), password_set_utc=data.get("password_set_utc"), + totp_secret=data.get("totp_secret"), + totp_enabled=bool(data.get("totp_enabled")), + totp_recovery_codes=list(data.get("totp_recovery_codes") or []), + totp_enrolled_utc=data.get("totp_enrolled_utc"), ) diff --git a/frontend/apps/web/lib/auth/AuthForm.tsx b/frontend/apps/web/lib/auth/AuthForm.tsx index 699023a..955688f 100644 --- a/frontend/apps/web/lib/auth/AuthForm.tsx +++ b/frontend/apps/web/lib/auth/AuthForm.tsx @@ -10,7 +10,17 @@ import { Logo } from '@petrobrain/ui'; import { useChatStore } from '@/lib/chat/store'; import { useSettingsStore } from '@/lib/chat/settings'; -import { AuthError, signin, signup } from './api'; +import { + AuthError, + enroll2fa, + isMfaChallenge, + signin, + signup, + verify2fa, + type AuthResponse, + type MfaChallenge, + type MfaEnrollData, +} from './api'; export type AuthMode = 'signin' | 'signup'; @@ -91,6 +101,34 @@ export function AuthForm({ mode }: AuthFormProps) { // knows the sign-in request is still in progress. const [slow, setSlow] = useState(false); const slowTimerRef = useRef | null>(null); + // Set when the password step succeeds but a second factor is still required; + // swaps the credentials card for the 2FA step. recoveryCodes/pendingRoute are + // the one-time codes shown right after enrollment, before we route on. + const [challenge, setChallenge] = useState(null); + const [recoveryCodes, setRecoveryCodes] = useState(null); + const [pendingRoute, setPendingRoute] = useState(null); + + function completeWithSession(res: AuthResponse) { + const cleanedEmail = email.trim(); + setSession(res.token, res.refresh_token, { + ...res.principal, + email: res.principal.email || cleanedEmail, + }); + clearSessionExpired(); + if (mode === 'signup' && name.trim()) setCallMeName(name.trim()); + if (mode === 'signup' && typeof window !== 'undefined') { + sessionStorage.removeItem('petrobrain-signup-account-type'); + } + const target = (res.onboarding_required ? '/onboarding' : '/chat') as Route; + // Recovery codes only come back on the turn that completes enrollment; show + // them once (the user must save them) before routing into the app. + if (res.recovery_codes && res.recovery_codes.length > 0) { + setRecoveryCodes(res.recovery_codes); + setPendingRoute(target); + } else { + router.push(target); + } + } useEffect(() => { return () => { @@ -147,16 +185,15 @@ export function AuthForm({ mode }: AuthFormProps) { ...(accountType ? { account_type: accountType } : {}), }, controller.signal) : await signin(apiBaseUrl, { email: cleanedEmail, password }, controller.signal); - setSession(res.token, res.refresh_token, { ...res.principal, email: res.principal.email || cleanedEmail }); - clearSessionExpired(); - // Capture the signup name in the same settings field the sidebar pill - // already prefers over the auto-generated user_id, so the user sees - // their name straight after creating the account. - if (mode === 'signup' && name.trim()) { - setCallMeName(name.trim()); + if (isMfaChallenge(res)) { + // Password accepted, second factor still required. Swap to the 2FA step; + // the credentials are not kept past this point. + setChallenge(res); + setBusy(false); + setSlow(false); + return; } - if (mode === 'signup') sessionStorage.removeItem('petrobrain-signup-account-type'); - router.push((res.onboarding_required ? '/onboarding' : '/chat') as Route); + completeWithSession(res); } catch (err) { if (err instanceof AuthError) { setError(err.message); @@ -197,10 +234,23 @@ export function AuthForm({ mode }: AuthFormProps) { {copy.title}

- {copy.subtitle} + {recoveryCodes ? 'Save your recovery codes' : challenge ? 'Two-factor authentication' : copy.subtitle}

+ {recoveryCodes ? ( + { if (pendingRoute) router.push(pendingRoute); }} + /> + ) : challenge ? ( + + ) : ( + <>
+ + )} ); } + +/** + * Second-factor step shown after the password is accepted. If the user is not + * yet enrolled it provisions a TOTP secret and shows the manual setup key plus + * an "open in authenticator" link (better than a QR on a mobile-only device, + * where you can't scan your own screen); enrolled users just enter a code. + */ +function TwoFactorStep({ + baseUrl, + challenge, + onComplete, +}: { + baseUrl: string; + challenge: MfaChallenge; + onComplete: (res: AuthResponse) => void; +}) { + const [enrollData, setEnrollData] = useState(null); + const [loadingEnroll, setLoadingEnroll] = useState(!challenge.enrolled); + const [code, setCode] = useState(''); + const [busy, setBusy] = useState(false); + const [err, setErr] = useState(null); + const [copied, setCopied] = useState(false); + + useEffect(() => { + if (challenge.enrolled) return; + let active = true; + enroll2fa(baseUrl, challenge.mfa_token) + .then((data) => { if (active) setEnrollData(data); }) + .catch((e) => { + if (active) { + setErr(e instanceof AuthError ? e.message : 'Could not start two-factor setup.'); + } + }) + .finally(() => { if (active) setLoadingEnroll(false); }); + return () => { active = false; }; + }, [baseUrl, challenge]); + + async function submitCode(e: FormEvent) { + e.preventDefault(); + if (busy) return; + const cleaned = code.trim(); + if (!cleaned) { setErr('Enter the code from your authenticator app.'); return; } + setBusy(true); + setErr(null); + try { + const res = await verify2fa(baseUrl, { mfa_token: challenge.mfa_token, code: cleaned }); + onComplete(res); + } catch (e) { + setErr(e instanceof AuthError ? e.message : 'Could not verify the code. Try again.'); + setBusy(false); + } + } + + async function copySecret() { + if (!enrollData) return; + try { + await navigator.clipboard.writeText(enrollData.secret); + setCopied(true); + window.setTimeout(() => setCopied(false), 1600); + } catch { + // Clipboard can be blocked; the key is visible to type manually anyway. + } + } + + const inputCls = + 'h-11 w-full rounded-xl border border-neutral-200 bg-white px-3.5 text-center text-lg tracking-[0.3em] shadow-[0_1px_2px_rgba(15,23,42,0.04)] transition-all hover:border-primary-300 focus:border-primary-400 focus:outline-none focus:ring-2 focus:ring-primary-200 dark:border-neutral-700 dark:bg-neutral-900 dark:text-neutral-100 dark:placeholder-neutral-500'; + + return ( + + {!challenge.enrolled ? ( +
+

+ Set up an authenticator app (Google Authenticator, Authy, 1Password). Add this + account, then enter the 6-digit code it shows. +

+ {loadingEnroll ? ( +

Preparing setup…

+ ) : enrollData ? ( +
+ + Open in authenticator app + +

+ Or enter this key manually +

+
+ + {enrollData.secret} + + +
+
+ ) : null} +
+ ) : ( +

+ Enter the 6-digit code from your authenticator app. You can also use one of your + recovery codes. +

+ )} + +
+ + setCode(e.target.value)} + placeholder="123456" + className={inputCls} + /> +
+ + {err ? ( +

+ {err} +

+ ) : null} + + + + ); +} + +/** + * One-time recovery codes shown immediately after enrollment. The user must + * save these; they are the only way back in if the authenticator is lost. + */ +function RecoveryCodesCard({ + codes, + onContinue, +}: { + codes: string[]; + onContinue: () => void; +}) { + const [ack, setAck] = useState(false); + + function copyAll() { + void navigator.clipboard?.writeText(codes.join('\n')).catch(() => {}); + } + + return ( +
+

+ Save these recovery codes somewhere safe. Each works once if you lose access to your + authenticator. They will not be shown again. +

+
    + {codes.map((c) => ( +
  • {c}
    • + ))} +
+ + + +
+ ); +} diff --git a/frontend/apps/web/lib/auth/api.ts b/frontend/apps/web/lib/auth/api.ts index 4c4b7af..c18c34e 100644 --- a/frontend/apps/web/lib/auth/api.ts +++ b/frontend/apps/web/lib/auth/api.ts @@ -20,6 +20,33 @@ export interface AuthResponse { refresh_token: string; principal: AuthPrincipalPayload; onboarding_required?: boolean; + // Present only on the response that completes 2FA enrollment - the one-time + // recovery codes, shown to the user once. + recovery_codes?: string[] | null; +} + +/** + * Returned by /auth/signin and /auth/signup when a second factor is required. + * No session is issued yet; the client enrols (if needed) then posts a code to + * /auth/2fa/verify to exchange `mfa_token` for a real AuthResponse. + */ +export interface MfaChallenge { + mfa_required: true; + enrolled: boolean; + mfa_token: string; +} + +export interface MfaEnrollData { + secret: string; + otpauth_uri: string; + issuer: string; + account: string; +} + +export type AuthResult = AuthResponse | MfaChallenge; + +export function isMfaChallenge(result: AuthResult): result is MfaChallenge { + return (result as MfaChallenge).mfa_required === true; } export interface AuthErrorBody { @@ -35,12 +62,24 @@ export class AuthError extends Error { } } +/** Pull `{detail}` off an error body, falling back to the status line. */ +async function errorDetail(res: Response): Promise { + let detail = `${res.status} ${res.statusText}`; + try { + const data = (await res.json()) as AuthErrorBody; + if (data && typeof data.detail === 'string') detail = data.detail; + } catch { + // Body wasn't JSON; keep the status-line fallback. + } + return detail; +} + async function postAuth( baseUrl: string, path: '/auth/signup' | '/auth/signin', body: { email: string; password: string; account_type?: 'individual' | 'company'; full_name?: string }, signal?: AbortSignal, -): Promise { +): Promise { const res = await fetch(`${baseUrl}${path}`, { method: 'POST', headers: { 'Content-Type': 'application/json' }, @@ -49,26 +88,17 @@ async function postAuth( }); if (res.ok) { - return (await res.json()) as AuthResponse; + // Either a full session or a 2FA challenge - the caller disambiguates. + return (await res.json()) as AuthResult; } - - // Pull a friendlier message off `{detail: "..."}` when the backend supplied - // one, fall back to the HTTP status otherwise. - let detail = `${res.status} ${res.statusText}`; - try { - const data = (await res.json()) as AuthErrorBody; - if (data && typeof data.detail === 'string') detail = data.detail; - } catch { - // Body wasn't JSON; keep the status-line fallback. - } - throw new AuthError(detail, res.status); + throw new AuthError(await errorDetail(res), res.status); } export function signup( baseUrl: string, body: { email: string; password: string; account_type?: 'individual' | 'company'; full_name?: string }, signal?: AbortSignal, -): Promise { +): Promise { return postAuth(baseUrl, '/auth/signup', body, signal); } @@ -76,10 +106,48 @@ export function signin( baseUrl: string, body: { email: string; password: string }, signal?: AbortSignal, -): Promise { +): Promise { return postAuth(baseUrl, '/auth/signin', body, signal); } +/** + * Begin TOTP enrollment using a challenge token from signin/signup. Returns the + * authenticator secret + otpauth URI; the user then proves a code via verify2fa. + */ +export async function enroll2fa( + baseUrl: string, + mfaToken: string, + signal?: AbortSignal, +): Promise { + const res = await fetch(`${baseUrl}/auth/2fa/enroll`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ mfa_token: mfaToken }), + ...(signal ? { signal } : {}), + }); + if (res.ok) return (await res.json()) as MfaEnrollData; + throw new AuthError(await errorDetail(res), res.status); +} + +/** + * Complete sign-in (or enrollment) by submitting a 6-digit code or a recovery + * code with the challenge token. Returns the real session on success. + */ +export async function verify2fa( + baseUrl: string, + body: { mfa_token: string; code: string }, + signal?: AbortSignal, +): Promise { + const res = await fetch(`${baseUrl}/auth/2fa/verify`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify(body), + ...(signal ? { signal } : {}), + }); + if (res.ok) return (await res.json()) as AuthResponse; + throw new AuthError(await errorDetail(res), res.status); +} + /** * Kick off a password reset for an email. The backend is enumeration-safe and * always answers 200 with a neutral message, so this resolves with that message diff --git a/requirements-tierb.txt b/requirements-tierb.txt index 5bc12d0..85941f1 100644 --- a/requirements-tierb.txt +++ b/requirements-tierb.txt @@ -19,6 +19,7 @@ httpx==0.28.1 pytest==9.0.3 pytest-asyncio==1.4.0 PyJWT==2.13.0 +pyotp==2.9.0 sentence-transformers==5.5.1 celery==5.5.3 boto3==1.40.59 diff --git a/requirements.txt b/requirements.txt index c83813d..74c9ef8 100644 --- a/requirements.txt +++ b/requirements.txt @@ -15,6 +15,7 @@ pytest-asyncio==1.4.0 PyJWT==2.13.0 cryptography==48.0.1 bcrypt==4.2.0 +pyotp==2.9.0 sentence-transformers==5.5.1 celery==5.5.3 boto3==1.40.59 diff --git a/tests/test_2fa.py b/tests/test_2fa.py new file mode 100644 index 0000000..6a24bb1 --- /dev/null +++ b/tests/test_2fa.py @@ -0,0 +1,205 @@ +"""Two-factor (TOTP) flow tests for /auth: enrollment, code login, recovery +codes, enforcement, and lockout.""" +import os +import sys +from types import SimpleNamespace + +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +import pyotp +import pytest +from fastapi.testclient import TestClient + +from app.api import deps, routes_auth +from app.db.tenants_repository import LocalJsonTenantsRepository +from app.db.users_repository import LocalJsonUsersRepository +from app.main import app +from tests.auth_helpers import JWT_AUDIENCE, JWT_ISSUER, JWT_SECRET + + +client = TestClient(app) + + +def _settings(**overrides): + base = { + "jwt_secret": JWT_SECRET, + "jwt_public_key": "", + "jwt_issuer": JWT_ISSUER, + "jwt_audience": JWT_AUDIENCE, + "jwt_ttl_hours": 1, + "refresh_token_ttl_days": 14, + "enable_self_signup": True, + "default_signup_tenant_id": "demo", + "default_signup_role": "engineer", + "password_min_length": 8, + "require_2fa": False, + "mfa_challenge_ttl_minutes": 10, + "totp_issuer": "PetroBrain", + } + base.update(overrides) + return SimpleNamespace(**base) + + +@pytest.fixture +def users_repo(tmp_path): + return LocalJsonUsersRepository(tmp_path / "users.jsonl") + + +@pytest.fixture +def tenants_repo(tmp_path): + return LocalJsonTenantsRepository(tmp_path / "tenants.jsonl") + + +# Mutable holder so individual tests can flip require_2fa for the same wiring. +_STATE: dict = {} + + +@pytest.fixture(autouse=True) +def wire(monkeypatch, users_repo, tenants_repo): + _STATE["require_2fa"] = False + + def current_settings(): + return _settings(require_2fa=_STATE["require_2fa"]) + + monkeypatch.setattr(deps, "get_settings", current_settings) + monkeypatch.setattr(routes_auth, "get_settings", current_settings) + monkeypatch.setattr(routes_auth, "get_users_repository", lambda: users_repo) + monkeypatch.setattr(routes_auth, "get_tenants_repository", lambda: tenants_repo) + from app.core import auth_lockout + auth_lockout.reset_for_tests() + from app.core.http_hardening import clear_rate_limits + clear_rate_limits() + + +def _code_for(secret: str) -> str: + return pyotp.TOTP(secret).now() + + +def _signup(email="user@example.com", password="correcthorse1"): + return client.post("/auth/signup", json={"email": email, "password": password}) + + +def test_enrollment_completes_and_returns_recovery_codes_and_session(): + _STATE["require_2fa"] = True + # Signup now returns a challenge instead of a session. + s = _signup() + assert s.status_code == 201, s.text + body = s.json() + assert body["mfa_required"] is True + assert body["enrolled"] is False + mfa_token = body["mfa_token"] + assert "token" not in body # no session issued yet + + # Enroll -> get a secret + otpauth URI. + enroll = client.post("/auth/2fa/enroll", json={"mfa_token": mfa_token}) + assert enroll.status_code == 200, enroll.text + secret = enroll.json()["secret"] + assert enroll.json()["otpauth_uri"].startswith("otpauth://totp/") + + # Verify with a real code -> full session + one-time recovery codes. + verify = client.post( + "/auth/2fa/verify", json={"mfa_token": mfa_token, "code": _code_for(secret)} + ) + assert verify.status_code == 200, verify.text + vb = verify.json() + assert vb["token"] + assert vb["refresh_token"] + assert isinstance(vb["recovery_codes"], list) and len(vb["recovery_codes"]) == 10 + # The session token works against a protected route (decodes to a principal). + me = client.get("/auth/me", headers={"Authorization": f"Bearer {vb['token']}"}) + assert me.status_code == 200 + + +def test_enrolled_user_is_challenged_on_signin_and_logs_in_with_code(): + _STATE["require_2fa"] = True + # Enroll first. + mfa_token = _signup("bob@example.com").json()["mfa_token"] + secret = client.post("/auth/2fa/enroll", json={"mfa_token": mfa_token}).json()["secret"] + client.post("/auth/2fa/verify", json={"mfa_token": mfa_token, "code": _code_for(secret)}) + + # Now sign in: password is right but a session is NOT issued; a challenge is. + signin = client.post( + "/auth/signin", json={"email": "bob@example.com", "password": "correcthorse1"} + ) + assert signin.status_code == 200, signin.text + sb = signin.json() + assert sb["mfa_required"] is True + assert sb["enrolled"] is True + assert "token" not in sb + + verify = client.post( + "/auth/2fa/verify", + json={"mfa_token": sb["mfa_token"], "code": _code_for(secret)}, + ) + assert verify.status_code == 200 + assert verify.json()["token"] + + +def test_enrolled_user_is_challenged_even_when_flag_off(): + # Enroll while required. + _STATE["require_2fa"] = True + mfa_token = _signup("carol@example.com").json()["mfa_token"] + secret = client.post("/auth/2fa/enroll", json={"mfa_token": mfa_token}).json()["secret"] + client.post("/auth/2fa/verify", json={"mfa_token": mfa_token, "code": _code_for(secret)}) + + # Turn the global flag off: an already-enrolled user must still pass 2FA. + _STATE["require_2fa"] = False + signin = client.post( + "/auth/signin", json={"email": "carol@example.com", "password": "correcthorse1"} + ) + assert signin.json().get("mfa_required") is True + assert signin.json()["enrolled"] is True + + +def test_recovery_code_logs_in_and_is_single_use(): + _STATE["require_2fa"] = True + mfa_token = _signup("dave@example.com").json()["mfa_token"] + secret = client.post("/auth/2fa/enroll", json={"mfa_token": mfa_token}).json()["secret"] + codes = client.post( + "/auth/2fa/verify", json={"mfa_token": mfa_token, "code": _code_for(secret)} + ).json()["recovery_codes"] + one = codes[0] + + # Sign in and use a recovery code instead of a TOTP. + ch = client.post( + "/auth/signin", json={"email": "dave@example.com", "password": "correcthorse1"} + ).json()["mfa_token"] + r1 = client.post("/auth/2fa/verify", json={"mfa_token": ch, "code": one}) + assert r1.status_code == 200, r1.text + + # The same recovery code cannot be reused. + ch2 = client.post( + "/auth/signin", json={"email": "dave@example.com", "password": "correcthorse1"} + ).json()["mfa_token"] + r2 = client.post("/auth/2fa/verify", json={"mfa_token": ch2, "code": one}) + assert r2.status_code == 401 + + +def test_wrong_code_is_rejected(): + _STATE["require_2fa"] = True + mfa_token = _signup("erin@example.com").json()["mfa_token"] + client.post("/auth/2fa/enroll", json={"mfa_token": mfa_token}) + bad = client.post( + "/auth/2fa/verify", json={"mfa_token": mfa_token, "code": "000000"} + ) + assert bad.status_code == 401 + + +def test_challenge_token_cannot_be_used_as_a_session_token(): + _STATE["require_2fa"] = True + mfa_token = _signup("mallory@example.com").json()["mfa_token"] + # The challenge token must NOT authenticate against a protected route. + me = client.get("/auth/me", headers={"Authorization": f"Bearer {mfa_token}"}) + assert me.status_code == 401 + + +def test_signin_without_2fa_is_unchanged_when_flag_off(): + _STATE["require_2fa"] = False + _signup("frank@example.com") + r = client.post( + "/auth/signin", json={"email": "frank@example.com", "password": "correcthorse1"} + ) + assert r.status_code == 200 + body = r.json() + assert body["token"] + assert "mfa_required" not in body