firegex/tests/integration/test_mcp.py

"""Real MCP over HTTP, then real traffic through the configured datapath."""
import asyncio
import hashlib
from contextlib import asynccontextmanager
import os
import time
import uuid

import pytest
import requests
from mcp import Client
from mcp.client.streamable_http import streamable_http_client, create_mcp_http_client

from helpers.net import free_port
from helpers.traffic import Channel
from integration.filter_code import BLOCK_MARKER
from mcp_api.store import Store

pytestmark = pytest.mark.instance


@pytest.fixture(scope="module")
def token_store():
    path = os.getenv("FIREGEX_MCP_TEST_DB")
    if not path:
        pytest.skip("Set FIREGEX_MCP_TEST_DB to the test instance's MCP database; tests/Dockerfile.mcp does this")
    store = Store(path)
    store.open()
    yield store
    store.close()


@pytest.fixture(scope="module")
def mcp_credentials(token_store):
    admin = token_store.token_create("pytest-admin", "admin", 1)
    observer = token_store.token_create("pytest-observer", "observer", 1)
    yield admin, observer
    token_store.token_revoke(admin["id"])
    token_store.token_revoke(observer["id"])


@asynccontextmanager
async def connect(address, credential):
    async with create_mcp_http_client(headers={"Authorization": "Bearer " + credential["token"]}) as http:
        async with Client(streamable_http_client(address.rstrip("/") + "/mcp/", http_client=http)) as client:
            yield client


async def call(client, name, ok=True, **arguments):
    response = await client.call_tool(name, arguments)
    value = response.structured_content
    assert value is not None, response
    assert value["ok"] is ok, value
    assert response.is_error is not ok, response
    return value


def options(revision=None, **kw):
    return {"request_id": uuid.uuid4().hex, **({"if_revision": revision} if revision else {}), **kw}


async def read_service(client, sid):
    return (await call(client, "fgx_service_get", service=sid))["data"]


def test_protocol_catalog_auth_and_legacy_clients(fg_address, mcp_credentials):
    admin, observer = mcp_credentials
    url = fg_address.rstrip("/") + "/mcp/"
    assert requests.post(url, json={}, timeout=5).status_code == 401
    headers = {"Authorization": "Bearer " + admin["token"], "Accept": "application/json, text/event-stream"}
    initialize = {"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {
        "protocolVersion": "2025-06-18", "capabilities": {}, "clientInfo": {"name": "legacy-test", "version": "1"}}}
    response = requests.post(url, json=initialize, headers=headers, timeout=5)
    assert response.status_code == 200, response.text
    assert response.json()["result"]["protocolVersion"] == "2025-06-18"
    assert requests.post(url, json=initialize, headers={**headers, "Origin": "https://untrusted.example"}, timeout=5).status_code == 403
    assert requests.post(url, json=initialize, headers={**headers, "Host": "untrusted.example"}, timeout=5).status_code == 421

    async def scenario():
        async with connect(fg_address, admin) as client:
            listed = await client.list_tools()
            assert len(listed.tools) == 28
            assert all(t.output_schema and "ok" in t.output_schema["properties"] for t in listed.tools)
            names = {t.name for t in listed.tools}
            assert {"fgx_filter_apply", "fgx_change_rollback", "fgx_instance_admin"} <= names
            await call(client, "fgx_overview", include_interfaces=True)
            await call(client, "fgx_help", topic="pyfilter", query="RawPacket")
        async with connect(fg_address, observer) as client:
            await call(client, "fgx_overview")
            denied = await call(client, "fgx_python_check", ok=False, code="pass")
            assert denied["error"]["code"] == "FORBIDDEN"
            await call(client, "fgx_config_export", ok=False)
    asyncio.run(scenario())


def test_revocation_takes_effect_on_next_request(fg_address, token_store):
    credential = token_store.token_create("revocation-test", "observer", 1)
    async def scenario():
        async with connect(fg_address, credential) as client:
            await call(client, "fgx_overview")
            token_store.token_revoke(credential["id"])
            with pytest.raises(Exception):
                await client.call_tool("fgx_overview", {})
    asyncio.run(scenario())


