WonderQ-Project/WonderQ-Admin/tests/test_admin_rbac.py

from datetime import datetime, timedelta, timezone
import importlib.util
from pathlib import Path

from fastapi.testclient import TestClient

from app.auth import decode_admin_access_token, hash_password
from app.database import get_db
from app.main import create_app
from app.models import AdminDepartment, AdminMenu, AdminRole, AdminRoleDepartment, AdminRoleMenu, AdminUser, AdminUserDepartment, AdminUserRole, Lead, MediaAsset
from app.rbac import DATA_SCOPE_VALUES, build_menu_tree
from app.redis_session import (
    InMemoryAdminSessionStore,
    get_admin_session_store,
    hash_refresh_token,
)


class LoginDb:
    def __init__(self, user: AdminUser):
        self.user = user

    def scalar(self, _statement):
        return self.user

    def get(self, model, identifier):
        return self.user if model is AdminUser and identifier == self.user.id else None

    class _EmptyResult:
        def all(self):
            return []

    def scalars(self, _statement):
        return self._EmptyResult()


def test_refresh_session_rotation_rejects_reuse():
    store = InMemoryAdminSessionStore()
    now = datetime.now(timezone.utc)
    store.create(
        session_id="session-1",
        user_id="admin-1",
        access_jti="access-1",
        refresh_hash=hash_refresh_token("refresh-1"),
        access_expires_at=now + timedelta(minutes=15),
        refresh_expires_at=now + timedelta(days=7),
    )

    assert store.is_access_active("session-1", "access-1")
    assert store.rotate(
        session_id="session-1",
        refresh_hash=hash_refresh_token("refresh-1"),
        access_jti="access-2",
        refresh_hash_next=hash_refresh_token("refresh-2"),
        access_expires_at=now + timedelta(minutes=15),
        refresh_expires_at=now + timedelta(days=7),
    )
    assert not store.is_access_active("session-1", "access-1")
    assert store.is_access_active("session-1", "access-2")
    assert not store.rotate(
        session_id="session-1",
        refresh_hash=hash_refresh_token("refresh-1"),
        access_jti="access-3",
        refresh_hash_next=hash_refresh_token("refresh-3"),
        access_expires_at=now + timedelta(minutes=15),
        refresh_expires_at=now + timedelta(days=7),
    )


def test_in_memory_store_supports_login_limit_and_permission_cache():
    store = InMemoryAdminSessionStore()
    assert store.allow_login_attempt("ip:admin@example.com", 2, 60)
    assert store.allow_login_attempt("ip:admin@example.com", 2, 60)
    assert not store.allow_login_attempt("ip:admin@example.com", 2, 60)
    store.set_permission_context("admin-1", {"permissions": ["admin:read"]}, 60)
    assert store.get_permission_context("admin-1") == {"permissions": ["admin:read"]}


def test_in_memory_captcha_is_single_use():
    store = InMemoryAdminSessionStore()
    store.create_captcha("captcha-1", "ABCD", 120)

    assert store.consume_captcha("captcha-1", "ABCD") is True
    assert store.consume_captcha("captcha-1", "ABCD") is False


def test_admin_captcha_endpoint_returns_graphical_challenge():
    store = InMemoryAdminSessionStore()
    app = create_app()
    app.dependency_overrides[get_admin_session_store] = lambda: store

    try:
        response = TestClient(app).get("/api/admin/auth/captcha")
    finally:
        app.dependency_overrides.clear()

    assert response.status_code == 200
    data = response.json()["data"]
    assert data["captchaEnabled"] is True
    assert data["captchaId"]
    assert data["image"].startswith("data:image/svg+xml;base64,")
    assert data["expiresIn"] == 120


