From 3c30bb5ae0cab939d6871cd2c87b1198d8eb5b6f Mon Sep 17 00:00:00 2001 From: Daniel Iulian Iacob Date: Thu, 8 Oct 2026 19:53:31 +0300 Subject: [PATCH 1/4] Send the common agent headers on all cloud requests --- aikido_zen/background_process/__init__.py | 3 +- .../aikido_background_process.py | 5 +- .../get_manager_info.py | 4 +- .../get_common_agent_headers.py | 18 ++++++ .../get_common_agent_headers_test.py | 62 +++++++++++++++++++ aikido_zen/background_process/init_test.py | 40 ++++++++++++ .../realtime/sse_client/__init__.py | 7 ++- .../realtime/sse_client/real_server_test.py | 10 +++ .../requests/make_request.py | 15 +++-- .../requests/make_request_test.py | 47 ++++++++++++++ aikido_zen/helpers/get_agent_session_id.py | 14 +++++ .../helpers/get_agent_session_id_test.py | 27 ++++++++ aikido_zen/helpers/get_hostname.py | 13 ++++ aikido_zen/helpers/get_hostname_test.py | 18 ++++++ aikido_zen/helpers/get_machine_ip.py | 5 +- aikido_zen/helpers/get_machine_ip_test.py | 32 ++++++++++ aikido_zen/helpers/uuid7.py | 18 ++++++ aikido_zen/helpers/uuid7_test.py | 30 +++++++++ end2end/flask_postgres_test.py | 22 ++++++- end2end/server/check_events_from_mock.py | 4 ++ end2end/server/mock_aikido_core.py | 11 ++++ 21 files changed, 390 insertions(+), 15 deletions(-) create mode 100644 aikido_zen/background_process/get_common_agent_headers.py create mode 100644 aikido_zen/background_process/get_common_agent_headers_test.py create mode 100644 aikido_zen/background_process/requests/make_request_test.py create mode 100644 aikido_zen/helpers/get_agent_session_id.py create mode 100644 aikido_zen/helpers/get_agent_session_id_test.py create mode 100644 aikido_zen/helpers/get_hostname.py create mode 100644 aikido_zen/helpers/get_hostname_test.py create mode 100644 aikido_zen/helpers/uuid7.py create mode 100644 aikido_zen/helpers/uuid7_test.py 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..43da069fe 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,7 +12,7 @@ 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", "ipAddress": h.get_ip(), 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..bdbd81b90 --- /dev/null +++ b/aikido_zen/background_process/get_common_agent_headers.py @@ -0,0 +1,18 @@ +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 + +_UNKNOWN = "unknown" +_UNRESOLVED_IP = "x.x.x.x" # What get_ip returns when it cannot resolve the address. + + +def get_common_agent_headers(): + ip_address = get_ip() + return { + "X-Agent-Platform": "python", + "X-Agent-Version": config.PKG_VERSION, + "X-Agent-Hostname": get_hostname() or _UNKNOWN, + "X-Agent-IP-Address": _UNKNOWN if ip_address == _UNRESOLVED_IP else ip_address, + "X-Agent-Session-Id": get_agent_session_id(), + } 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..968fe404e --- /dev/null +++ b/aikido_zen/background_process/get_common_agent_headers_test.py @@ -0,0 +1,62 @@ +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_ip_address(): + get_ip.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-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_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: "x.x.x.x") + + assert get_common_agent_headers()["X-Agent-IP-Address"] == "unknown" 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..60e274c24 100644 --- a/aikido_zen/background_process/requests/make_request.py +++ b/aikido_zen/background_process/requests/make_request.py @@ -4,16 +4,19 @@ 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 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) + req = urllib.request.Request( + url, + data=data, + method=method, + headers={**get_common_agent_headers(), **(headers or {})}, + ) 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..4436cd71d --- /dev/null +++ b/aikido_zen/background_process/requests/make_request_test.py @@ -0,0 +1,47 @@ +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 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..97142e185 --- /dev/null +++ b/aikido_zen/helpers/get_hostname.py @@ -0,0 +1,13 @@ +""" +Helper function file, see function docstring +""" + +import socket + + +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..0e6ff7b98 --- /dev/null +++ b/aikido_zen/helpers/get_hostname_test.py @@ -0,0 +1,18 @@ +import socket +from aikido_zen.helpers.get_hostname import get_hostname + + +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() == "" diff --git a/aikido_zen/helpers/get_machine_ip.py b/aikido_zen/helpers/get_machine_ip.py index d36eef729..0f678f5b2 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""" + # Cached: this queries DNS and runs on every report. try: return socket.gethostbyname(socket.gethostname()) - except Exception: + except Exception: # pylint: disable=broad-exception-caught return "x.x.x.x" diff --git a/aikido_zen/helpers/get_machine_ip_test.py b/aikido_zen/helpers/get_machine_ip_test.py index a68eaffea..f7a4aea37 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") @@ -23,3 +28,30 @@ def test_get_ip_failure(monkeypatch): ) assert get_ip() == "x.x.x.x" + + +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() == "x.x.x.x" + 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..475a6a1e0 100644 --- a/end2end/flask_postgres_test.py +++ b/end2end/flask_postgres_test.py @@ -2,7 +2,7 @@ import pytest import json import requests -from .server.check_events_from_mock import fetch_events_from_mock, validate_started_event, filter_on_event_type, \ +from .server.check_events_from_mock import fetch_events_from_mock, fetch_agent_headers_from_mock, validate_started_event, filter_on_event_type, \ clear_events_from_mock # e2e tests for flask_postgres sample app @@ -144,3 +144,23 @@ 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_events_carry_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 + agent_headers = fetch_agent_headers_from_mock("http://localhost:5000") + + assert len(agent_headers) > 0 + for headers in agent_headers: + assert headers["x-agent-platform"] == "python" + assert headers["x-agent-version"] == "1.0-REPLACE-VERSION" + assert headers["x-agent-hostname"] + assert headers["x-agent-ip-address"] + + session_ids = {headers["x-agent-session-id"] for headers in agent_headers} + 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..9381aee6d 100644 --- a/end2end/server/check_events_from_mock.py +++ b/end2end/server/check_events_from_mock.py @@ -9,6 +9,10 @@ def fetch_events_from_mock(url): json_events = json.loads(res.content.decode("utf-8")) return json_events +def fetch_agent_headers_from_mock(url): + res = requests.get(f"{url}/mock/agent_headers", timeout=5) + return json.loads(res.content.decode("utf-8")) + def clear_events_from_mock(url): mock_events_url = f"{url}/mock/reset" res = requests.get(mock_events_url, timeout=5) diff --git a/end2end/server/mock_aikido_core.py b/end2end/server/mock_aikido_core.py index 9665e8295..bfcb1eb22 100644 --- a/end2end/server/mock_aikido_core.py +++ b/end2end/server/mock_aikido_core.py @@ -45,6 +45,7 @@ } events = [] +agent_headers = [] @app.route('/config', methods=['GET']) @@ -77,6 +78,11 @@ 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. + agent_headers.append({ + name.lower(): value for name, value in request.headers.items() + if name.lower().startswith('x-agent-') + }) return jsonify(responses["config"]) @@ -93,6 +99,10 @@ def mock_set_config(): def mock_get_events(): return jsonify(events) +@app.route('/mock/agent_headers', methods=['GET']) +def mock_get_agent_headers(): + return jsonify(agent_headers) + @app.route('/timeout5/api/runtime/events', methods=['POST']) def mock_timeout_5_secs(): time.sleep(5) @@ -103,6 +113,7 @@ def mock_timeout_5_secs(): @app.route('/mock/reset', methods=['GET']) def mock_reset_events(): events.clear() + agent_headers.clear() return jsonify({"msg": "OK"}) From 732c2a265e04a431543da4f8907c7ab3ce95abec Mon Sep 17 00:00:00 2001 From: Daniel Iulian Iacob Date: Fri, 9 Oct 2026 10:56:56 +0300 Subject: [PATCH 2/4] Keep the original request construction in make_request --- .../background_process/requests/make_request.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/aikido_zen/background_process/requests/make_request.py b/aikido_zen/background_process/requests/make_request.py index 60e274c24..541375fdd 100644 --- a/aikido_zen/background_process/requests/make_request.py +++ b/aikido_zen/background_process/requests/make_request.py @@ -11,12 +11,11 @@ def make_request(method, url, timeout, data=None, headers=None): - req = urllib.request.Request( - url, - data=data, - method=method, - headers={**get_common_agent_headers(), **(headers or {})}, - ) + req = urllib.request.Request(url, data=data, method=method) + + # Add headers + 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) From e8bca90923963195e2afbf04eed7e170d0d5e8a9 Mon Sep 17 00:00:00 2001 From: Daniel Iulian Iacob Date: Fri, 9 Oct 2026 12:43:17 +0300 Subject: [PATCH 3/4] Harden the agent headers and add X-Agent-Library --- .../get_manager_info.py | 2 +- .../get_manager_info_test.py | 2 ++ .../get_common_agent_headers.py | 16 +++++++++---- .../get_common_agent_headers_test.py | 19 +++++++++++++-- .../requests/make_request_test.py | 24 +++++++++++++++++++ aikido_zen/config.py | 1 + aikido_zen/helpers/get_hostname.py | 2 ++ aikido_zen/helpers/get_hostname_test.py | 16 +++++++++++++ aikido_zen/helpers/get_machine_ip.py | 6 ++--- aikido_zen/helpers/get_machine_ip_test.py | 4 ++-- 10 files changed, 80 insertions(+), 12 deletions(-) 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 43da069fe..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 @@ -14,7 +14,7 @@ def get_manager_info(connection_manager): "dryMode": not connection_manager.block, "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 index bdbd81b90..f165dda3e 100644 --- a/aikido_zen/background_process/get_common_agent_headers.py +++ b/aikido_zen/background_process/get_common_agent_headers.py @@ -4,15 +4,23 @@ from aikido_zen.helpers.get_machine_ip import get_ip _UNKNOWN = "unknown" -_UNRESOLVED_IP = "x.x.x.x" # What get_ip returns when it cannot resolve the address. def get_common_agent_headers(): - ip_address = get_ip() return { "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": _UNKNOWN if ip_address == _UNRESOLVED_IP else ip_address, + "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 index 968fe404e..ace227316 100644 --- a/aikido_zen/background_process/get_common_agent_headers_test.py +++ b/aikido_zen/background_process/get_common_agent_headers_test.py @@ -17,8 +17,9 @@ @pytest.fixture(autouse=True) -def reset_cached_ip_address(): +def reset_cached_instance_values(): get_ip.cache_clear() + get_hostname.cache_clear() def test_common_agent_headers(monkeypatch): @@ -26,6 +27,7 @@ def test_common_agent_headers(monkeypatch): 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", @@ -50,6 +52,12 @@ def test_hostname_and_ip_match_the_event_payload(monkeypatch): 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: "") @@ -57,6 +65,13 @@ def test_missing_hostname_is_unknown(monkeypatch): def test_unresolved_ip_address_is_unknown(monkeypatch): - monkeypatch.setattr(headers_module, "get_ip", lambda: "x.x.x.x") + 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/requests/make_request_test.py b/aikido_zen/background_process/requests/make_request_test.py index 4436cd71d..5d4f301b0 100644 --- a/aikido_zen/background_process/requests/make_request_test.py +++ b/aikido_zen/background_process/requests/make_request_test.py @@ -1,3 +1,5 @@ +import http.server +import threading from unittest.mock import patch from aikido_zen.background_process.get_common_agent_headers import ( @@ -45,3 +47,25 @@ def fake_urlopen(request, timeout): 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_hostname.py b/aikido_zen/helpers/get_hostname.py index 97142e185..feab27bcb 100644 --- a/aikido_zen/helpers/get_hostname.py +++ b/aikido_zen/helpers/get_hostname.py @@ -3,8 +3,10 @@ """ 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: diff --git a/aikido_zen/helpers/get_hostname_test.py b/aikido_zen/helpers/get_hostname_test.py index 0e6ff7b98..5dc49d525 100644 --- a/aikido_zen/helpers/get_hostname_test.py +++ b/aikido_zen/helpers/get_hostname_test.py @@ -1,7 +1,13 @@ +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") @@ -16,3 +22,13 @@ def test_get_hostname_failure(monkeypatch): ) 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 0f678f5b2..a8b080031 100644 --- a/aikido_zen/helpers/get_machine_ip.py +++ b/aikido_zen/helpers/get_machine_ip.py @@ -8,9 +8,9 @@ @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: # pylint: disable=broad-exception-caught - return "x.x.x.x" + except Exception: + return "" diff --git a/aikido_zen/helpers/get_machine_ip_test.py b/aikido_zen/helpers/get_machine_ip_test.py index f7a4aea37..10ee682d5 100644 --- a/aikido_zen/helpers/get_machine_ip_test.py +++ b/aikido_zen/helpers/get_machine_ip_test.py @@ -27,7 +27,7 @@ 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): @@ -53,5 +53,5 @@ def test_get_ip_resolves_a_failure_once(monkeypatch): or (_ for _ in ()).throw(Exception("Mocked exception")), ) - assert get_ip() == get_ip() == "x.x.x.x" + assert get_ip() == get_ip() == "" assert len(calls) == 1 From 989fd50d14f9056e6977d26666ca6a3937309895 Mon Sep 17 00:00:00 2001 From: Daniel Iulian Iacob Date: Fri, 9 Oct 2026 14:10:50 +0300 Subject: [PATCH 4/4] Record the agent headers alongside each event in the mock core --- .../get_common_agent_headers.py | 6 ++---- end2end/flask_postgres_test.py | 14 +++++++------ end2end/server/check_events_from_mock.py | 9 +++----- end2end/server/mock_aikido_core.py | 21 ++++++++----------- 4 files changed, 22 insertions(+), 28 deletions(-) diff --git a/aikido_zen/background_process/get_common_agent_headers.py b/aikido_zen/background_process/get_common_agent_headers.py index f165dda3e..74c56f0fa 100644 --- a/aikido_zen/background_process/get_common_agent_headers.py +++ b/aikido_zen/background_process/get_common_agent_headers.py @@ -3,8 +3,6 @@ from aikido_zen.helpers.get_hostname import get_hostname from aikido_zen.helpers.get_machine_ip import get_ip -_UNKNOWN = "unknown" - def get_common_agent_headers(): return { @@ -22,5 +20,5 @@ def _header_value(value): try: value.encode("latin-1") except UnicodeEncodeError: - return _UNKNOWN - return value or _UNKNOWN + return "unknown" + return value or "unknown" diff --git a/end2end/flask_postgres_test.py b/end2end/flask_postgres_test.py index 475a6a1e0..a0ecf92e7 100644 --- a/end2end/flask_postgres_test.py +++ b/end2end/flask_postgres_test.py @@ -2,7 +2,7 @@ import pytest import json import requests -from .server.check_events_from_mock import fetch_events_from_mock, fetch_agent_headers_from_mock, validate_started_event, filter_on_event_type, \ +from .server.check_events_from_mock import fetch_events_from_mock, validate_started_event, filter_on_event_type, \ clear_events_from_mock # e2e tests for flask_postgres sample app @@ -146,21 +146,23 @@ def test_track_sends_no_event_without_firewall(): assert filter_on_event_type(events, "custom") == [] -def test_events_carry_the_agent_headers(): +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 - agent_headers = fetch_agent_headers_from_mock("http://localhost:5000") + captured = fetch_events_from_mock("http://localhost:5000", include_headers=True) - assert len(agent_headers) > 0 - for headers in agent_headers: + 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 = {headers["x-agent-session-id"] for headers in agent_headers} + 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 9381aee6d..578eaec3e 100644 --- a/end2end/server/check_events_from_mock.py +++ b/end2end/server/check_events_from_mock.py @@ -3,16 +3,13 @@ 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 -def fetch_agent_headers_from_mock(url): - res = requests.get(f"{url}/mock/agent_headers", timeout=5) - return json.loads(res.content.decode("utf-8")) - def clear_events_from_mock(url): mock_events_url = f"{url}/mock/reset" res = requests.get(mock_events_url, timeout=5) diff --git a/end2end/server/mock_aikido_core.py b/end2end/server/mock_aikido_core.py index bfcb1eb22..bd3924741 100644 --- a/end2end/server/mock_aikido_core.py +++ b/end2end/server/mock_aikido_core.py @@ -44,8 +44,7 @@ "configUpdatedAt": {}, } -events = [] -agent_headers = [] +events = [] # [{ 'event': ..., 'requestHeaders': ... }] @app.route('/config', methods=['GET']) @@ -77,11 +76,12 @@ 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. - agent_headers.append({ - name.lower(): value for name, value in request.headers.items() - if name.lower().startswith('x-agent-') + events.append({ + "event": request.get_json(), + "requestHeaders": { + name.lower(): value for name, value in request.headers.items() + }, }) return jsonify(responses["config"]) @@ -97,11 +97,9 @@ def mock_set_config(): @app.route('/mock/events', methods=['GET']) def mock_get_events(): - return jsonify(events) - -@app.route('/mock/agent_headers', methods=['GET']) -def mock_get_agent_headers(): - return jsonify(agent_headers) + 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(): @@ -113,7 +111,6 @@ def mock_timeout_5_secs(): @app.route('/mock/reset', methods=['GET']) def mock_reset_events(): events.clear() - agent_headers.clear() return jsonify({"msg": "OK"})