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
"""

import json
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,
    AssetType,
    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      # 未绑定 → 保持空


def _seed_orphan_ai_asset(
    engine, slug: str = "openai", name: str = "OpenAI", api_key: str = "sk-test-123"
) -> int:
    """造迁移前现场：一个没挂账号的 AI 资产 + 其 ai_accounts 详情"""
    with Session(engine) as s:
        provider = s.exec(select(Provider).where(Provider.slug == slug)).first()
        if not provider:
            provider = Provider(slug=slug, name=name, category="ai_agent", sdk_type=f"{slug}-api")
            s.add(provider)
            s.commit()
            s.refresh(provider)
        asset = Asset(
            name=f"{name} 账号",
            asset_type=AssetType.AI_AGENT,
            provider=slug,
            provider_id=provider.id,
            external_id=f"ai-{slug}",
        )
        s.add(asset)
        s.commit()
        s.refresh(asset)
        s.add(
            AIAccount(
                asset_id=asset.id,
                provider=slug,
                api_key_encrypted=crypto.encrypt(api_key),
            )
        )
        s.commit()
        assert asset.account_id is None  # 前置：确实还没挂靠
        return asset.id


def test_backfill_attaches_ai_assets_using_enum_member_name(legacy_engine):
    """AI 资产挂靠迁移必须用 Enum 成员名 'AI_AGENT'

    SQLAlchemy Enum 列存的是成员名而不是枚举值（'ai_agent'），
    迁移 SQL 写值会静默零命中——AI 资产永远挂不上账号，且没有任何报错。
    """
    engine = legacy_engine
    asset_id = _seed_orphan_ai_asset(engine)

    # 先把存储形态钉住：真库里 asset_type 存的就是成员名
    with engine.connect() as conn:
        assert conn.execute(sa.text("SELECT asset_type FROM assets")).scalar() == "AI_AGENT"

    db._migrate_assets_db()

    with Session(engine) as s:
        asset = s.get(Asset, asset_id)
        accounts = s.exec(select(Account)).all()
        assert len(accounts) == 1
        acc = accounts[0]
        # 按 ai_accounts.provider 找/建账号（沿用「{slug}-默认」命名），回填外键与历史字符串字段
        assert (acc.name, acc.platform) == ("openai-默认", "openai")
        assert asset.account_id == acc.id
        assert asset.account == "openai"
        # 账号尚无 API 配置时，把资产的 api_key 重组 JSON 后加密写入
        assert json.loads(crypto.decrypt(acc.api_config_encrypted)) == {"api_key": "sk-test-123"}


def test_backfill_reuses_existing_platform_account(legacy_engine):
    """同平台已有账号（步骤 1 由平台凭证下沉建出）时必须复用，不建重复账号

    否则会同时存在「deepseek」与「deepseek-默认」两个账号，
    新建那个没凭证，在 UI 上点同步/测试只会得到「缺少凭证配置」。
    """
    engine = legacy_engine
    with Session(engine) as s:
        s.add(
            Provider(
                slug="deepseek",
                name="DeepSeek",
                category="ai_agent",
                sdk_type="deepseek-api",
                api_config_encrypted=crypto.encrypt(json.dumps({"api_key": "sk-platform"})),
            )
        )
        s.commit()
    asset_id = _seed_orphan_ai_asset(engine, "deepseek", "DeepSeek", "sk-asset")

    db._migrate_assets_db()

    with Session(engine) as s:
        accounts = s.exec(select(Account)).all()
        assert len(accounts) == 1  # 复用而不是另建一个
        assert accounts[0].name == "deepseek-默认"
        assert s.get(Asset, asset_id).account_id == accounts[0].id
        # 账号已有平台级配置时不得被资产级 key 覆写
        assert json.loads(crypto.decrypt(accounts[0].api_config_encrypted)) == {
            "api_key": "sk-platform"
        }


def test_backfill_ai_assets_is_idempotent(legacy_engine):
    """重跑不得重复建账号，也不得覆写已有的账号级 API 配置"""
    engine = legacy_engine
    _seed_orphan_ai_asset(engine)

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

    with Session(engine) as s:
        assert len(s.exec(select(Account)).all()) == 1
        assert len(s.exec(select(Asset)).all()) == 1
        assert json.loads(crypto.decrypt(s.exec(select(Account)).first().api_config_encrypted)) == {
            "api_key": "sk-test-123"
        }