def test_admin_access_token_contains_scoped_session_claims():
    user = AdminUser(
        id="admin-1",
        email="admin@example.com",
        name="Admin",
        role="admin",
        passwordHash=hash_password("ChangeMe123!", rounds=4),
        isActive=True,
    )

    from app.auth import create_access_token

    token = create_access_token(user, session_id="session-1", access_jti="access-1")
    payload = decode_admin_access_token(token)

    assert payload["typ"] == "admin_access"
    assert payload["aud"] == "admin"
    assert payload["sid"] == "session-1"
    assert payload["jti"] == "access-1"
    assert payload["role"] == "admin"


def test_admin_login_keeps_legacy_fields_and_sets_http_only_refresh_cookie():
    user = AdminUser(
        id="admin-1",
        email="admin@example.com",
        name="Admin",
        role="admin",
        passwordHash=hash_password("ChangeMe123!", rounds=4),
        isActive=True,
    )
    app = create_app()
    app.dependency_overrides[get_db] = lambda: LoginDb(user)
    store = InMemoryAdminSessionStore()
    store.create_captcha("captcha-login", "ABCD", 120)
    app.dependency_overrides[get_admin_session_store] = lambda: store

    try:
        response = TestClient(app).post(
            "/api/admin/auth/login",
            json={
                "email": "admin@example.com",
                "password": "ChangeMe123!",
                "captchaId": "captcha-login",
                "captchaCode": "ABCD",
                "rememberMe": False,
            },
        )
    finally:
        app.dependency_overrides.clear()

    assert response.status_code == 200
    body = response.json()["data"]
    assert body["token"] == body["accessToken"]
    assert body["expiresIn"] > 0
    assert body["user"]["id"] == "admin-1"
    cookie = response.headers["set-cookie"]
    assert "HttpOnly" in cookie
    assert "Path=/api/admin/auth" in cookie
    assert "Max-Age=" not in cookie


def test_refresh_rotates_access_token_and_logout_revokes_it():
    user = AdminUser(
        id="admin-1",
        email="admin@example.com",
        name="Admin",
        role="admin",
        passwordHash=hash_password("ChangeMe123!", rounds=4),
        isActive=True,
    )
    store = InMemoryAdminSessionStore()
    app = create_app()
    app.dependency_overrides[get_db] = lambda: LoginDb(user)
    app.dependency_overrides[get_admin_session_store] = lambda: store

    try:
        client = TestClient(app)
        store.create_captcha("captcha-refresh", "ABCD", 120)
        login_response = client.post(
            "/api/admin/auth/login",
            json={
                "email": "admin@example.com",
                "password": "ChangeMe123!",
                "captchaId": "captcha-refresh",
                "captchaCode": "ABCD",
                "rememberMe": False,
            },
        )
        old_token = login_response.json()["data"]["accessToken"]

        refresh_response = client.post("/api/admin/auth/refresh")
        new_token = refresh_response.json()["data"]["accessToken"]

        assert refresh_response.status_code == 200
        assert "Max-Age=" not in refresh_response.headers["set-cookie"]
        assert new_token != old_token
        assert client.get("/api/admin/me", headers={"Authorization": f"Bearer {old_token}"}).status_code == 401
        assert client.get("/api/admin/me", headers={"Authorization": f"Bearer {new_token}"}).status_code == 200

        logout_response = client.post("/api/admin/auth/logout")
        assert logout_response.status_code == 200
        assert client.get("/api/admin/me", headers={"Authorization": f"Bearer {new_token}"}).status_code == 401
    finally:
        app.dependency_overrides.clear()


def test_remember_me_keeps_persistent_cookie_after_refresh_rotation():
    user = AdminUser(
        id="admin-1",
        email="admin@example.com",
        name="Admin",
        role="admin",
        passwordHash=hash_password("ChangeMe123!", rounds=4),
        isActive=True,
    )
    store = InMemoryAdminSessionStore()
    app = create_app()
    app.dependency_overrides[get_db] = lambda: LoginDb(user)
    app.dependency_overrides[get_admin_session_store] = lambda: store

    try:
        client = TestClient(app)
        store.create_captcha("captcha-remember", "ABCD", 120)
        login_response = client.post(
            "/api/admin/auth/login",
            json={
                "email": "admin@example.com",
                "password": "ChangeMe123!",
                "captchaId": "captcha-remember",
                "captchaCode": "ABCD",
                "rememberMe": True,
            },
        )

        assert login_response.status_code == 200
        assert "Max-Age=" in login_response.headers["set-cookie"]

        refresh_response = client.post("/api/admin/auth/refresh")

        assert refresh_response.status_code == 200
        assert "Max-Age=" in refresh_response.headers["set-cookie"]
    finally:
        app.dependency_overrides.clear()


