CIAgent/apps/api/tests/unit/test_db_viewer.py

"""Settings -> Database viewer: mint (admin-gated) -> bootstrap (token ->
session cookie) -> verify (cookie -> live is_admin check) round trip.

Cookies are extracted from Set-Cookie headers and passed explicitly on
follow-up requests rather than relying on TestClient's cookie jar - the
session cookie is marked Secure, and TestClient's base_url is plain http,
so a real cookie jar wouldn't resend it anyway."""

from __future__ import annotations

import uuid
from datetime import UTC, datetime, timedelta

import jwt
from fastapi.testclient import TestClient
from sqlalchemy.ext.asyncio import AsyncSession

from app.core.config import get_settings
from app.repositories.user_repository import UserRepository


def _unique_email() -> str:
    return f"user-{uuid.uuid4().hex[:12]}@example.com"


async def _register_admin_and_login(client: TestClient, db_session: AsyncSession) -> dict[str, str]:
    email = f"admin-{uuid.uuid4().hex[:12]}@example.com"
    client.post(
        "/api/v1/auth/register",
        json={"email": email, "password": "correct-horse-1", "display_name": "Admin"},
    )
    user = await UserRepository(db_session).get_by_email(email)
    user.is_admin = True
    await db_session.commit()
    tokens = client.post(
        "/api/v1/auth/login", json={"email": email, "password": "correct-horse-1"}
    ).json()
    return {"Authorization": f"Bearer {tokens['access_token']}"}


async def _register_non_admin_and_login(client: TestClient) -> dict[str, str]:
    email = _unique_email()
    client.post(
        "/api/v1/auth/register",
        json={"email": email, "password": "correct-horse-1", "display_name": "T"},
    )
    tokens = client.post(
        "/api/v1/auth/login", json={"email": email, "password": "correct-horse-1"}
    ).json()
    return {"Authorization": f"Bearer {tokens['access_token']}"}


def _extract_cookie(set_cookie_header: str) -> tuple[str, str]:
    first_pair = set_cookie_header.split(";", 1)[0]
    name, value = first_pair.split("=", 1)
    return name, value


# --- /db-viewer/session (mint) --------------------------------------------


def test_create_session_requires_auth(client: TestClient):
    resp = client.post("/api/v1/db-viewer/session")
    assert resp.status_code == 401


async def test_create_session_non_admin_forbidden(client: TestClient):
    headers = await _register_non_admin_and_login(client)
    resp = client.post("/api/v1/db-viewer/session", headers=headers)
    assert resp.status_code == 403


async def test_admin_can_mint_a_bootstrap_token(client: TestClient, db_session: AsyncSession):
    headers = await _register_admin_and_login(client, db_session)
    resp = client.post("/api/v1/db-viewer/session", headers=headers)
    assert resp.status_code == 200
    assert resp.json()["token"]


async def test_minting_a_session_is_logged_to_the_admins_account_activity(
    client: TestClient, db_session: AsyncSession
):
    headers = await _register_admin_and_login(client, db_session)
    client.post("/api/v1/db-viewer/session", headers=headers)

    events = client.get("/api/v1/auth/security-events", headers=headers).json()
    assert any(e["event_type"] == "db_viewer_session_created" for e in events)


# --- /db-viewer/bootstrap (token -> cookie) -------------------------------


async def test_bootstrap_with_a_valid_token_sets_a_session_cookie_and_redirects(
    client: TestClient, db_session: AsyncSession
):
    headers = await _register_admin_and_login(client, db_session)
    token = client.post("/api/v1/db-viewer/session", headers=headers).json()["token"]

    resp = client.get(
        "/api/v1/db-viewer/bootstrap", params={"token": token}, follow_redirects=False
    )
    assert resp.status_code == 302
    assert resp.headers["location"] == "/?pgsql=postgres"
    assert "db_viewer_session" in resp.headers["set-cookie"]
    assert "HttpOnly" in resp.headers["set-cookie"]
    assert "Secure" in resp.headers["set-cookie"]