def test_service_create_preview_retry_conflict_and_delete(fg_address, mcp_credentials, api):
    created = []
    async def scenario():
        async with connect(fg_address, mcp_credentials[0]) as client:
            name = "mcp-" + uuid.uuid4().hex[:8]
            spec = {"name": name, "addresses": [{"ip_int": "127.0.0.1", "port": free_port()}]}
            preview = await call(client, "fgx_service_apply", spec=spec, options=options(dry_run=True))
            assert preview["data"]["dry_run"]
            assert not (await call(client, "fgx_service_list", query=name))["data"]["items"]
            request_options = options()
            first = await call(client, "fgx_service_apply", spec=spec, options=request_options)
            sid = first["data"]["service_id"]
            created.append(sid)
            repeat = await call(client, "fgx_service_apply", spec=spec, options=request_options)
            assert repeat["data"]["service_id"] == sid
            assert repeat["change"]["id"] == first["change"]["id"]
            mismatch = await call(client, "fgx_service_apply", ok=False, spec={**spec, "name": "different"}, options=request_options)
            assert mismatch["error"]["code"] == "REQUEST_ID_CONFLICT"
            before = await read_service(client, sid)
            updated = await call(client, "fgx_service_apply", service=sid, spec={"name": name + "-edited"}, options=options(before["revision"]))
            stale = await call(client, "fgx_service_apply", ok=False, service=sid, spec={"name": "stale"}, options=options(before["revision"]))
            assert stale["error"]["code"] == "REVISION_CONFLICT"
            await call(client, "fgx_change_rollback", change_id=updated["change"]["id"])
            restored = await read_service(client, sid)
            assert restored["name"] == name
            await call(client, "fgx_service_delete", service=sid, options=options(restored["revision"]))
            await call(client, "fgx_service_get", ok=False, service=sid)
    try:
        asyncio.run(scenario())
    finally:
        for sid in created:
            api.services_delete(sid)


def test_batch_filter_real_traffic_rejection_and_rollback(fg_address, mcp_credentials, protected, filtering_layer):
    sid, server, port = protected(filtering_layer, name="mcp-rx")
    channel = Channel(server, port, filtering_layer.ipv6, filtering_layer.tls)
    async def scenario():
        async with connect(fg_address, mcp_credentials[0]) as client:
            before = await read_service(client, sid)
            result = await call(client, "fgx_filter_apply", service=sid, spec={"kind": "regex", "name": "guard",
                "patterns": [{"pattern": "MCPBLOCK"}, {"pattern": "SECOND", "case_sensitive": False}]}, options=options(before["revision"]))
            fid = result["data"]["filter_id"]
            active = await call(client, "fgx_service_set_state", service=sid, state="running", options=options(result["data"]["revision"]))
            await asyncio.sleep(0.5)
            assert await asyncio.to_thread(channel.gets_through, b"ordinary request")
            assert await asyncio.to_thread(channel.is_blocked, b"MCPBLOCK")
            assert await asyncio.to_thread(channel.is_blocked, b"second")
            current = await read_service(client, sid)
            bad = await call(client, "fgx_filter_apply", ok=False, service=sid, filter=fid,
                             spec={"patterns": [{"pattern": "(unclosed"}]}, options=options(current["revision"]))
            assert bad["error"]["code"] == "FILTER_COMPILE_ERROR"
            assert (await read_service(client, sid))["revision"] == current["revision"]
            update = await call(client, "fgx_filter_apply", service=sid, filter=fid,
                                spec={"patterns": [{"pattern": "CHANGED"}]}, options=options(current["revision"]))
            assert await asyncio.to_thread(channel.gets_through, b"MCPBLOCK")
            assert await asyncio.to_thread(channel.is_blocked, b"CHANGED")
            await call(client, "fgx_change_rollback", change_id=update["change"]["id"])
            assert await asyncio.to_thread(channel.is_blocked, b"MCPBLOCK")
            assert await asyncio.to_thread(channel.gets_through, b"CHANGED")
            stats = await call(client, "fgx_stats_get", service=sid)
            assert stats["data"]["total"] >= 1
            logs = await call(client, "fgx_logs_get", service=sid, limit=2)
            await call(client, "fgx_logs_get", service=sid, cursor=logs["data"]["next_cursor"])
    asyncio.run(scenario())


