vps-manager/tests/test_credential_migration.py

"""凭据中心化迁移测试（accounts.login_password → credentials）

模拟旧库（accounts 无 credential_id 列）→ 跑迁移 → 断言：
密文原样搬入、site 取平台显示名、账号关联回填、原字段置空、重跑幂等、
无密码账号不产生空条目，以及旧 credentials 表补 otp_verified_at 列并回填。

运行：.venv/bin/pytest tests/test_credential_migration.py -v
"""

from datetime import datetime, timezone

import pytest
import sqlalchemy as sa
from cryptography.fernet import Fernet
from sqlmodel import Session, SQLModel, create_engine, select

import app.database as db
from app.core import crypto
from app.core.config import settings
from app.models.asset import AIAccount, Account, Asset, CloudflareDetail, DomainDetail, VPSDetail
from app.models.credential import Credential, LoginType
from app.models.provider import Provider
from app.models.ssl import SiteCert, Subdomain

# 旧结构 accounts 表：无 credential_id 列（迁移应自动补列）
OLD_ACCOUNTS_DDL = """
CREATE TABLE accounts (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    name VARCHAR NOT NULL,
    platform VARCHAR,
    remark VARCHAR,
    login_user VARCHAR,
    login_password_encrypted VARCHAR,
    api_config_encrypted VARCHAR,
    last_synced_at DATETIME,
    created_at DATETIME NOT NULL
)
"""


@pytest.fixture()
def legacy_engine(tmp_path, monkeypatch):
    """旧结构测试库：新表齐备，但 accounts 缺 credential_id 列"""
    monkeypatch.setattr(settings, "MASTER_KEY", Fernet.generate_key().decode())
    crypto._get_fernet.cache_clear()  # 避免复用其它测试缓存的 Fernet 实例
    engine = create_engine(
        f"sqlite:///{tmp_path / 'test.db'}", connect_args={"check_same_thread": False}
    )
    # 除 Account 外按当前模型建表（含新 credentials 表；accounts 用旧 DDL）
    SQLModel.metadata.create_all(
        engine,
        tables=[
            Provider.__table__, Asset.__table__, VPSDetail.__table__, DomainDetail.__table__,
            AIAccount.__table__, CloudflareDetail.__table__, Subdomain.__table__,
            SiteCert.__table__, Credential.__table__,
        ],
    )
    with engine.begin() as conn:
        conn.execute(sa.text(OLD_ACCOUNTS_DDL))
    # 迁移函数引用模块级 assets_engine → 指向测试库
    monkeypatch.setattr(db, "assets_engine", engine)
    yield engine
    crypto._get_fernet.cache_clear()


def _seed(engine, pwd_aliyun: str, pwd_forum: str) -> None:
    """造迁移前现场：平台账号（有密码）/ 无平台账号（有密码）/ 无密码账号"""
    with Session(engine) as s:
        s.add(Provider(slug="aliyun", name="阿里云", category="vps"))
        s.commit()
    now = datetime.now(timezone.utc).replace(tzinfo=None)
    with engine.begin() as conn:
        conn.execute(
            sa.text(
                "INSERT INTO accounts (name, platform, login_user, login_password_encrypted, created_at) "
                "VALUES ('me@gmail.com', 'aliyun', 'me@gmail.com', :pwd, :now)"
            ),
            {"pwd": pwd_aliyun, "now": now},
        )
        # platform 为空串：site 应回退 '未分类'；无 login_user：username 回退 name
        conn.execute(
            sa.text(
                "INSERT INTO accounts (name, platform, login_user, login_password_encrypted, created_at) "
                "VALUES ('forum-user', '', NULL, :pwd, :now)"
            ),
            {"pwd": pwd_forum, "now": now},
        )
        # 无密码账号：不应产生凭据条目
        conn.execute(
            sa.text(
                "INSERT INTO accounts (name, platform, created_at) VALUES ('no-pwd', 'vultr', :now)"
            ),
            {"now": now},
        )


