test(pipeline): update end-to-end integration test for mTLS and dynamic PKI enrollment in test_pipeline.py

tests/test_pipeline.py

@@ -3,9 +3,11 @@ import sys
 import json
 import time
 import socket
+import ssl
 import struct
 import sqlite3
 import urllib.request
+import urllib.error
 import warnings
 from datetime import datetime, timezone, timedelta
 
@@ -16,7 +18,9 @@ sys.path.insert(0, ROOT_DIR)
 sys.path.insert(0, SRC_DIR)
 
 warnings.filterwarnings("ignore")
-import pgpy
+
+import Win_Client
+import Linux_Client
 
 # Test server endpoints
 TCP_HOST = "127.0.0.1"
@@ -24,6 +28,7 @@ TCP_PORT = 9443
 HERMES_HOST = "127.0.0.1"
 HERMES_PORT = 8443
 
+
 def run_tests():
     print("=== [1] Verifying server_config.json & client_config.json ===")
     server_cfg_path = "server_config.json" if os.path.exists("server_config.json") else os.path.join(ROOT_DIR, "server_config.json")
@@ -45,56 +50,117 @@ def run_tests():
     print(f"[OK] Verified client_config.json contains no machine/server/site name.")
     print(f"[OK] Fingerprint verified: {client_conf['server_fingerprint']}")
 
-    # Load public key
-    pub_key, _ = pgpy.PGPKey.from_blob(client_conf["server_public_key"])
+    cert_dir = os.path.join(ROOT_DIR, "test_pipeline_certs")
+    os.makedirs(cert_dir, exist_ok=True)
+    client_id = "test-edge-node.corp.internal"
+    enrollment_secret = server_conf.get("enrollment_secret") or client_conf.get("enrollment_secret")
 
-    def send_socket_batch(logs, auth_token=client_conf["auth_token"]):
+    print("\n=== [2] Testing Client Dynamic PKI Enrollment API (/api/client/enroll) ===")
+    enroll_url = f"http://{HERMES_HOST}:{HERMES_PORT}/api/client/enroll"
+
+    # 2a. Test rejection on invalid enrollment secret
+    bad_enroll_payload = {
+        "client_id": client_id,
+        "hostname": client_id,
+        "os": "linux",
+        "enrollment_secret": "invalid-secret-xyz"
+    }
+    req_bad = urllib.request.Request(
+        enroll_url,
+        data=json.dumps(bad_enroll_payload).encode("utf-8"),
+        headers={"Content-Type": "application/json"}
+    )
+    try:
+        with urllib.request.urlopen(req_bad, timeout=5):
+            assert False, "Expected HTTP 403 on invalid secret"
+    except urllib.error.HTTPError as e:
+        assert e.code == 403, f"Expected HTTP 403, got {e.code}"
+        print("[OK] Invalid enrollment secret rejected with HTTP 403.")
+
+    # 2b. Test valid client enrollment
+    valid_enroll_payload = {
+        "client_id": client_id,
+        "hostname": client_id,
+        "os": "linux",
+        "enrollment_secret": enrollment_secret
+    }
+    req_valid = urllib.request.Request(
+        enroll_url,
+        data=json.dumps(valid_enroll_payload).encode("utf-8"),
+        headers={"Content-Type": "application/json"}
+    )
+    with urllib.request.urlopen(req_valid, timeout=5) as resp:
+        assert resp.status == 200, f"Expected 200, got {resp.status}"
+        enroll_data = json.loads(resp.read().decode("utf-8"))
+        assert "ca_cert" in enroll_data
+        assert "client_cert" in enroll_data
+        assert "client_key" in enroll_data
+
+    ca_path = os.path.join(cert_dir, "ca.crt")
+    cert_path = os.path.join(cert_dir, "client.crt")
+    key_path = os.path.join(cert_dir, "client.key")
+
+    with open(ca_path, "w", encoding="utf-8") as f:
+        f.write(enroll_data["ca_cert"])
+    with open(cert_path, "w", encoding="utf-8") as f:
+        f.write(enroll_data["client_cert"])
+    with open(key_path, "w", encoding="utf-8") as f:
+        f.write(enroll_data["client_key"])
+    print(f"[OK] Client enrolled successfully. Certificates stored in {cert_dir}")
+
+    # 2c. Verify client shows in /api/clients
+    clients_req = urllib.request.Request(f"http://{HERMES_HOST}:{HERMES_PORT}/api/clients")
+    with urllib.request.urlopen(clients_req, timeout=5) as resp:
+        clients_data = json.loads(resp.read().decode("utf-8"))
+        assert clients_data["active_seats"] >= 1
+        found_c = any(c["client_id"] == client_id for c in clients_data["clients"])
+        assert found_c, f"Client {client_id} should be listed in /api/clients"
+    print(f"[OK] Verified client in /api/clients: Active Seats: {clients_data['active_seats']}/{clients_data['max_seats']}")
+
+    def send_mtls_batch(logs):
+        ctx = ssl.create_default_context(ssl.Purpose.SERVER_AUTH, cafile=ca_path)
+        ctx.load_cert_chain(certfile=cert_path, keyfile=key_path)
+        ctx.minimum_version = ssl.TLSVersion.TLSv1_3
+        ctx.check_hostname = False
+
+        raw_sock = socket.create_connection((TCP_HOST, TCP_PORT), timeout=10)
+        with ctx.wrap_socket(raw_sock, server_hostname=TCP_HOST) as s:
             payload = {
-            "server": "test-edge-node.corp.internal",
+                "server": client_id,
                 "timestamp": datetime.now(timezone.utc).isoformat(),
                 "logs": logs
             }