def test_python_check_apply_toggle_functions_and_dry_run(fg_address, mcp_credentials, protected, inspecting_layer):
    sid, server, port = protected(inspecting_layer, name="mcp-py")
    channel = Channel(server, port, inspecting_layer.ipv6)
    async def scenario():
        async with connect(fg_address, mcp_credentials[0]) as client:
            checked = await call(client, "fgx_python_check", service=sid, code=BLOCK_MARKER)
            assert checked["data"]["ok"] and checked["data"]["filters"] == ["refuse_marker"], checked["data"]
            rejected = await call(client, "fgx_python_check", code="def broken(:")
            assert not rejected["data"]["ok"]
            before = await read_service(client, sid)
            spec = {"kind": "pyfilter", "name": "python", "code": BLOCK_MARKER}
            await call(client, "fgx_filter_apply", service=sid, spec=spec, options=options(before["revision"], dry_run=True))
            assert not (await read_service(client, sid))["filters"]
            result = await call(client, "fgx_filter_apply", service=sid, spec=spec, options=options(before["revision"]))
            fid = result["data"]["filter_id"]
            await call(client, "fgx_service_set_state", service=sid, state="running", options=options(result["data"]["revision"]))
            await asyncio.sleep(0.5)
            assert await asyncio.to_thread(channel.gets_through, b"ordinary request")
            assert await asyncio.to_thread(channel.is_blocked, b"PYBLOCK")
            current = await read_service(client, sid)
            await call(client, "fgx_filter_apply", service=sid, filter=fid, spec={"functions": {"refuse_marker": False}},
                       options=options(current["revision"]))
            assert await asyncio.to_thread(channel.gets_through, b"PYBLOCK")
            detail = await call(client, "fgx_filter_get", service=sid, filter=fid)
            assert not detail["data"]["functions"][0]["active"]
    asyncio.run(scenario())


def test_rest_edits_invalidate_mcp_revision(fg_address, mcp_credentials, service, api):
    sid = service("mcp-conflict", "127.0.0.1", free_port(), "proxy")
    async def scenario():
        async with connect(fg_address, mcp_credentials[0]) as client:
            before = await read_service(client, sid)
            assert api.services_edit(sid, name="changed-in-ui")
            result = await call(client, "fgx_service_apply", ok=False, service=sid,
                                spec={"name": "overwrite"}, options=options(before["revision"]))
            assert result["error"]["code"] == "REVISION_CONFLICT"
            assert (await read_service(client, sid))["name"] == "changed-in-ui"
    asyncio.run(scenario())


def test_firewall_confirmation_and_expiry(fg_address, mcp_credentials):
    async def scenario():
        async with connect(fg_address, mcp_credentials[0]) as client:
            before = (await call(client, "fgx_firewall_get"))["data"]
            change = await call(client, "fgx_firewall_apply", patch={"settings": {"allow_icmp": not before["settings"]["allow_icmp"]}},
                                options=options(before["revision"], confirm_within_seconds=5))
            await asyncio.sleep(5.8)
            assert (await call(client, "fgx_firewall_get"))["data"]["revision"] == before["revision"]
            status = await call(client, "fgx_change_get", change_id=change["change"]["id"])
            assert status["data"]["status"] == "rolled_back"
            change = await call(client, "fgx_firewall_apply", patch={"settings": {"allow_icmp": not before["settings"]["allow_icmp"]}},
                                options=options(before["revision"], confirm_within_seconds=5))
            await call(client, "fgx_change_confirm", change_id=change["change"]["id"])
            await call(client, "fgx_change_rollback", change_id=change["change"]["id"])
            assert (await call(client, "fgx_firewall_get"))["data"]["revision"] == before["revision"]
    asyncio.run(scenario())


