| 1 | """TOTP two-factor authentication models and schemas.""" |
| 2 | |
| 3 | import datetime |
| 4 | from typing import Optional |
| 5 | |
| 6 | from pydantic import BaseModel |
| 7 | from sqlmodel import JSON |
| 8 | from sqlmodel import Column |
| 9 | from sqlmodel import Field |
| 10 | from sqlmodel import SQLModel |
| 11 | from sqlmodel import Text |
| 12 | |
| 13 | |
| 14 | class UserTOTP(SQLModel, table=True): |
| 15 | """Stores per-user TOTP 2FA configuration.""" |
| 16 | |
| 17 | __tablename__ = "user_totp" |
| 18 | |
| 19 | id: Optional[int] = Field(default=None, primary_key=True) |
| 20 | user_id: int = Field(unique=True, index=True) |
| 21 | |
| 22 | # Fernet-encrypted TOTP secret |
| 23 | secret_enc: str = Field(sa_column=Column(Text, nullable=False)) |
| 24 | |
| 25 | # Not enabled until user verifies with a valid code |
| 26 | enabled: bool = Field(default=False) |
| 27 | |
| 28 | # List of {hash: str, used: bool} — bcrypt hashed backup codes |
| 29 | backup_codes: list = Field(default=[], sa_column=Column(JSON, nullable=False)) |
| 30 | |
| 31 | # Last TOTP counter used — prevents replay of same code |
| 32 | last_used_at: Optional[int] = Field(default=None) |
| 33 | |
| 34 | created_at: datetime.datetime = Field(default_factory=datetime.datetime.utcnow) |
| 35 | |
| 36 | |
| 37 | # ── Pydantic request / response schemas ────────────────────────────────────── |
| 38 | |
| 39 | |
| 40 | class TOTPSetupResponse(BaseModel): |
| 41 | """Returned on setup — contains secret, QR data, and one-time backup codes.""" |
| 42 | |
| 43 | secret: str |
| 44 | otpauth_url: str |
| 45 | qr_data_uri: str |
| 46 | backup_codes: list[str] |
| 47 | message: str = "Scan the QR code with your authenticator app, then verify with a code." |
| 48 | |
| 49 | |
| 50 | class TOTPVerifyRequest(BaseModel): |
| 51 | code: str |
| 52 | |
| 53 | |
| 54 | class TOTPDisableRequest(BaseModel): |
| 55 | code: Optional[str] = None |
| 56 | backup_code: Optional[str] = None |
| 57 | |
| 58 | |
| 59 | class TOTPValidateRequest(BaseModel): |
| 60 | """Used during login — temp_token + TOTP code or backup code.""" |
| 61 | |
| 62 | temp_token: str |
| 63 | code: Optional[str] = None |
| 64 | backup_code: Optional[str] = None |
| 65 | |
| 66 | |
| 67 | class TOTPStatusResponse(BaseModel): |
| 68 | enabled: bool |
| 69 | message: str = "2FA status retrieved" |
| 70 | success: bool = True |
| 71 | |
| 72 | |
| 73 | class TOTPBackupCodesResponse(BaseModel): |
| 74 | backup_codes: list[str] |
| 75 | message: str = "New backup codes generated. Save them — they are shown only once." |
| 76 | success: bool = True |