apps/backend-api/tests/test_api_security.py

from __future__ import annotations

from fastapi.testclient import TestClient

from app import main
from app.config import settings
from app.repository import InMemoryRepo


def setup_function() -> None:
    main.repo = InMemoryRepo()


def api_headers(**extra: str) -> dict[str, str]:
    headers = dict(extra)
    if settings.api_shared_secret:
        headers["x-api-key"] = settings.api_shared_secret
    return headers


def create_session(client: TestClient) -> str:
    response = client.post("/v1/auth/guest", headers=api_headers(), json={})
    assert response.status_code == 200
    return response.json()["session_id"]


def create_project(client: TestClient, session_id: str) -> str:
    response = client.post(
        "/v1/projects",
        headers=api_headers(**{"x-session-id": session_id}),
        json={"title": "Private project", "mode": "part", "output_type": "3d_solid"},
    )
    assert response.status_code == 200
    return response.json()["id"]


def override_settings(**updates: int):
    previous = {key: getattr(settings, key) for key in updates}
    for key, value in updates.items():
        object.__setattr__(settings, key, value)
    return previous


def restore_settings(previous: dict[str, int]) -> None:
    for key, value in previous.items():
        object.__setattr__(settings, key, value)


def test_project_access_is_scoped_to_session() -> None:
    with TestClient(main.app) as client:
        owner = create_session(client)
        attacker = create_session(client)
        project_id = create_project(client, owner)
        owner_response = client.get(f"/v1/projects/{project_id}", headers=api_headers(**{"x-session-id": owner}))
        assert "owner_session_id" not in owner_response.json()["project"]
        assert owner not in owner_response.text
        response = client.get(f"/v1/projects/{project_id}", headers=api_headers(**{"x-session-id": attacker}))
        assert response.status_code == 403


def test_state_change_rejects_untrusted_browser_origin() -> None:
    with TestClient(main.app) as client:
        response = client.post(
            "/v1/auth/guest", headers=api_headers(origin="https://attacker.example"), json={},
        )
        assert response.status_code == 403


def test_generation_uses_explicit_parent_and_is_idempotent() -> None:
    with TestClient(main.app) as client:
        session = create_session(client)
        project_id = create_project(client, session)
        payload = {
            "message": "A bracket 80 width and 6 thickness",
            "parent_version_id": None,
            "attachment_ids": [],
            "profile": "balanced",
            "idempotency_key": "same-request-key",
        }
        headers = api_headers(**{"x-session-id": session})
        first = client.post(f"/v1/projects/{project_id}/generations", headers=headers, json=payload)
        second = client.post(f"/v1/projects/{project_id}/generations", headers=headers, json=payload)
        assert first.status_code == 202
        assert second.status_code == 202
        assert first.json()["id"] == second.json()["id"]


def test_generation_rejects_message_beyond_profile_limit() -> None:
    with TestClient(main.app) as client:
        session = create_session(client)
        project_id = create_project(client, session)
        payload = {
            "message": "x" * 1201,
            "parent_version_id": None,
            "attachment_ids": [],
            "profile": "balanced",
            "idempotency_key": "too-long-balanced",
        }
        headers = api_headers(**{"x-session-id": session})
        response = client.post(f"/v1/projects/{project_id}/generations", headers=headers, json=payload)
        assert response.status_code == 422
        detail = response.json()["detail"]
        assert detail["profile"] == "balanced"
        assert detail["max_prompt_chars"] == 1200


def test_legacy_generate_rejects_prompt_beyond_profile_limit() -> None:
    with TestClient(main.app) as client:
        session = create_session(client)
        project_id = create_project(client, session)
        headers = api_headers(**{"x-session-id": session})
        response = client.post(
            f"/v1/projects/{project_id}/generate",
            headers=headers,
            json={"prompt": "x" * 701, "profile": "fast"},
        )
        assert response.status_code == 422
        detail = response.json()["detail"]
        assert detail["profile"] == "fast"
        assert detail["max_prompt_chars"] == 700


def test_clarification_rejects_message_beyond_profile_limit() -> None:
    with TestClient(main.app) as client:
        session = create_session(client)
        project_id = create_project(client, session)
        run, _ = main.repo.create_run(
            project_id=project_id,
            session_id=session,
            parent_version_id=None,
            idempotency_key="awaiting-clarification",
            message="x" * 1190,
            attachment_ids=[],
            profile="balanced",
        )
        main.repo.update_run(project_id, run.id, status="awaiting_clarification", clarification_questions=["Need one detail"])
        headers = api_headers(**{"x-session-id": session})
        response = client.post(
            f"/v1/projects/{project_id}/generations/{run.id}/clarification",
            headers=headers,
            json={"answer": "y" * 50},
        )
        assert response.status_code == 422
        detail = response.json()["detail"]
        assert detail["profile"] == "balanced"
        assert detail["max_prompt_chars"] == 1200