-        pgp_msg = pgpy.PGPMessage.new(json.dumps(payload))
-        enc = pub_key.encrypt(pgp_msg)
-        
-        envelope = {
-            "auth_token": auth_token,
-            "timestamp": datetime.now(timezone.utc).isoformat(),
-            "encrypted_payload": str(enc)
-        }
-        envelope_bytes = json.dumps(envelope).encode("utf-8")
-        
-        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
-            s.settimeout(5.0)
-            s.connect((TCP_HOST, TCP_PORT))
-            frame = struct.pack(">I", len(envelope_bytes)) + envelope_bytes
+            payload_bytes = json.dumps(payload).encode("utf-8")
+            frame = struct.pack(">I", len(payload_bytes)) + payload_bytes
             s.sendall(frame)
 
             resp_len_bytes = s.recv(4)
+            if not resp_len_bytes:
+                raise ConnectionError("Server closed connection without response.")
             resp_len = struct.unpack(">I", resp_len_bytes)[0]
-            resp_bytes = s.recv(resp_len)
+            resp_bytes = bytearray()
+            while len(resp_bytes) < resp_len:
+                chunk = s.recv(min(4096, resp_len - len(resp_bytes)))
+                if not chunk:
+                    break
+                resp_bytes.extend(chunk)
             return json.loads(resp_bytes.decode("utf-8"))
 
-    print("\n=== [2] Testing Socket Authentication Failure ===")
-    bad_resp = send_socket_batch([], auth_token="invalid-token-12345")
-    assert bad_resp.get("status") == "error", f"Expected error, got: {bad_resp}"
-    print(f"[OK] Bad auth rejected correctly: {bad_resp['message']}")
-
+    print("\n=== [3] Testing Temporal Persistence & 4-Run Rule for Warnings over mTLS ===")
     test_signature = "TestServiceDegraded"
     candidate_log = [{
-        "server": "test-edge-node",
+        "server": client_id,
         "os_type": "linux",
         "signature": test_signature,
         "severity": "WARNING",
         "message": "Resource usage high warning"
     }]
 
-    print("\n=== [3] Testing Temporal Persistence & 4-Run Rule for Warnings ===")
     for run_num in range(1, 5):
-        resp = send_socket_batch(candidate_log)
+        resp = send_mtls_batch(candidate_log)
         assert resp.get("status") == "success", f"Run {run_num} failed: {resp}"
         promoted = resp.get("promoted_verified", 0)
