From 5026967cee480dd07d4eff0acc586c0580ae931b Mon Sep 17 00:00:00 2001 From: Tianyu Liu Date: Sun, 21 Jun 2026 21:38:22 +0200 Subject: [PATCH] M4-T04: add AuthUser TOTP columns and auth_recovery_code table --- alembic_app/env.py | 2 +- alembic_app/versions/20260621_08_totp.py | 61 +++++ app/models/auth.py | 22 ++ docs/design/m4-login-hardening.md | 2 +- scripts/app_db_adopt.py | 2 +- tests/test_auth_throttle_model.py | 9 +- tests/test_totp_models.py | 280 +++++++++++++++++++++++ 7 files changed, 373 insertions(+), 5 deletions(-) create mode 100644 alembic_app/versions/20260621_08_totp.py create mode 100644 tests/test_totp_models.py diff --git a/alembic_app/env.py b/alembic_app/env.py index 7760158..c23bde5 100644 --- a/alembic_app/env.py +++ b/alembic_app/env.py @@ -6,7 +6,7 @@ from sqlalchemy import engine_from_config, pool from app.config import get_settings from app.db import Base from app.models.config import AppConfigEntry # noqa: F401 -from app.models.auth import AuthSession, AuthUser # noqa: F401 +from app.models.auth import AuthSession, AuthUser, RecoveryCode # noqa: F401 from app.models.auth_throttle import LoginThrottle # noqa: F401 from app.models.public_ip import PublicIPHistory, PublicIPState # noqa: F401 from app.models.location import Location # noqa: F401 diff --git a/alembic_app/versions/20260621_08_totp.py b/alembic_app/versions/20260621_08_totp.py new file mode 100644 index 0000000..2609811 --- /dev/null +++ b/alembic_app/versions/20260621_08_totp.py @@ -0,0 +1,61 @@ +"""add TOTP fields to auth_users and create auth_recovery_code table + +Revision ID: 20260621_08_totp +Revises: 20260621_07_auth_login_throttle +Create Date: 2026-06-21 00:00:00.000000 +""" + +from typing import Sequence, Union + +import sqlalchemy as sa +from alembic import op + + +revision: str = "20260621_08_totp" +down_revision: Union[str, None] = "20260621_07_auth_login_throttle" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # Add totp_secret (nullable) to auth_users — existing rows get NULL, which is correct. + op.add_column("auth_users", sa.Column("totp_secret", sa.String(length=64), nullable=True)) + # Add totp_enabled (NOT NULL) with server_default="0" so existing rows default to false. + op.add_column( + "auth_users", + sa.Column( + "totp_enabled", + sa.Boolean(), + nullable=False, + server_default="0", + ), + ) + + # Create auth_recovery_code table for one-time TOTP recovery codes. + op.create_table( + "auth_recovery_code", + sa.Column("id", sa.Integer(), autoincrement=True, nullable=False), + sa.Column("user_id", sa.Integer(), nullable=False), + sa.Column("code_hash", sa.String(length=255), nullable=False), + sa.Column("used_at", sa.DateTime(timezone=True), nullable=True), + sa.ForeignKeyConstraint(["user_id"], ["auth_users.id"]), + sa.PrimaryKeyConstraint("id"), + ) + op.create_index( + op.f("ix_auth_recovery_code_user_id"), + "auth_recovery_code", + ["user_id"], + unique=False, + ) + + +def downgrade() -> None: + # Drop auth_recovery_code table first (has FK to auth_users). + op.drop_index(op.f("ix_auth_recovery_code_user_id"), table_name="auth_recovery_code") + op.drop_table("auth_recovery_code") + + # Drop the two TOTP columns from auth_users. + # Use batch_alter_table for SQLite compatibility (alembic's portable column-drop path). + with op.batch_alter_table("auth_users") as batch_op: + batch_op.drop_column("totp_enabled") + batch_op.drop_column("totp_secret") diff --git a/app/models/auth.py b/app/models/auth.py index 08da8e5..6e3c300 100644 --- a/app/models/auth.py +++ b/app/models/auth.py @@ -15,8 +15,12 @@ class AuthUser(Base): is_active: Mapped[bool] = mapped_column(Boolean, nullable=False, default=True) force_password_change: Mapped[bool] = mapped_column(Boolean, nullable=False, default=True) created_at: Mapped[datetime] = mapped_column(DateTime(timezone=True), nullable=False) + # TOTP fields (Phase B, M4-T04) — nullable/false by default so existing users are unaffected + totp_secret: Mapped[str | None] = mapped_column(String(64), nullable=True) + totp_enabled: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False) sessions: Mapped[list["AuthSession"]] = relationship(back_populates="user") + recovery_codes: Mapped[list["RecoveryCode"]] = relationship(back_populates="user") class AuthSession(Base): @@ -31,3 +35,21 @@ class AuthSession(Base): revoked_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True) user: Mapped[AuthUser] = relationship(back_populates="sessions") + + +class RecoveryCode(Base): + """One-time TOTP recovery codes for AuthUser. + + ``code_hash`` stores the Argon2 hash of the plaintext code (plaintext is + returned only at setup time and never stored). ``used_at`` is NULL while + the code is still valid; set to the consumption timestamp when consumed. + """ + + __tablename__ = "auth_recovery_code" + + id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) + user_id: Mapped[int] = mapped_column(ForeignKey("auth_users.id"), nullable=False, index=True) + code_hash: Mapped[str] = mapped_column(String(255), nullable=False) + used_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True), nullable=True) + + user: Mapped[AuthUser] = relationship(back_populates="recovery_codes") diff --git a/docs/design/m4-login-hardening.md b/docs/design/m4-login-hardening.md index 6e50d75..721cd01 100644 --- a/docs/design/m4-login-hardening.md +++ b/docs/design/m4-login-hardening.md @@ -188,7 +188,7 @@ Phase B（TOTP，可选配） - **Reviewer checklist**: 密码可交互输入且不回显；只动 auth 行；幂等可重跑。 ### M4-T04 — AuthUser TOTP 字段 + `auth_recovery_code` 表 `[schema]` -- **Status**: `todo` · **Depends**: none（与 Phase A 并行的 schema，但 T06 依赖 T02） +- **Status**: `done` · **Depends**: none（与 Phase A 并行的 schema，但 T06 依赖 T02） - **Files**: `modify app/models/auth.py`（`AuthUser.totp_secret`/`totp_enabled`）；`create` 模型 `RecoveryCode`；`create alembic_app/versions/_NN_totp.py`；`modify alembic_app/env.py`、`scripts/app_db_adopt.py`；`create tests/test_totp_models.py` - **Steps**: `totp_secret` nullable、`totp_enabled` bool default false；`auth_recovery_code(user_id FK, code_hash, used_at nullable)`；revision 建列 + 表；更新 baseline。 - **Out of scope / 不要碰**: 不写 TOTP 逻辑（T05）。 diff --git a/scripts/app_db_adopt.py b/scripts/app_db_adopt.py index 0b44d93..257b021 100644 --- a/scripts/app_db_adopt.py +++ b/scripts/app_db_adopt.py @@ -15,7 +15,7 @@ if str(PROJECT_ROOT) not in sys.path: from app.config import get_settings -APP_BASELINE_REVISION = "20260621_07_auth_login_throttle" +APP_BASELINE_REVISION = "20260621_08_totp" class AppDatabaseAdoptionError(RuntimeError): diff --git a/tests/test_auth_throttle_model.py b/tests/test_auth_throttle_model.py index 59eafbd..d3b5fa3 100644 --- a/tests/test_auth_throttle_model.py +++ b/tests/test_auth_throttle_model.py @@ -171,12 +171,17 @@ def test_next_allowed_at_persists_when_set(throttle_db): def test_downgrade_removes_table(tmp_path: Path): - """downgrade -1 from head must cleanly drop auth_login_throttle.""" + """Downgrading the auth_login_throttle revision must cleanly drop the table. + + We upgrade to the specific throttle revision (not necessarily head, since later + revisions may have been added), then downgrade one step to verify the table is gone. + """ db_path = tmp_path / "down_test.db" db_url = f"sqlite:///{db_path}" alembic_cfg = _make_app_alembic_config(db_url) + throttle_revision = "20260621_07_auth_login_throttle" - command.upgrade(alembic_cfg, "head") + command.upgrade(alembic_cfg, throttle_revision) engine = create_engine(db_url, connect_args={"check_same_thread": False}) assert "auth_login_throttle" in inspect(engine).get_table_names() engine.dispose() diff --git a/tests/test_totp_models.py b/tests/test_totp_models.py new file mode 100644 index 0000000..8805c9d --- /dev/null +++ b/tests/test_totp_models.py @@ -0,0 +1,280 @@ +"""Tests for M4-T04: AuthUser TOTP fields and auth_recovery_code migration.""" +from __future__ import annotations + +from datetime import datetime, timezone +from pathlib import Path + +import pytest +from alembic import command +from alembic.config import Config +from sqlalchemy import create_engine, inspect, text +from sqlalchemy.orm import Session + +from app.models.auth import AuthUser, RecoveryCode + + +def _make_app_alembic_config(database_url: str) -> Config: + config = Config("alembic_app.ini") + config.set_main_option("sqlalchemy.url", database_url) + return config + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture() +def totp_db(tmp_path: Path): + """Temporary SQLite DB upgraded to the current Alembic head.""" + db_path = tmp_path / "totp_test.db" + db_url = f"sqlite:///{db_path}" + alembic_cfg = _make_app_alembic_config(db_url) + command.upgrade(alembic_cfg, "head") + engine = create_engine(db_url, connect_args={"check_same_thread": False}) + yield engine + engine.dispose() + + +# --------------------------------------------------------------------------- +# Migration shape tests +# --------------------------------------------------------------------------- + + +def test_auth_users_has_totp_columns(totp_db): + """After upgrade to head, auth_users must have totp_secret and totp_enabled columns.""" + inspector = inspect(totp_db) + columns = {col["name"]: col for col in inspector.get_columns("auth_users")} + + assert "totp_secret" in columns, "totp_secret column missing from auth_users" + assert "totp_enabled" in columns, "totp_enabled column missing from auth_users" + + # totp_secret must be nullable + assert columns["totp_secret"]["nullable"], "totp_secret should be nullable" + # totp_enabled must be NOT NULL + assert not columns["totp_enabled"]["nullable"], "totp_enabled should be NOT NULL" + + +def test_auth_recovery_code_table_exists(totp_db): + """auth_recovery_code table must exist after upgrade to head.""" + inspector = inspect(totp_db) + assert "auth_recovery_code" in inspector.get_table_names() + + +def test_auth_recovery_code_has_expected_columns(totp_db): + """auth_recovery_code must have id, user_id, code_hash, used_at with correct nullability.""" + inspector = inspect(totp_db) + columns = {col["name"]: col for col in inspector.get_columns("auth_recovery_code")} + + expected_non_nullable = {"id", "user_id", "code_hash"} + expected_nullable = {"used_at"} + + for col_name in expected_non_nullable: + assert col_name in columns, f"Missing column: {col_name}" + assert not columns[col_name]["nullable"], f"Column {col_name} should be NOT NULL" + + for col_name in expected_nullable: + assert col_name in columns, f"Missing column: {col_name}" + assert columns[col_name]["nullable"], f"Column {col_name} should be nullable" + + +def test_auth_recovery_code_has_fk_to_auth_users(totp_db): + """auth_recovery_code.user_id must have a FK constraint pointing to auth_users.id.""" + inspector = inspect(totp_db) + fks = inspector.get_foreign_keys("auth_recovery_code") + assert len(fks) == 1, f"Expected 1 FK on auth_recovery_code, got {len(fks)}" + fk = fks[0] + assert fk["referred_table"] == "auth_users" + assert "user_id" in fk["constrained_columns"] + assert "id" in fk["referred_columns"] + + +def test_auth_recovery_code_user_id_index_exists(totp_db): + """ix_auth_recovery_code_user_id index must exist.""" + inspector = inspect(totp_db) + indexes = inspector.get_indexes("auth_recovery_code") + index_names = {idx["name"] for idx in indexes} + assert "ix_auth_recovery_code_user_id" in index_names + + +# --------------------------------------------------------------------------- +# Default value tests +# --------------------------------------------------------------------------- + + +def test_totp_enabled_defaults_to_false(totp_db): + """A newly inserted AuthUser must have totp_enabled=False and totp_secret=None.""" + now = datetime.now(tz=timezone.utc) + with Session(totp_db) as session: + user = AuthUser( + username="testuser", + password_hash="$argon2id$v=19$m=65536,t=3,p=4$fakehash", + is_active=True, + force_password_change=False, + created_at=now, + ) + session.add(user) + session.commit() + user_id = user.id + + with Session(totp_db) as session: + fetched = session.get(AuthUser, user_id) + assert fetched is not None + assert fetched.totp_enabled is False, "totp_enabled should default to False" + assert fetched.totp_secret is None, "totp_secret should default to None" + + +def test_totp_secret_can_be_set(totp_db): + """totp_secret can be set to a base32 string and retrieved correctly.""" + now = datetime.now(tz=timezone.utc) + secret = "JBSWY3DPEHPK3PXP" + with Session(totp_db) as session: + user = AuthUser( + username="totp_user", + password_hash="$argon2id$v=19$m=65536,t=3,p=4$fakehash", + is_active=True, + force_password_change=False, + created_at=now, + totp_secret=secret, + totp_enabled=True, + ) + session.add(user) + session.commit() + user_id = user.id + + with Session(totp_db) as session: + fetched = session.get(AuthUser, user_id) + assert fetched is not None + assert fetched.totp_secret == secret + assert fetched.totp_enabled is True + + +def test_recovery_code_insert_and_used_at_nullable(totp_db): + """A RecoveryCode can be inserted with used_at=None and later marked consumed.""" + now = datetime.now(tz=timezone.utc) + with Session(totp_db) as session: + user = AuthUser( + username="rc_user", + password_hash="$argon2id$v=19$m=65536,t=3,p=4$fakehash", + is_active=True, + force_password_change=False, + created_at=now, + ) + session.add(user) + session.flush() + + code = RecoveryCode( + user_id=user.id, + code_hash="$argon2id$v=19$m=65536,t=3,p=4$rcfakehash", + used_at=None, + ) + session.add(code) + session.commit() + code_id = code.id + + with Session(totp_db) as session: + fetched = session.get(RecoveryCode, code_id) + assert fetched is not None + assert fetched.used_at is None + + # Mark as consumed + fetched.used_at = now + session.commit() + + with Session(totp_db) as session: + fetched = session.get(RecoveryCode, code_id) + assert fetched is not None + assert fetched.used_at is not None + + +# --------------------------------------------------------------------------- +# Existing user migration test (not-break-login guarantee) +# --------------------------------------------------------------------------- + + +def test_existing_user_defaults_to_totp_disabled_after_upgrade(tmp_path: Path): + """An existing user row inserted before the TOTP migration must have totp_enabled=0 after upgrade. + + Simulates the real-world scenario: a pre-existing row in auth_users gets migrated + to the TOTP revision and the server_default ensures totp_enabled=false (0) for it. + """ + db_path = tmp_path / "existing_user_test.db" + db_url = f"sqlite:///{db_path}" + alembic_cfg = _make_app_alembic_config(db_url) + + # Upgrade to the TOTP revision's predecessor (the throttle revision). + command.upgrade(alembic_cfg, "20260621_07_auth_login_throttle") + + # Insert a user row directly (simulates a user that existed before the TOTP migration). + engine = create_engine(db_url, connect_args={"check_same_thread": False}) + with engine.connect() as conn: + conn.execute( + text( + "INSERT INTO auth_users (username, password_hash, is_active, force_password_change, created_at)" + " VALUES ('legacy_admin', 'hash', 1, 0, '2026-01-01T00:00:00+00:00')" + ) + ) + conn.commit() + row = conn.execute( + text("SELECT id FROM auth_users WHERE username = 'legacy_admin'") + ).fetchone() + legacy_user_id = row[0] + engine.dispose() + + # Now upgrade to head (applies the TOTP migration). + command.upgrade(alembic_cfg, "head") + + engine = create_engine(db_url, connect_args={"check_same_thread": False}) + with engine.connect() as conn: + row = conn.execute( + text( + "SELECT totp_enabled, totp_secret FROM auth_users WHERE id = :uid" + ), + {"uid": legacy_user_id}, + ).fetchone() + engine.dispose() + + assert row is not None + # totp_enabled should be 0 (false) for the legacy row — server_default="0" ensures this. + assert row[0] == 0, f"Legacy user should have totp_enabled=0 after migration, got {row[0]}" + # totp_secret should be NULL for the legacy row. + assert row[1] is None, f"Legacy user should have totp_secret=NULL after migration, got {row[1]}" + + +# --------------------------------------------------------------------------- +# Downgrade / upgrade reversibility tests +# --------------------------------------------------------------------------- + + +def test_downgrade_removes_totp_columns_and_table(tmp_path: Path): + """downgrade -1 from head must drop auth_recovery_code and remove TOTP columns from auth_users.""" + db_path = tmp_path / "down_totp_test.db" + db_url = f"sqlite:///{db_path}" + alembic_cfg = _make_app_alembic_config(db_url) + + command.upgrade(alembic_cfg, "head") + + engine = create_engine(db_url, connect_args={"check_same_thread": False}) + inspector = inspect(engine) + assert "auth_recovery_code" in inspector.get_table_names() + auth_user_cols = {col["name"] for col in inspector.get_columns("auth_users")} + assert "totp_secret" in auth_user_cols + assert "totp_enabled" in auth_user_cols + engine.dispose() + + # Downgrade one step (removes TOTP migration). + command.downgrade(alembic_cfg, "-1") + + engine = create_engine(db_url, connect_args={"check_same_thread": False}) + inspector = inspect(engine) + assert "auth_recovery_code" not in inspector.get_table_names(), ( + "auth_recovery_code table should be gone after downgrade" + ) + auth_user_cols_after = {col["name"] for col in inspector.get_columns("auth_users")} + assert "totp_secret" not in auth_user_cols_after, ( + "totp_secret column should be removed after downgrade" + ) + assert "totp_enabled" not in auth_user_cols_after, ( + "totp_enabled column should be removed after downgrade" + ) + engine.dispose()