fischer-agentkit/tests/unit/test_saml_provider.py

"""SAML 2.0 provider 单元测试 (U4 SSO 集成)。

验证：
- ACS NameID 提取 + JIT 用户创建
- 按 NameID 关联（R1a: 禁止邮箱合并）
- 相同 NameID 复用现有用户
"""

from __future__ import annotations

import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

import aiosqlite
import pytest

from agentkit.server.auth.idp_registry import (
    IDPConfig,
    IDPRegistry,
    SAMLConfig,
)
from agentkit.server.auth.models import init_auth_db
from agentkit.server.auth.providers.saml import SAMLProvider


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------


@pytest.fixture
def saml_idp(monkeypatch: pytest.MonkeyPatch) -> IDPRegistry:
    """注册一个测试 SAML IDP 并替换全局 registry。"""
    registry = IDPRegistry()
    registry.register(
        IDPConfig(
            name="test_saml",
            type="saml",
            display_name="Test SAML",
            saml=SAMLConfig(
                entity_id="https://idp.test/entity",
                sso_url="https://idp.test/sso",
                idp_cert="-----BEGIN CERTIFICATE-----\nfakecert\n-----END CERTIFICATE-----",
                sp_entity_id="https://app.test/sp",
                acs_url="https://app.test/api/v1/auth/saml/test_saml/acs",
            ),
        )
    )
    monkeypatch.setattr("agentkit.server.auth.providers.saml.get_idp_registry", lambda: registry)
    return registry


@pytest.fixture
async def auth_db(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path:
    db_path = tmp_path / "auth.db"
    await init_auth_db(db_path)
    monkeypatch.setenv("AGENTKIT_AUTH_DB", str(db_path))
    return db_path


@pytest.fixture
def mock_saml_parse(monkeypatch: pytest.MonkeyPatch) -> None:
    """替换 _process_saml 为假实现 — 返回固定 NameID + attributes。"""

    def _fake_process(
        self: SAMLProvider, req: dict, cfg: Any, provider_name: str
    ) -> tuple[str, dict[str, Any]]:
        return "saml-nameid-456", {"email": "saml_user@example.com", "name": "SAML User"}

    monkeypatch.setattr(SAMLProvider, "_process_saml", _fake_process)


# ---------------------------------------------------------------------------
# Tests
# ---------------------------------------------------------------------------


class TestHandleAcs:
    async def test_creates_user_by_nameid(
        self,
        saml_idp: IDPRegistry,
        auth_db: Path,
        mock_saml_parse: None,
    ) -> None:
        """ACS 消费 SAMLResponse + NameID 提取 + JIT 用户创建。"""
        user = await SAMLProvider().handle_acs("test_saml", "fake-saml-response")
        assert user["username"] == "SAML User"
        assert user["email"] == "saml_user@example.com"

        # 验证 user_idp_links 按 NameID 关联
        async with aiosqlite.connect(str(auth_db)) as db:
            db.row_factory = aiosqlite.Row
            cursor = await db.execute(
                "SELECT * FROM user_idp_links WHERE idp_name = ? AND idp_subject = ?",
                ("test_saml", "saml-nameid-456"),
            )
            link = await cursor.fetchone()
        assert link is not None
        assert link["idp_type"] == "saml"
        assert link["user_id"] == user["id"]

    async def test_same_nameid_reuses_user(
        self,
        saml_idp: IDPRegistry,
        auth_db: Path,
        mock_saml_parse: None,
    ) -> None:
        """相同 NameID 第二次登录复用现有用户。"""
        provider = SAMLProvider()
        user1 = await provider.handle_acs("test_saml", "resp1")
        user2 = await provider.handle_acs("test_saml", "resp2")
        assert user1["id"] == user2["id"]

    async def test_no_email_merge(
        self,
        saml_idp: IDPRegistry,
        auth_db: Path,
        mock_saml_parse: None,
    ) -> None:
        """R1a: 不同 NameID 相同 email — 禁止合并。"""
        now = datetime.now(timezone.utc).isoformat()
        local_id = str(uuid.uuid4())
        async with aiosqlite.connect(str(auth_db)) as db:
            await db.execute(
                "INSERT INTO users (id, username, email, password_hash, role, is_active, "
                "is_terminal_authorized, is_server_terminal_authorized, created_at, updated_at) "
                "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
                (
                    local_id,
                    "local_saml_user",
                    "saml_user@example.com",
                    "$2b$12$dummyhash",
                    "member",
                    1,
                    0,
                    0,
                    now,
                    now,
                ),
            )
            await db.commit()

        # SAML NameID=saml-nameid-456（不同于本地用户）
        saml_user = await SAMLProvider().handle_acs("test_saml", "resp")
        assert saml_user["id"] != local_id

    async def test_unknown_provider_raises(
        self,
        saml_idp: IDPRegistry,
        auth_db: Path,
        mock_saml_parse: None,
    ) -> None:
        with pytest.raises(ValueError, match="未配置"):
            await SAMLProvider().handle_acs("nonexistent", "resp")

    async def test_missing_nameid_raises(
        self,
        saml_idp: IDPRegistry,
        auth_db: Path,
        monkeypatch: pytest.MonkeyPatch,
    ) -> None:
        """SAML assertion 缺少 NameID 应拒绝。"""
        monkeypatch.setattr(
            SAMLProvider,
            "_process_saml",
            lambda self, req, cfg, name: ("", {}),
        )
        with pytest.raises(ValueError, match="NameID"):
            await SAMLProvider().handle_acs("test_saml", "resp")


class TestAuthenticateNotImplemented:
    async def test_raises(self) -> None:
        from agentkit.server.auth.providers.exceptions import ProviderNotImplemented

        with pytest.raises(ProviderNotImplemented):
            await SAMLProvider().authenticate(username="x", password="y")


class TestRevokeUser:
    async def test_disables_local_user(
        self,
        saml_idp: IDPRegistry,
        auth_db: Path,
        mock_saml_parse: None,
    ) -> None:
        provider = SAMLProvider()
        user = await provider.handle_acs("test_saml", "resp")
        await provider.revoke_user(user["id"])
        async with aiosqlite.connect(str(auth_db)) as db:
            db.row_factory = aiosqlite.Row
            cursor = await db.execute("SELECT is_active FROM users WHERE id = ?", (user["id"],))
            row = await cursor.fetchone()
        assert bool(row["is_active"]) is False