-        print(f"[Run {run_num}/4] Ingested successfully. Promoted to verified: {promoted}")
+        print(f"[Run {run_num}/4] Ingested successfully via mTLS. Promoted to verified: {promoted}")
         if run_num < 4:
             assert promoted == 0, f"Expected 0 promoted on run {run_num} for warning, got {promoted}"
         else:
@@ -112,21 +178,21 @@ def run_tests():
     print(f"[DB Verification] Issue '{test_signature}' -> run_count: {run_count}, status: {status}")
     assert run_count >= 4, f"Expected run_count >= 4, got {run_count}"
     assert status == "VERIFIED", f"Expected status 'VERIFIED', got {status}"
-    print("[OK] 4-Run Rule verified: Warning promoted to VERIFIED anomaly on 4th run!")
+    print("[OK] 4-Run Rule verified: Warning promoted to VERIFIED anomaly on 4th run over mTLS!")
 
-    print("\n=== [3b] Testing Immediate Pass for Errors ===")
+    print("\n=== [4] Testing Immediate Pass for Errors over mTLS ===")
     error_signature = "TestServiceCrashImmediate"
     error_log = [{
-        "server": "test-edge-node",
+        "server": client_id,
         "os_type": "linux",
         "signature": error_signature,
         "severity": "ERROR",
         "message": "Fatal process crash occurred"
     }]
-    err_resp = send_socket_batch(error_log)
+    err_resp = send_mtls_batch(error_log)
     assert err_resp.get("status") == "success", f"Error run failed: {err_resp}"
     print(f"[Run 1/1] Error ingested successfully. Promoted to verified: {err_resp.get('promoted_verified')}")
-    assert err_resp.get("promoted_verified") == 1, f"Expected error to be promoted to verified immediately on run 1, got {err_resp.get('promoted_verified')}"
+    assert err_resp.get("promoted_verified") == 1, f"Expected error to be promoted to verified immediately, got {err_resp.get('promoted_verified')}"
 
     conn = sqlite3.connect(server_conf.get("db_path", "logar_state.db"))
     cursor = conn.cursor()
@@ -140,7 +206,7 @@ def run_tests():
     assert err_status == "VERIFIED", f"Expected status 'VERIFIED', got {err_status}"
     print("[OK] Immediate pass verified: Error promoted to VERIFIED anomaly immediately!")
 
-    print("\n=== [4] Testing Hermes Reporting Endpoint (/api/hermes/report) ===")
+    print("\n=== [5] Testing Hermes Reporting Endpoint (/api/hermes/report) ===")
     req = urllib.request.Request(f"http://{HERMES_HOST}:{HERMES_PORT}/api/hermes/report")
     with urllib.request.urlopen(req, timeout=5) as response:
         assert response.status == 200, f"Expected 200, got {response.status}"
@@ -164,19 +230,24 @@ def run_tests():
     assert found_error, f"Error issue {error_signature} should be in Hermes report"
     print("[OK] Hermes reporting validated!")
 
-    print("\n=== [5] Testing Windows Client Script Integration ===")
+    print("\n=== [6] Testing Windows Client Script Integration ===")
     from Win_Client import get_recent_windows_logs
     win_logs = get_recent_windows_logs(hours=24)
     print(f"[Win_Client] Successfully queried Windows logs: {len(win_logs)} candidate entries.")
 
-    print("\n=== [6] Testing Linux Client Script Integration ===")
+    print("\n=== [7] Testing Linux Client Script Integration ===")
     from Linux_Client import get_recent_linux_logs
     linux_logs = get_recent_linux_logs(hours=24)
     print(f"[Linux_Client] Successfully queried Linux logs: {len(linux_logs)} candidate entries.")
 
-    print("\n==========================================")
-    print("  ALL VERIFICATION TESTS PASSED SUCCESSFULLY! ")
-    print("==========================================")
+    import shutil
+    if os.path.exists(cert_dir):
+        shutil.rmtree(cert_dir, ignore_errors=True)
+
+    print("\n=======================================================")
+    print("  ALL VERIFICATION TESTS (mTLS + PKI + PIPELINE) PASSED! ")
+    print("=======================================================")
+
 
 if __name__ == "__main__":
     run_tests()