def test_rbac_models_define_normalized_association_tables():
    assert AdminRole.__tablename__ == "AdminRole"
    assert AdminMenu.__tablename__ == "AdminMenu"
    assert AdminDepartment.__tablename__ == "AdminDepartment"
    assert {AdminUserRole.__table__.c.userId.name, AdminUserRole.__table__.c.roleId.name} == {"userId", "roleId"}
    assert {AdminRoleMenu.__table__.c.roleId.name, AdminRoleMenu.__table__.c.menuId.name} == {"roleId", "menuId"}
    assert {AdminRoleDepartment.__table__.c.roleId.name, AdminRoleDepartment.__table__.c.deptId.name} == {"roleId", "deptId"}
    assert {AdminUserDepartment.__table__.c.userId.name, AdminUserDepartment.__table__.c.deptId.name} == {"userId", "deptId"}


def test_rbac_migration_is_appended_to_current_head():
    migration_path = Path(__file__).parents[1] / "alembic" / "versions" / "0025_admin_rbac.py"
    spec = importlib.util.spec_from_file_location("admin_rbac_migration", migration_path)
    assert spec and spec.loader
    migration = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(migration)

    assert migration.down_revision == "0024_vehicle_demand"
    assert set(migration.RBAC_TABLES) == {
        "AdminRole",
        "AdminMenu",
        "AdminDepartment",
        "AdminUserRole",
        "AdminRoleMenu",
        "AdminUserDepartment",
        "AdminRoleDepartment",
    }


def test_ownership_migration_is_appended_after_rbac_and_models_expose_scope_fields():
    migration_path = Path(__file__).parents[1] / "alembic" / "versions" / "0026_admin_ownership.py"
    spec = importlib.util.spec_from_file_location("admin_ownership_migration", migration_path)
    assert spec and spec.loader
    migration = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(migration)

    assert migration.down_revision == "0025_admin_rbac"
    assert {"deptId", "createdById"}.issubset(Lead.__table__.c.keys())
    assert {"deptId", "createdById"}.issubset(MediaAsset.__table__.c.keys())
    assert "Lead" in migration.OWNED_TABLES


def test_dynamic_menu_tree_preserves_parent_order_and_button_permissions():
    menus = [
        AdminMenu(id="child", parentId="root", name="按钮", type="button", permissionCode="home:edit", sortOrder=20, isVisible=True, isActive=True),
        AdminMenu(id="root", parentId=None, name="首页", type="page", path="/home", componentKey="HomePage", permissionCode="home:read", sortOrder=10, isVisible=True, isActive=True),
        AdminMenu(id="hidden", parentId=None, name="隐藏", type="page", sortOrder=30, isVisible=False, isActive=True),
    ]

    tree = build_menu_tree(menus)

    assert DATA_SCOPE_VALUES == {"all", "dept", "dept_and_children", "custom_dept", "self"}
    assert [item["id"] for item in tree] == ["root"]
    assert tree[0]["children"][0]["permission"] == "home:edit"
    assert tree[0]["isVisible"] is True
    assert tree[0]["isActive"] is True


def test_dynamic_menu_tree_keeps_granted_menu_when_parent_is_not_in_scope():
    menus = [
        AdminMenu(id="page", parentId="missing-parent", name="受限页面", type="page", path="/restricted", sortOrder=10, isVisible=True, isActive=True),
    ]

    tree = build_menu_tree(menus)

    assert [item["id"] for item in tree] == ["page"]