def test_guest_project_generation_cap_is_enforced_per_project() -> None:
    previous = override_settings(guest_project_generation_cap=1, guest_runs_per_window=10)
    try:
        with TestClient(main.app) as client:
            session = create_session(client)
            project_id = create_project(client, session)
            headers = api_headers(**{"x-session-id": session})
            first = client.post(
                f"/v1/projects/{project_id}/generations",
                headers=headers,
                json={
                    "message": "Base bracket 80x50x6",
                    "parent_version_id": None,
                    "attachment_ids": [],
                    "profile": "balanced",
                    "idempotency_key": "project-cap-first",
                },
            )
            assert first.status_code == 202

            second = client.post(
                f"/v1/projects/{project_id}/generations",
                headers=headers,
                json={
                    "message": "Refine the same bracket with 4 holes",
                    "parent_version_id": None,
                    "attachment_ids": [],
                    "profile": "balanced",
                    "idempotency_key": "project-cap-second",
                },
            )
            assert second.status_code == 429
            detail = second.json()["detail"]
            assert detail["limit_type"] == "project_generation_cap"
            assert detail["project_generation_cap"] == 1
    finally:
        restore_settings(previous)


def test_per_ip_session_cap_blocks_new_guest_sessions() -> None:
    previous = override_settings(
        ip_sessions_per_window=1,
        guest_runs_per_window=10,
    )
    try:
        with TestClient(main.app) as client:
            headers = {**api_headers(), "x-forwarded-for": "203.0.113.42"}
            first = client.post("/v1/auth/guest", headers=headers, json={})
            assert first.status_code == 200
            second = client.post("/v1/auth/guest", headers=headers, json={})
            assert second.status_code == 429
            assert second.json()["detail"]["limit_type"] == "ip_session_cap"
    finally:
        restore_settings(previous)


def test_per_ip_run_cap_blocks_additional_generations() -> None:
    previous = override_settings(
        ip_runs_per_window=1,
        guest_runs_per_window=10,
        guest_project_generation_cap=0,
        guest_project_token_cap=0,
    )
    try:
        with TestClient(main.app) as client:
            headers = {**api_headers(), "x-forwarded-for": "203.0.113.42"}
            session = client.post("/v1/auth/guest", headers=headers, json={}).json()["session_id"]
            project_id = create_project(client, session)
            run_headers = {**headers, "x-session-id": session}
            first = client.post(
                f"/v1/projects/{project_id}/generations",
                headers=run_headers,
                json={
                    "message": "Base bracket 80x50x6",
                    "parent_version_id": None,
                    "attachment_ids": [],
                    "profile": "balanced",
                    "idempotency_key": "ip-run-1",
                },
            )
            assert first.status_code == 202
            second = client.post(
                f"/v1/projects/{project_id}/generations",
                headers=run_headers,
                json={
                    "message": "Add 4 holes",
                    "parent_version_id": None,
                    "attachment_ids": [],
                    "profile": "balanced",
                    "idempotency_key": "ip-run-2",
                },
            )
            assert second.status_code == 429
            assert second.json()["detail"]["limit_type"] == "ip_run_cap"
    finally:
        restore_settings(previous)


def test_guest_project_token_cap_blocks_additional_runs() -> None:
    previous = override_settings(
        guest_project_token_cap=1000,
        vision_summary_max_tokens=220,
        guest_runs_per_window=10,
    )
    try:
        with TestClient(main.app) as client:
            session = create_session(client)
            project_id = create_project(client, session)
            run, _ = main.repo.create_run(
                project_id=project_id,
                session_id=session,
                parent_version_id=None,
                idempotency_key="token-cap-existing-run",
                message="Initial bracket",
                attachment_ids=[],
                profile="balanced",
            )
            main.repo.update_run(
                project_id,
                run.id,
                status="completed",
                telemetry={"spec_usage": {"total_tokens": 650}, "cad_usage": {"total_tokens": 500}},
            )

            response = client.post(
                f"/v1/projects/{project_id}/generations",
                headers=api_headers(**{"x-session-id": session}),
                json={
                    "message": "Make it a little taller",
                    "parent_version_id": None,
                    "attachment_ids": [],
                    "profile": "balanced",
                    "idempotency_key": "token-cap-next-run",
                },
            )
            assert response.status_code == 429
            detail = response.json()["detail"]
            assert detail["limit_type"] == "project_token_cap"
            assert detail["project_token_cap"] == 1000
            assert detail["used_tokens"] == 1150
    finally:
        restore_settings(previous)