def test_migration_moves_passwords_to_vault(legacy_engine):
    engine = legacy_engine
    pwd_aliyun = crypto.encrypt("p@ss-阿里云")
    pwd_forum = crypto.encrypt("forum-pass-中文")
    _seed(engine, pwd_aliyun, pwd_forum)

    db._migrate_assets_db()

    with Session(engine) as s:
        creds = s.exec(select(Credential).order_by(Credential.id)).all()
        # 仅两条有密码的账号被迁移，无密码账号不产生空条目
        assert len(creds) == 2

        c1, c2 = creds
        # 密文原样搬入（不重新加密，避免中间态明文暴露），仍可解密还原
        assert c1.password_encrypted == pwd_aliyun
        assert crypto.decrypt(c1.password_encrypted) == "p@ss-阿里云"
        # site 取平台显示名（join providers），username 取 login_user
        assert c1.site == "阿里云"
        assert c1.username == "me@gmail.com"
        assert c1.login_type == LoginType.PASSWORD
        # 无 provider 匹配 + platform 空 → site 回退 '未分类'；username 回退 name
        assert c2.password_encrypted == pwd_forum
        assert c2.site == "未分类"
        assert c2.username == "forum-user"

        # 账号侧：credential_id 回填、原密码字段置空（唯一事实源，防双份漂移）
        accs = {a.name: a for a in s.exec(select(Account)).all()}
        assert accs["me@gmail.com"].credential_id == c1.id
        assert accs["me@gmail.com"].login_password_encrypted is None
        assert accs["forum-user"].credential_id == c2.id
        assert accs["forum-user"].login_password_encrypted is None
        assert accs["no-pwd"].credential_id is None


def test_migration_is_idempotent(legacy_engine):
    """重跑迁移不得重复建条目（幂等条件：credential_id 已非空则跳过）"""
    engine = legacy_engine
    _seed(engine, crypto.encrypt("pw1"), crypto.encrypt("pw2"))

    db._migrate_assets_db()
    db._migrate_assets_db()
    db._migrate_assets_db()

    with Session(engine) as s:
        assert len(s.exec(select(Credential)).all()) == 2


def test_migration_adds_credential_id_column(legacy_engine):
    """旧库缺 credential_id 列时迁移应自动补列"""
    engine = legacy_engine
    with engine.connect() as conn:
        cols = {c["name"] for c in sa.inspect(conn).get_columns("accounts")}
    assert "credential_id" not in cols  # 前置：确实是旧结构

    db._migrate_assets_db()

    with engine.connect() as conn:
        cols = {c["name"] for c in sa.inspect(conn).get_columns("accounts")}
        indexes = {ix["name"] for ix in sa.inspect(conn).get_indexes("accounts")}
    assert "credential_id" in cols

    db._migrate_indexes()
    with engine.connect() as conn:
        indexes = {ix["name"] for ix in sa.inspect(conn).get_indexes("accounts")}
    assert "ix_accounts_credential_id" in indexes


# 旧结构 credentials 表：无 otp_verified_at 列（当时绑定强制校验过动态码）
OLD_CREDENTIALS_DDL = """
CREATE TABLE credentials (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    site VARCHAR NOT NULL,
    username VARCHAR,
    login_type VARCHAR NOT NULL,
    oauth_provider VARCHAR,
    password_encrypted VARCHAR,
    otp_secret_encrypted VARCHAR,
    url VARCHAR,
    note VARCHAR,
    created_at DATETIME NOT NULL,
    updated_at DATETIME NOT NULL
)
"""


def test_migration_backfills_otp_verified(legacy_engine):
    """旧库 credentials 缺列时自动补列，并把已绑定条目回填为「已验证」

    这些条目在旧版是校验过当前动态码才落库的，若回填为 NULL 会被前端全线误报。
    """
    engine = legacy_engine
    now = datetime.now(timezone.utc).replace(tzinfo=None)
    with engine.begin() as conn:
        conn.execute(sa.text("DROP TABLE credentials"))
        conn.execute(sa.text(OLD_CREDENTIALS_DDL))
        conn.execute(
            sa.text(
                "INSERT INTO credentials (site, login_type, otp_secret_encrypted, created_at, updated_at) "
                "VALUES ('GitHub', 'PASSWORD', 'legacy-cipher-text', :now, :now)"
            ),
            {"now": now},
        )
        conn.execute(
            sa.text(
                "INSERT INTO credentials (site, login_type, otp_secret_encrypted, created_at, updated_at) "
                "VALUES ('Forum', 'PASSWORD', NULL, :now, :now)"
            ),
            {"now": now},
        )

    db._migrate_assets_db()

    with engine.connect() as conn:
        cols = {c["name"] for c in sa.inspect(conn).get_columns("credentials")}
        assert "otp_verified_at" in cols
        rows = conn.execute(
            sa.text("SELECT site, otp_verified_at FROM credentials ORDER BY id")
        ).fetchall()
    assert rows[0][1] is not None  # 已绑 2FA → 视为已验证
    assert rows[1][1] is None      # 未绑定 → 保持空