def test_snapshot_export_import_and_jobs(fg_address, mcp_credentials, service, datapath_in_our_path):
    if not datapath_in_our_path:
        pytest.skip("Import requires a working Firegex datapath initialization")
    sid = service("mcp-snapshot", "127.0.0.1", free_port(), "proxy")
    async def scenario():
        async with connect(fg_address, mcp_credentials[0]) as client:
            before = await read_service(client, sid)
            saved = await call(client, "fgx_snapshot_create", scope="service", service=sid, label="test")
            await call(client, "fgx_snapshot_list")
            job = await call(client, "fgx_service_apply", service=sid, spec={"name": "mcp-background"},
                             options=options(before["revision"], background=True))
            job_id = job["data"]["job_id"]
            for _ in range(50):
                state = await call(client, "fgx_job", job_id=job_id)
                if state["data"]["status"] not in {"queued", "applying"}:
                    break
                await asyncio.sleep(0.1)
            assert state["data"]["status"] == "applied", state
            current = await read_service(client, sid)
            await call(client, "fgx_snapshot_restore", snapshot_id=saved["data"]["id"], options=options(current["revision"]))
            assert (await read_service(client, sid))["name"] == before["name"]
            artifact = (await call(client, "fgx_config_export"))["data"]
            url = fg_address.rstrip("/") + artifact["download_path"]
            assert requests.get(url, timeout=5).status_code == 401
            downloaded = await asyncio.to_thread(requests.get, url, headers={"Authorization": "Bearer " + mcp_credentials[0]["token"]}, timeout=5)
            assert downloaded.status_code == 200
            assert hashlib.sha256(downloaded.content).hexdigest() == artifact["sha256"]
            assert len(downloaded.content) == artifact["bytes"]
            assert "control.sqlite3" not in downloaded.text
            overview = (await call(client, "fgx_overview"))["data"]
            await call(client, "fgx_config_import", artifact_id=artifact["id"], options=options(overview["instance_revision"], dry_run=True))
            imported = await call(client, "fgx_config_import", artifact_id=artifact["id"], options=options(overview["instance_revision"]))
            await call(client, "fgx_change_confirm", change_id=imported["change"]["id"])
            assert (await read_service(client, sid))["name"] == before["name"]
    asyncio.run(scenario())


def test_chain_order_delete_and_binary_regex_samples(fg_address, mcp_credentials, service):
    sid = service("mcp-order", "127.0.0.1", free_port(), "proxy")
    async def scenario():
        async with connect(fg_address, mcp_credentials[0]) as client:
            sample = await call(client, "fgx_regex_test", patterns=[{"pattern": "MARK"}],
                                samples=[{"encoding": "hex", "data": "004d41524bff"}])
            assert sample["data"]["samples"][0]["matches"]
            for name in ("first", "second"):
                current = await read_service(client, sid)
                await call(client, "fgx_filter_apply", service=sid, spec={"kind": "regex", "name": name},
                           options=options(current["revision"]))
            current = await read_service(client, sid)
            result = await call(client, "fgx_chain_reorder", service=sid, filters=["second", "first"],
                                options=options(current["revision"]))
            assert [f["name"] for f in (await read_service(client, sid))["filters"]] == ["second", "first"]
            result = await call(client, "fgx_filter_delete", service=sid, filter="second", options=options(result["data"]["revision"]))
            assert [f["name"] for f in (await read_service(client, sid))["filters"]] == ["first"]
            await call(client, "fgx_change_rollback", change_id=result["change"]["id"])
            assert [f["name"] for f in (await read_service(client, sid))["filters"]] == ["second", "first"]
    asyncio.run(scenario())


def test_invalid_import_and_admin_role_boundary(fg_address, mcp_credentials, token_store):
    operator = token_store.token_create("operator-boundary", "operator", 1)
    async def scenario():
        async with connect(fg_address, operator) as client:
            before = (await call(client, "fgx_overview"))["data"]["instance_revision"]
            await call(client, "fgx_config_import", ok=False, configuration={"../escape.db": {}}, options=options(before))
            assert (await call(client, "fgx_overview"))["data"]["instance_revision"] == before
            denied = await call(client, "fgx_instance_admin", ok=False, action="factory_reset", options=options(before))
            assert denied["error"]["code"] == "FORBIDDEN"
    try:
        asyncio.run(scenario())
    finally:
        token_store.token_revoke(operator["id"])


