diff --git a/aikido_zen/background_process/__init__.py b/aikido_zen/background_process/__init__.py index 23d1d69b3..e145de8e1 100644 --- a/aikido_zen/background_process/__init__.py +++ b/aikido_zen/background_process/__init__.py @@ -11,6 +11,7 @@ from aikido_zen.helpers.token import get_token_from_env from aikido_zen.helpers.get_temp_dir import get_temp_dir from aikido_zen.helpers.hash_aikido_token import hash_aikido_token +from aikido_zen.helpers.get_agent_session_id import get_agent_session_id from aikido_zen.helpers.logging import logger from aikido_zen.background_process.comms import ( @@ -71,7 +72,7 @@ def start_background_process(): # Daemon is set to True so that the process kills itself when the main process dies background_process = get_process_factory()( target=AikidoBackgroundProcess, - args=(comms.address, comms.key), + args=(comms.address, comms.key, get_agent_session_id()), name="zen-agent-process", daemon=True, ) diff --git a/aikido_zen/background_process/aikido_background_process.py b/aikido_zen/background_process/aikido_background_process.py index b4b1b3f37..57387a574 100644 --- a/aikido_zen/background_process/aikido_background_process.py +++ b/aikido_zen/background_process/aikido_background_process.py @@ -11,6 +11,7 @@ import platform from threading import Thread from queue import Queue +from aikido_zen.helpers.get_agent_session_id import set_agent_session_id from aikido_zen.helpers.logging import logger from aikido_zen.background_process.cloud_connection_manager import ( CloudConnectionManager, @@ -33,7 +34,9 @@ class AikidoBackgroundProcess: - (spawned) reporting thread which collects IPC data and send it to a CloudConnectionManager """ - def __init__(self, address, key): + def __init__(self, address, key, session_id): + # Spawn re-imports this module and would otherwise create another session id. + set_agent_session_id(session_id) try: listener = con.Listener(address, authkey=None) except OSError: diff --git a/aikido_zen/background_process/cloud_connection_manager/get_manager_info.py b/aikido_zen/background_process/cloud_connection_manager/get_manager_info.py index d733b6010..5446277bf 100644 --- a/aikido_zen/background_process/cloud_connection_manager/get_manager_info.py +++ b/aikido_zen/background_process/cloud_connection_manager/get_manager_info.py @@ -1,9 +1,9 @@ """Exports get_manager_info function""" -import socket import platform import aikido_zen.config as config import aikido_zen.helpers.get_machine_ip as h +from aikido_zen.helpers.get_hostname import get_hostname def get_manager_info(connection_manager): @@ -12,9 +12,9 @@ def get_manager_info(connection_manager): """ return { "dryMode": not connection_manager.block, - "hostname": socket.gethostname(), + "hostname": get_hostname(), "version": config.PKG_VERSION, - "library": "firewall-python", + "library": config.LIBRARY_NAME, "ipAddress": h.get_ip(), "serverless": bool(connection_manager.serverless), "stack": ( diff --git a/aikido_zen/background_process/cloud_connection_manager/get_manager_info_test.py b/aikido_zen/background_process/cloud_connection_manager/get_manager_info_test.py index 8c1105867..032fa6b5b 100644 --- a/aikido_zen/background_process/cloud_connection_manager/get_manager_info_test.py +++ b/aikido_zen/background_process/cloud_connection_manager/get_manager_info_test.py @@ -1,5 +1,6 @@ import pytest import socket +from aikido_zen.helpers.get_hostname import get_hostname from unittest.mock import MagicMock, patch from .get_manager_info import get_manager_info @@ -22,6 +23,7 @@ def mock_connection_manager(): # Fixture to patch socket.gethostname @pytest.fixture(autouse=True) def patch_gethostname(monkeypatch): + get_hostname.cache_clear() monkeypatch.setattr(socket, "gethostname", lambda: "test-hostname") diff --git a/aikido_zen/background_process/get_common_agent_headers.py b/aikido_zen/background_process/get_common_agent_headers.py new file mode 100644 index 000000000..74c56f0fa --- /dev/null +++ b/aikido_zen/background_process/get_common_agent_headers.py @@ -0,0 +1,24 @@ +import aikido_zen.config as config +from aikido_zen.helpers.get_agent_session_id import get_agent_session_id +from aikido_zen.helpers.get_hostname import get_hostname +from aikido_zen.helpers.get_machine_ip import get_ip + + +def get_common_agent_headers(): + return { + "X-Agent-Platform": "python", + "X-Agent-Library": config.LIBRARY_NAME, + "X-Agent-Version": config.PKG_VERSION, + "X-Agent-Hostname": _header_value(get_hostname()), + "X-Agent-IP-Address": _header_value(get_ip()), + "X-Agent-Session-Id": get_agent_session_id(), + } + + +def _header_value(value): + # Header values are latin-1, and one that cannot be encoded fails the request. + try: + value.encode("latin-1") + except UnicodeEncodeError: + return "unknown" + return value or "unknown" diff --git a/aikido_zen/background_process/get_common_agent_headers_test.py b/aikido_zen/background_process/get_common_agent_headers_test.py new file mode 100644 index 000000000..ace227316 --- /dev/null +++ b/aikido_zen/background_process/get_common_agent_headers_test.py @@ -0,0 +1,77 @@ +import socket +from unittest.mock import MagicMock + +import pytest + +import aikido_zen.config as config +import aikido_zen.background_process.get_common_agent_headers as headers_module +from aikido_zen.background_process.get_common_agent_headers import ( + get_common_agent_headers, +) +from aikido_zen.helpers.get_agent_session_id import get_agent_session_id +from aikido_zen.helpers.get_hostname import get_hostname +from aikido_zen.helpers.get_machine_ip import get_ip +from aikido_zen.background_process.cloud_connection_manager.get_manager_info import ( + get_manager_info, +) + + +@pytest.fixture(autouse=True) +def reset_cached_instance_values(): + get_ip.cache_clear() + get_hostname.cache_clear() + + +def test_common_agent_headers(monkeypatch): + monkeypatch.setattr(headers_module, "get_ip", lambda: "10.0.0.1") + + assert get_common_agent_headers() == { + "X-Agent-Platform": "python", + "X-Agent-Library": config.LIBRARY_NAME, + "X-Agent-Version": config.PKG_VERSION, + "X-Agent-Hostname": get_hostname() or "unknown", + "X-Agent-IP-Address": "10.0.0.1", + "X-Agent-Session-Id": get_agent_session_id(), + } + + +def test_common_agent_headers_keep_the_same_session_id(): + assert ( + get_common_agent_headers()["X-Agent-Session-Id"] + == get_common_agent_headers()["X-Agent-Session-Id"] + ) + + +def test_hostname_and_ip_match_the_event_payload(monkeypatch): + monkeypatch.setattr(socket, "gethostname", lambda: "instance-app-7c9d") + monkeypatch.setattr(socket, "gethostbyname", lambda _hostname: "10.0.1.42") + manager_info = get_manager_info(MagicMock()) + headers = get_common_agent_headers() + + assert headers["X-Agent-Hostname"] == manager_info["hostname"] + assert headers["X-Agent-IP-Address"] == manager_info["ipAddress"] + + +def test_hostname_that_cannot_be_a_header_is_unknown(monkeypatch): + monkeypatch.setattr(socket, "gethostname", lambda: "p\u0101yments") + + assert get_common_agent_headers()["X-Agent-Hostname"] == "unknown" + + +def test_missing_hostname_is_unknown(monkeypatch): + monkeypatch.setattr(socket, "gethostname", lambda: "") + + assert get_common_agent_headers()["X-Agent-Hostname"] == "unknown" + + +def test_unresolved_ip_address_is_unknown(monkeypatch): + monkeypatch.setattr(headers_module, "get_ip", lambda: "") + + assert get_common_agent_headers()["X-Agent-IP-Address"] == "unknown" + + +def test_library_matches_the_event_payload(): + assert ( + get_common_agent_headers()["X-Agent-Library"] + == get_manager_info(MagicMock())["library"] + ) diff --git a/aikido_zen/background_process/init_test.py b/aikido_zen/background_process/init_test.py index 7483efe06..702bc7cae 100644 --- a/aikido_zen/background_process/init_test.py +++ b/aikido_zen/background_process/init_test.py @@ -4,6 +4,13 @@ import pytest import aikido_zen.background_process as background_process +from aikido_zen.background_process.aikido_background_process import ( + AikidoBackgroundProcess, +) +from aikido_zen.helpers.get_agent_session_id import ( + get_agent_session_id, + set_agent_session_id, +) from .comms_test import reset_comms_after_test @@ -164,3 +171,36 @@ def test_stale_socket_removed_by_another_worker(monkeypatch, mocker): background_process.start_background_process() process.return_value.start.assert_called_once_with() + + +def test_started_process_receives_the_current_session_id(monkeypatch, mocker): + process = mocker.patch( + "aikido_zen.background_process.get_process_factory" + ).return_value + monkeypatch.setenv("AIKIDO_TOKEN", "AIK_RUNTIME_TEST") + monkeypatch.setattr(background_process.platform, "system", lambda: "Linux") + monkeypatch.setattr( + background_process, "get_uds_filename", lambda: "/tmp/aikido-session.sock" + ) + monkeypatch.setattr(background_process.os.path, "exists", lambda _path: False) + + background_process.start_background_process() + + assert process.call_args.kwargs["args"][2] == get_agent_session_id() + + +def test_background_process_adopts_the_parent_session_id(monkeypatch): + def listener(*_args, **_kwargs): + raise OSError("address in use") + + original = get_agent_session_id() + monkeypatch.setattr( + "aikido_zen.background_process.aikido_background_process.con.Listener", + listener, + ) + try: + with pytest.raises(SystemExit): + AikidoBackgroundProcess("/tmp/zen-session.sock", b"key", "parent-session") + assert get_agent_session_id() == "parent-session" + finally: + set_agent_session_id(original) diff --git a/aikido_zen/background_process/realtime/sse_client/__init__.py b/aikido_zen/background_process/realtime/sse_client/__init__.py index d94aa50e2..132118ea1 100644 --- a/aikido_zen/background_process/realtime/sse_client/__init__.py +++ b/aikido_zen/background_process/realtime/sse_client/__init__.py @@ -13,9 +13,11 @@ import urllib.request import aikido_zen.background_process.realtime as realtime +from aikido_zen.background_process.get_common_agent_headers import ( + get_common_agent_headers, +) from aikido_zen.helpers.logging import logger from .parser import SSEParser -import aikido_zen.config as config INITIAL_RECONNECT_SECS = 5 MAX_RECONNECT_SECS = 60 @@ -92,11 +94,10 @@ def _connect(token, on_event, read_timeout_secs): url, method="GET", headers={ + **get_common_agent_headers(), "Authorization": str(token), "Accept": "text/event-stream", "Cache-Control": "no-cache", - "X-Agent-Platform": "python", - "X-Agent-Version": config.PKG_VERSION, }, ) diff --git a/aikido_zen/background_process/realtime/sse_client/real_server_test.py b/aikido_zen/background_process/realtime/sse_client/real_server_test.py index 7a216879a..7f5f3bc06 100644 --- a/aikido_zen/background_process/realtime/sse_client/real_server_test.py +++ b/aikido_zen/background_process/realtime/sse_client/real_server_test.py @@ -10,6 +10,9 @@ import time from unittest.mock import patch +from aikido_zen.background_process.get_common_agent_headers import ( + get_common_agent_headers, +) from aikido_zen.helpers.token import Token from . import _connect, _reconnect_loop @@ -65,6 +68,9 @@ def do_get(handler): captured["cache"] = handler.headers.get("Cache-Control") captured["x-agent-platform"] = handler.headers.get("X-Agent-Platform") captured["x-agent-version"] = handler.headers.get("X-Agent-Version") + captured["x-agent-hostname"] = handler.headers.get("X-Agent-Hostname") + captured["x-agent-ip-address"] = handler.headers.get("X-Agent-IP-Address") + captured["x-agent-session-id"] = handler.headers.get("X-Agent-Session-Id") handler.send_response(200) handler.send_header("Content-Type", "text/event-stream") handler.send_header("Cache-Control", "no-cache") @@ -81,6 +87,7 @@ def do_get(handler): handler.wfile.flush() server, url = start_server(make_handler(do_get)) + expected_headers = get_common_agent_headers() try: with patch(REALTIME_URL_PATCH_TARGET, return_value=url): outcome, status_code = _connect( @@ -96,6 +103,9 @@ def do_get(handler): assert captured["cache"] == "no-cache" assert captured["x-agent-platform"] == "python" assert captured["x-agent-version"] == "1.0-REPLACE-VERSION" + assert captured["x-agent-hostname"] == expected_headers["X-Agent-Hostname"] + assert captured["x-agent-ip-address"] == expected_headers["X-Agent-IP-Address"] + assert captured["x-agent-session-id"] == expected_headers["X-Agent-Session-Id"] assert len(events) == 2 assert events[0].event == "config-updated" assert json.loads(events[0].data) == {"configUpdatedAt": 100} diff --git a/aikido_zen/background_process/requests/make_request.py b/aikido_zen/background_process/requests/make_request.py index e89b73295..541375fdd 100644 --- a/aikido_zen/background_process/requests/make_request.py +++ b/aikido_zen/background_process/requests/make_request.py @@ -4,6 +4,9 @@ import urllib.request import urllib.error +from aikido_zen.background_process.get_common_agent_headers import ( + get_common_agent_headers, +) from aikido_zen.background_process.requests.errors import TimeoutExceeded @@ -11,9 +14,8 @@ def make_request(method, url, timeout, data=None, headers=None): req = urllib.request.Request(url, data=data, method=method) # Add headers - if headers: - for key, value in headers.items(): - req.add_header(key, value) + for key, value in {**get_common_agent_headers(), **(headers or {})}.items(): + req.add_header(key, value) try: with urllib.request.urlopen(req, timeout=timeout) as response: return Response(response) diff --git a/aikido_zen/background_process/requests/make_request_test.py b/aikido_zen/background_process/requests/make_request_test.py new file mode 100644 index 000000000..5d4f301b0 --- /dev/null +++ b/aikido_zen/background_process/requests/make_request_test.py @@ -0,0 +1,71 @@ +import http.server +import threading +from unittest.mock import patch + +from aikido_zen.background_process.get_common_agent_headers import ( + get_common_agent_headers, +) +from aikido_zen.background_process.requests.make_request import make_request + + +class _Response: + def __enter__(self): + return self + + def __exit__(self, *args): + return False + + def getcode(self): + return 200 + + def info(self): + return {} + + def read(self): + return b"" + + +def test_make_request_sends_agent_headers(): + captured = {} + + def fake_urlopen(request, timeout): + captured["request"] = request + return _Response() + + expected = get_common_agent_headers() + with patch("urllib.request.urlopen", side_effect=fake_urlopen): + make_request( + "GET", + "https://example.test/api", + 1, + headers={"Authorization": "token"}, + ) + + headers = { + name.lower(): value for name, value in captured["request"].header_items() + } + assert headers["authorization"] == "token" + for name, value in expected.items(): + assert headers[name.lower()] == value + + +def test_make_request_with_a_hostname_that_cannot_be_a_header(): + class Handler(http.server.BaseHTTPRequestHandler): + def do_GET(self): + self.send_response(200) + self.end_headers() + + def log_message(self, *args): + pass + + server = http.server.HTTPServer(("127.0.0.1", 0), Handler) + threading.Thread(target=server.serve_forever, daemon=True).start() + try: + with patch("socket.gethostname", return_value="p\u0101yments"): + response = make_request( + "GET", "http://127.0.0.1:%d/" % server.server_address[1], 5 + ) + finally: + server.shutdown() + + assert response.status_code == 200 diff --git a/aikido_zen/config.py b/aikido_zen/config.py index 3bab3ce55..b669b5207 100644 --- a/aikido_zen/config.py +++ b/aikido_zen/config.py @@ -1,3 +1,4 @@ """Contains package versions""" PKG_VERSION = "1.0-REPLACE-VERSION" +LIBRARY_NAME = "firewall-python" diff --git a/aikido_zen/helpers/get_agent_session_id.py b/aikido_zen/helpers/get_agent_session_id.py new file mode 100644 index 000000000..53134423c --- /dev/null +++ b/aikido_zen/helpers/get_agent_session_id.py @@ -0,0 +1,14 @@ +from aikido_zen.helpers.uuid7 import uuid7 + +# Fork copies this value, so workers forked after import share one run. +_session_id = uuid7() + + +def get_agent_session_id(): + return _session_id + + +def set_agent_session_id(session_id): + # pylint: disable=global-statement + global _session_id + _session_id = session_id diff --git a/aikido_zen/helpers/get_agent_session_id_test.py b/aikido_zen/helpers/get_agent_session_id_test.py new file mode 100644 index 000000000..82e7a33b0 --- /dev/null +++ b/aikido_zen/helpers/get_agent_session_id_test.py @@ -0,0 +1,27 @@ +import os + +import pytest + +from aikido_zen.helpers.get_agent_session_id import get_agent_session_id + + +def test_session_id_stays_the_same(): + assert get_agent_session_id() + assert get_agent_session_id() == get_agent_session_id() + + +@pytest.mark.skipif(not hasattr(os, "fork"), reason="fork is not available") +def test_forked_workers_share_the_session_id(): + read_fd, write_fd = os.pipe() + pid = os.fork() + if pid == 0: + os.close(read_fd) + os.write(write_fd, get_agent_session_id().encode()) + os._exit(0) # pylint: disable=protected-access + + os.close(write_fd) + session_id = os.read(read_fd, 64).decode() + os.close(read_fd) + os.waitpid(pid, 0) + + assert session_id == get_agent_session_id() diff --git a/aikido_zen/helpers/get_hostname.py b/aikido_zen/helpers/get_hostname.py new file mode 100644 index 000000000..feab27bcb --- /dev/null +++ b/aikido_zen/helpers/get_hostname.py @@ -0,0 +1,15 @@ +""" +Helper function file, see function docstring +""" + +import socket +from functools import lru_cache + + +@lru_cache(maxsize=1) +def get_hostname(): + """Tries to fetch the hostname and returns an empty string on failure""" + try: + return socket.gethostname() + except OSError: + return "" diff --git a/aikido_zen/helpers/get_hostname_test.py b/aikido_zen/helpers/get_hostname_test.py new file mode 100644 index 000000000..5dc49d525 --- /dev/null +++ b/aikido_zen/helpers/get_hostname_test.py @@ -0,0 +1,34 @@ +import pytest +import socket +from aikido_zen.helpers.get_hostname import get_hostname + + +@pytest.fixture(autouse=True) +def reset_cached_hostname(): + get_hostname.cache_clear() + + +def test_get_hostname_success(monkeypatch): + monkeypatch.setattr(socket, "gethostname", lambda: "mocked_hostname") + + assert get_hostname() == "mocked_hostname" + + +def test_get_hostname_failure(monkeypatch): + monkeypatch.setattr( + socket, + "gethostname", + lambda: (_ for _ in ()).throw(OSError("Mocked exception")), + ) + + assert get_hostname() == "" + + +def test_get_hostname_is_read_once(monkeypatch): + calls = [] + monkeypatch.setattr( + socket, "gethostname", lambda: calls.append(None) or "mocked_hostname" + ) + + assert get_hostname() == get_hostname() + assert len(calls) == 1 diff --git a/aikido_zen/helpers/get_machine_ip.py b/aikido_zen/helpers/get_machine_ip.py index d36eef729..a8b080031 100644 --- a/aikido_zen/helpers/get_machine_ip.py +++ b/aikido_zen/helpers/get_machine_ip.py @@ -3,11 +3,14 @@ """ import socket +from functools import lru_cache +@lru_cache(maxsize=1) def get_ip(): - """Tries to fetch the IP and returns x.x.x.x on failure""" + """Tries to fetch the IP and returns an empty string on failure""" + # Cached: this queries DNS and runs on every report. try: return socket.gethostbyname(socket.gethostname()) except Exception: - return "x.x.x.x" + return "" diff --git a/aikido_zen/helpers/get_machine_ip_test.py b/aikido_zen/helpers/get_machine_ip_test.py index a68eaffea..10ee682d5 100644 --- a/aikido_zen/helpers/get_machine_ip_test.py +++ b/aikido_zen/helpers/get_machine_ip_test.py @@ -3,6 +3,11 @@ from aikido_zen.helpers.get_machine_ip import get_ip +@pytest.fixture(autouse=True) +def reset_cached_ip_address(): + get_ip.cache_clear() + + def test_get_ip_success(monkeypatch): # Mock the socket.gethostname to return a specific hostname monkeypatch.setattr(socket, "gethostname", lambda: "mocked_hostname") @@ -22,4 +27,31 @@ def test_get_ip_failure(monkeypatch): lambda hostname: (_ for _ in ()).throw(Exception("Mocked exception")), ) - assert get_ip() == "x.x.x.x" + assert get_ip() == "" + + +def test_get_ip_is_resolved_once(monkeypatch): + calls = [] + monkeypatch.setattr(socket, "gethostname", lambda: "mocked_hostname") + monkeypatch.setattr( + socket, + "gethostbyname", + lambda hostname: calls.append(hostname) or "192.168.1.1", + ) + + assert get_ip() == get_ip() + assert len(calls) == 1 + + +def test_get_ip_resolves_a_failure_once(monkeypatch): + calls = [] + monkeypatch.setattr(socket, "gethostname", lambda: "mocked_hostname") + monkeypatch.setattr( + socket, + "gethostbyname", + lambda hostname: calls.append(hostname) + or (_ for _ in ()).throw(Exception("Mocked exception")), + ) + + assert get_ip() == get_ip() == "" + assert len(calls) == 1 diff --git a/aikido_zen/helpers/uuid7.py b/aikido_zen/helpers/uuid7.py new file mode 100644 index 000000000..1196f8ba6 --- /dev/null +++ b/aikido_zen/helpers/uuid7.py @@ -0,0 +1,18 @@ +import os +import time +import uuid + + +def uuid7(): + if hasattr(uuid, "uuid7"): + return str(uuid.uuid7()) + return uuid7_polyfill() + + +def uuid7_polyfill(): + value = bytearray(os.urandom(16)) + timestamp_ms = time.time_ns() // 1_000_000 + value[0:6] = timestamp_ms.to_bytes(6, "big") + value[6] = (value[6] & 0x0F) | 0x70 # version 7 + value[8] = (value[8] & 0x3F) | 0x80 # RFC 9562 variant + return str(uuid.UUID(bytes=bytes(value))) diff --git a/aikido_zen/helpers/uuid7_test.py b/aikido_zen/helpers/uuid7_test.py new file mode 100644 index 000000000..f6bf2ab54 --- /dev/null +++ b/aikido_zen/helpers/uuid7_test.py @@ -0,0 +1,30 @@ +import re +import time + +from aikido_zen.helpers.uuid7 import uuid7, uuid7_polyfill + +_UUID_V7 = re.compile( + r"[0-9a-f]{8}-[0-9a-f]{4}-7[0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}" +) + + +def _timestamp_ms(value): + return int(value.replace("-", "")[:12], 16) + + +def _assert_uuid7(generate): + before = time.time_ns() // 1_000_000 + value = generate() + after = time.time_ns() // 1_000_000 + + assert _UUID_V7.fullmatch(value) + assert before <= _timestamp_ms(value) <= after + assert generate() != value + + +def test_uuid7_is_version_7(): + _assert_uuid7(uuid7) + + +def test_uuid7_polyfill_is_version_7(): + _assert_uuid7(uuid7_polyfill) diff --git a/end2end/flask_postgres_test.py b/end2end/flask_postgres_test.py index edb75a810..a0ecf92e7 100644 --- a/end2end/flask_postgres_test.py +++ b/end2end/flask_postgres_test.py @@ -144,3 +144,25 @@ def test_track_sends_no_event_without_firewall(): time.sleep(5) # Wait, in case an event would be reported events = fetch_events_from_mock("http://localhost:5000") assert filter_on_event_type(events, "custom") == [] + + +def test_every_event_carries_the_agent_headers(): + clear_events_from_mock("http://localhost:5000") + res = requests.get(track_url_fw, headers={"User-Agent": "e2e-test"}) + assert res.status_code == 200 + + time.sleep(5) # Wait for the event to be reported + captured = fetch_events_from_mock("http://localhost:5000", include_headers=True) + + assert len(captured) > 0 + for entry in captured: + headers = entry["requestHeaders"] + assert headers["x-agent-platform"] == "python" + assert headers["x-agent-library"] == "firewall-python" + assert headers["x-agent-version"] == "1.0-REPLACE-VERSION" + assert headers["x-agent-hostname"] + assert headers["x-agent-ip-address"] + + session_ids = {entry["requestHeaders"]["x-agent-session-id"] for entry in captured} + assert len(session_ids) == 1 + assert session_ids.pop() diff --git a/end2end/server/check_events_from_mock.py b/end2end/server/check_events_from_mock.py index 7911c66ad..578eaec3e 100644 --- a/end2end/server/check_events_from_mock.py +++ b/end2end/server/check_events_from_mock.py @@ -3,9 +3,10 @@ import json import time -def fetch_events_from_mock(url): +def fetch_events_from_mock(url, include_headers=False): mock_events_url = f"{url}/mock/events" - res = requests.get(mock_events_url, timeout=5) + params = {"includeHeaders": "true"} if include_headers else None + res = requests.get(mock_events_url, params=params, timeout=5) json_events = json.loads(res.content.decode("utf-8")) return json_events diff --git a/end2end/server/mock_aikido_core.py b/end2end/server/mock_aikido_core.py index 9665e8295..bd3924741 100644 --- a/end2end/server/mock_aikido_core.py +++ b/end2end/server/mock_aikido_core.py @@ -44,7 +44,7 @@ "configUpdatedAt": {}, } -events = [] +events = [] # [{ 'event': ..., 'requestHeaders': ... }] @app.route('/config', methods=['GET']) @@ -76,7 +76,13 @@ def get_fw_lists(): def post_events(): print("Got event: ", request.get_json()) if request.get_json(): - events.append(request.get_json()) + # The WSGI environ loses the original casing, so store them lowercase. + events.append({ + "event": request.get_json(), + "requestHeaders": { + name.lower(): value for name, value in request.headers.items() + }, + }) return jsonify(responses["config"]) @@ -91,7 +97,9 @@ def mock_set_config(): @app.route('/mock/events', methods=['GET']) def mock_get_events(): - return jsonify(events) + if request.args.get('includeHeaders') == 'true': + return jsonify(events) + return jsonify([captured["event"] for captured in events]) @app.route('/timeout5/api/runtime/events', methods=['POST']) def mock_timeout_5_secs():