Agent/tests/test_local_api_boundary.py

"""Trace API single-user local-development boundary."""

import os
import unittest
from unittest.mock import AsyncMock, patch

from fastapi.testclient import TestClient
from starlette.websockets import WebSocketDisconnect

import api_server
from cyber_agent.trace import websocket as trace_websocket


LOCAL_HOST = {"host": "127.0.0.1:8000"}
LOCAL_ORIGIN = "http://127.0.0.1:3000"
LOCAL_CLIENT = ("127.0.0.1", 50000)


def make_client(peer=LOCAL_CLIENT):
    return TestClient(api_server.app, client=peer)


class LocalApiBoundaryTest(unittest.TestCase):
    def setUp(self) -> None:
        self.list_patch = patch.object(
            api_server.trace_store,
            "list_traces",
            new=AsyncMock(return_value=[]),
        )
        self.list_patch.start()

    def tearDown(self) -> None:
        self.list_patch.stop()

    def test_health_declares_local_development_boundary(self) -> None:
        with make_client() as client:
            response = client.get("/health", headers=LOCAL_HOST)
        self.assertEqual(200, response.status_code)
        self.assertEqual("local_development", response.json()["deployment_mode"])
        self.assertEqual("loopback_only", response.json()["network_scope"])
        self.assertEqual("none", response.json()["authentication"])

    def test_http_rejects_non_local_host(self) -> None:
        for host in ("example.com", "api.internal.example"):
            with self.subTest(host=host), make_client() as client:
                response = client.get("/health", headers={"host": host})
            self.assertEqual(400, response.status_code)

    def test_http_rejects_non_loopback_peer_even_with_local_headers(self) -> None:
        with make_client(("192.0.2.10", 50000)) as client:
            response = client.get(
                "/health",
                headers={**LOCAL_HOST, "origin": LOCAL_ORIGIN},
            )
        self.assertEqual(403, response.status_code)
        self.assertEqual("Loopback client required", response.json()["detail"])

    def test_http_accepts_ipv6_loopback_peer(self) -> None:
        with make_client(("::1", 50000)) as client:
            response = client.get("/health", headers=LOCAL_HOST)
        self.assertEqual(200, response.status_code)

    def test_cors_allows_only_declared_local_frontends(self) -> None:
        for origin in ("http://127.0.0.1:3000", "http://localhost:3000"):
            with self.subTest(origin=origin), make_client() as client:
                response = client.options(
                    "/health",
                    headers={
                        **LOCAL_HOST,
                        "origin": origin,
                        "access-control-request-method": "GET",
                    },
                )
            self.assertEqual(200, response.status_code)
            self.assertEqual(origin, response.headers["access-control-allow-origin"])
            self.assertNotIn("access-control-allow-credentials", response.headers)

    def test_cors_rejects_non_local_origin(self) -> None:
        with make_client() as client:
            response = client.options(
                "/health",
                headers={
                    **LOCAL_HOST,
                    "origin": "https://attacker.example",
                    "access-control-request-method": "GET",
                },
            )
        self.assertEqual(403, response.status_code)

        with make_client() as client:
            simple = client.get(
                "/health",
                headers={
                    **LOCAL_HOST,
                    "origin": "https://attacker.example",
                },
            )
        self.assertEqual(403, simple.status_code)

    def test_websockets_reject_missing_or_non_local_origin(self) -> None:
        for path in (
            "/api/logs/watch",
            "/api/traces/not-a-real-trace/watch",
            "/ws_ping",
        ):
            for origin in (None, "https://attacker.example"):
                headers = dict(LOCAL_HOST)
                if origin is not None:
                    headers["origin"] = origin
                with self.subTest(path=path, origin=origin), make_client() as client:
                    with self.assertRaises(WebSocketDisconnect) as caught:
                        with client.websocket_connect(path, headers=headers):
                            pass
                self.assertEqual(1008, caught.exception.code)

    def test_websockets_reject_non_loopback_peer_even_with_local_headers(self) -> None:
        for path in (
            "/api/logs/watch",
            "/api/traces/not-a-real-trace/watch",
            "/ws_ping",
        ):
            with self.subTest(path=path), make_client(("192.0.2.10", 50000)) as client:
                with self.assertRaises(WebSocketDisconnect) as caught:
                    with client.websocket_connect(
                        path,
                        headers={**LOCAL_HOST, "origin": LOCAL_ORIGIN},
                    ):
                        pass
            self.assertEqual(1008, caught.exception.code)
            self.assertEqual("Loopback client required", caught.exception.reason)

    def test_local_logs_websocket_is_accepted(self) -> None:
        with make_client() as client:
            with client.websocket_connect(
                "/api/logs/watch",
                headers={**LOCAL_HOST, "origin": LOCAL_ORIGIN},
            ) as websocket:
                message = websocket.receive_json()
        self.assertEqual("logs_websocket", message["name"])

    def test_local_trace_websocket_is_accepted(self) -> None:
        class _Trace:
            trace_id = "trace-1"
            last_event_id = 0
            status = "completed"
            parent_trace_id = None

            def to_dict(self):
                return {"trace_id": self.trace_id}

        class _Store:
            async def get_trace(self, trace_id):
                return _Trace() if trace_id == "trace-1" else None

            async def get_goal_tree(self, trace_id):
                return None

            async def list_traces(self, *, limit):
                return []

            async def get_events(self, trace_id, since_event_id):
                return []

        trace_websocket.set_trace_store(_Store())
        try:
            with make_client() as client:
                with client.websocket_connect(
                    "/api/traces/trace-1/watch",
                    headers={**LOCAL_HOST, "origin": LOCAL_ORIGIN},
                ) as websocket:
                    message = websocket.receive_json()
        finally:
            trace_websocket.set_trace_store(api_server.trace_store)
        self.assertEqual("connected", message["event"])

    def test_local_ping_websocket_is_accepted(self) -> None:
        with make_client() as client:
            with client.websocket_connect(
                "/ws_ping",
                headers={**LOCAL_HOST, "origin": LOCAL_ORIGIN},
            ) as websocket:
                self.assertEqual("pong", websocket.receive_text())

    def test_reload_requires_explicit_true(self) -> None:
        for value, expected in (
            (None, False),
            ("true", True),
            ("1", True),
            ("yes", True),
            ("on", True),
            ("false", False),
            ("unexpected", False),
        ):
            environment = {} if value is None else {"TRACE_API_RELOAD": value}
            with self.subTest(value=value), patch.dict(
                os.environ,
                environment,
                clear=False,
            ):
                if value is None:
                    os.environ.pop("TRACE_API_RELOAD", None)
                self.assertIs(expected, api_server._env_flag("TRACE_API_RELOAD"))


if __name__ == "__main__":
    unittest.main()