def test_enabled_firewall_rules_and_rollback(fg_address, mcp_credentials):
    import socket
    from helpers.tcpserver import TcpServer
    port = free_port()
    server = TcpServer(port, False)
    server.start()
    for _ in range(100):
        try:
            with socket.create_connection(("127.0.0.1", port), timeout=1):
                break
        except ConnectionRefusedError:
            time.sleep(0.1)
    async def scenario():
        async with connect(fg_address, mcp_credentials[0]) as client:
            saved = (await call(client, "fgx_snapshot_create", scope="firewall"))["data"]
            try:
                assert await asyncio.to_thread(server.sendCheckData, b"baseline")
                before = (await call(client, "fgx_firewall_get"))["data"]
                change = await call(client, "fgx_firewall_apply", patch={"policy": "accept",
                    "settings": {"allow_loopback": False, "allow_established": False},
                    "rules": [{"name": "test-port", "mode": "out", "dst": "127.0.0.1",
                               "port_dst_from": port, "port_dst_to": port, "action": "reject"}]},
                    options=options(before["revision"]))
                await call(client, "fgx_change_confirm", change_id=change["change"]["id"])
                enabled = await call(client, "fgx_firewall_set_state", enabled=True, options=options(change["data"]["revision"]))
                await call(client, "fgx_change_confirm", change_id=enabled["change"]["id"])
                with pytest.raises((ConnectionRefusedError, ConnectionResetError, TimeoutError)):
                    socket.create_connection(("127.0.0.1", port), timeout=3)
                read = (await call(client, "fgx_firewall_get"))["data"]
                assert read["enabled"] and len(read["rules"]) == 1
                assert read["rules"][0]["id"]
            finally:
                current = (await call(client, "fgx_firewall_get"))["data"]
                restored = await call(client, "fgx_snapshot_restore", snapshot_id=saved["id"], options=options(current["revision"]))
                await call(client, "fgx_change_confirm", change_id=restored["change"]["id"])
            assert await asyncio.to_thread(server.sendCheckData, b"restored")
    try:
        asyncio.run(scenario())
    finally:
        server.stop()


def test_mcp_service_address_and_tls_edit(fg_address, mcp_credentials, service, stand_in, certificate, datapath_in_our_path):
    if not datapath_in_our_path:
        pytest.skip("firegex's datapath is not in front of this host's loopback")
    server = stand_in()
    sid = service("mcp-address-edit", "127.0.0.1", server.port, "proxy")
    cert, key = certificate("127.0.0.1")
    channel = Channel(server, server.port, False, False)
    async def scenario():
        async with connect(fg_address, mcp_credentials[0]) as client:
            current = await read_service(client, sid)
            filtered = await call(client, "fgx_filter_apply", service=sid,
                spec={"kind": "regex", "patterns": [{"pattern": "BLOCK_TLS_EDIT"}]}, options=options(current["revision"]))
            await call(client, "fgx_service_set_state", service=sid, state="running", options=options(filtered["data"]["revision"]))
            assert await asyncio.to_thread(channel.gets_through, b"ordinary plain")
            assert await asyncio.to_thread(channel.is_blocked, b"BLOCK_TLS_EDIT")
            current = await read_service(client, sid)
            public_port = free_port()
            edited = await call(client, "fgx_service_apply", service=sid,
                spec={"proto": "tls", "tls_cert": cert, "tls_key": key, "addresses": [
                    {"ip_int": "127.0.0.1", "port": public_port, "edge": "tls", "target_port": server.port, "upstream": "tcp"}]},
                options=options(current["revision"]))
            tls_channel = Channel(server, public_port, False, True)
            assert await asyncio.to_thread(tls_channel.gets_through, b"ordinary tls")
            assert await asyncio.to_thread(tls_channel.is_blocked, b"BLOCK_TLS_EDIT")
            await call(client, "fgx_change_rollback", change_id=edited["change"]["id"])
            assert await asyncio.to_thread(channel.gets_through, b"ordinary plain")
            assert await asyncio.to_thread(channel.is_blocked, b"BLOCK_TLS_EDIT")
    asyncio.run(scenario())