async def test_bootstrap_success_is_logged_to_the_admins_account_activity(
    client: TestClient, db_session: AsyncSession
):
    headers = await _register_admin_and_login(client, db_session)
    token = client.post("/api/v1/db-viewer/session", headers=headers).json()["token"]

    client.get("/api/v1/db-viewer/bootstrap", params={"token": token}, follow_redirects=False)

    events = client.get("/api/v1/auth/security-events", headers=headers).json()
    assert any(e["event_type"] == "db_viewer_accessed" for e in events)


def test_bootstrap_rejects_a_garbage_token(client: TestClient):
    resp = client.get(
        "/api/v1/db-viewer/bootstrap", params={"token": "not-a-real-token"}, follow_redirects=False
    )
    assert resp.status_code == 403


def test_bootstrap_rejects_an_expired_token(client: TestClient):
    settings = get_settings()
    now = datetime.now(UTC)
    expired_token = jwt.encode(
        {
            "sub": str(uuid.uuid4()),
            "type": "db_viewer_bootstrap",
            "iat": now - timedelta(minutes=10),
            "exp": now - timedelta(minutes=5),
            "jti": "x",
        },
        settings.jwt_secret,
        algorithm="HS256",
    )
    resp = client.get(
        "/api/v1/db-viewer/bootstrap", params={"token": expired_token}, follow_redirects=False
    )
    assert resp.status_code == 403


def test_bootstrap_rejects_a_session_type_token_used_as_a_bootstrap_token(client: TestClient):
    """Type confusion guard: a DB_VIEWER_SESSION token must not work as a
    DB_VIEWER_BOOTSTRAP token, even though both are signed with the same
    jwt_secret."""
    settings = get_settings()
    now = datetime.now(UTC)
    session_typed_token = jwt.encode(
        {
            "sub": str(uuid.uuid4()),
            "type": "db_viewer_session",
            "iat": now,
            "exp": now + timedelta(minutes=2),
            "jti": "x",
        },
        settings.jwt_secret,
        algorithm="HS256",
    )
    resp = client.get(
        "/api/v1/db-viewer/bootstrap",
        params={"token": session_typed_token},
        follow_redirects=False,
    )
    assert resp.status_code == 403


# --- /db-viewer/verify (cookie -> live is_admin check) --------------------


async def test_verify_succeeds_for_a_live_admin_with_a_valid_session_cookie(
    client: TestClient, db_session: AsyncSession
):
    headers = await _register_admin_and_login(client, db_session)
    token = client.post("/api/v1/db-viewer/session", headers=headers).json()["token"]
    bootstrap_resp = client.get(
        "/api/v1/db-viewer/bootstrap", params={"token": token}, follow_redirects=False
    )
    cookie_name, cookie_value = _extract_cookie(bootstrap_resp.headers["set-cookie"])

    verify_resp = client.get("/api/v1/db-viewer/verify", cookies={cookie_name: cookie_value})
    assert verify_resp.status_code == 200


def test_verify_rejects_a_missing_cookie(client: TestClient):
    resp = client.get("/api/v1/db-viewer/verify")
    assert resp.status_code == 401


async def test_verify_rejects_a_session_cookie_once_admin_is_revoked(
    client: TestClient, db_session: AsyncSession
):
    headers = await _register_admin_and_login(client, db_session)
    token = client.post("/api/v1/db-viewer/session", headers=headers).json()["token"]
    bootstrap_resp = client.get(
        "/api/v1/db-viewer/bootstrap", params={"token": token}, follow_redirects=False
    )
    cookie_name, cookie_value = _extract_cookie(bootstrap_resp.headers["set-cookie"])

    # Sanity: still works before revocation.
    assert (
        client.get("/api/v1/db-viewer/verify", cookies={cookie_name: cookie_value}).status_code
        == 200
    )

    admin_email_resp = client.get("/api/v1/auth/me", headers=headers)
    admin = await UserRepository(db_session).get_by_email(admin_email_resp.json()["email"])
    admin.is_admin = False
    await db_session.commit()

    assert (
        client.get("/api/v1/db-viewer/verify", cookies={cookie_name: cookie_value}).status_code
        == 401
    )