Send the common agent headers on all cloud requests

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,
     )

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:

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": (

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")
 
 

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"

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"]
+    )

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)

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,
         },
     )
 

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}

aikido_zen/background_process/requests/make_request.py

@@ -4,16 +4,18 @@
 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)
+    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)

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

aikido_zen/config.py

@@ -1,3 +1,4 @@
 """Contains package versions"""
 
 PKG_VERSION = "1.0-REPLACE-VERSION"
+LIBRARY_NAME = "firewall-python"