LOGAR/tests/test_pipeline.py

import os
import sys
import json
import time
import socket
import ssl
import struct
import sqlite3
import urllib.request
import urllib.error
import shutil
import warnings
from datetime import datetime, timezone, timedelta

# Ensure repository root and src/ directory are in sys.path
ROOT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
SRC_DIR = os.path.join(ROOT_DIR, "src")
sys.path.insert(0, ROOT_DIR)
sys.path.insert(0, SRC_DIR)

warnings.filterwarnings("ignore")

import Win_Client
import Linux_Client

# Test server endpoints
TCP_HOST = "127.0.0.1"
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")
    client_cfg_path = "client_config.json" if os.path.exists("client_config.json") else os.path.join(ROOT_DIR, "client_config.json")

    assert os.path.exists(server_cfg_path), f"{server_cfg_path} must exist"
    assert os.path.exists(client_cfg_path), f"{client_cfg_path} must exist"
    
    with open(client_cfg_path, "r", encoding="utf-8") as f:
        client_conf = json.load(f)
    
    with open(server_cfg_path, "r", encoding="utf-8") as f:
        server_conf = json.load(f)

    assert "server_name" not in client_conf, "client_config.json must NOT contain server_name"
    assert "name" not in client_conf, "client_config.json must NOT contain name"
    assert "site_name" not in client_conf, "client_config.json must NOT contain site_name"
    assert client_conf["server_fingerprint"] == server_conf["server_fingerprint"], "Fingerprints must match"
    print(f"[OK] Verified client_config.json contains no machine/server/site name.")
    print(f"[OK] Fingerprint verified: {client_conf['server_fingerprint']}")

    cert_dir = os.path.join(ROOT_DIR, "test_pipeline_certs")
    shutil.rmtree(cert_dir, ignore_errors=True)
    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")
    auth_token = server_conf.get("auth_token", "")

    print("\n=== [1.5] Testing Server License Activation via Web API (/api/license/import) ===")
    import server_enrollment
    
    # 1. Check current license status
    req_status_init = urllib.request.Request(f"http://{HERMES_HOST}:{HERMES_PORT}/api/license/status")
    with urllib.request.urlopen(req_status_init, timeout=5) as resp:
        init_lic_stat = json.loads(resp.read().decode("utf-8"))
    
    if init_lic_stat.get("status") == "UNLICENSED":
        unlicensed_enroll_payload = {
            "client_id": client_id,
            "hostname": client_id,
            "os": "linux",
            "enrollment_secret": enrollment_secret
        }
        req_unlicensed = urllib.request.Request(
            f"http://{HERMES_HOST}:{HERMES_PORT}/api/client/enroll",
            data=json.dumps(unlicensed_enroll_payload).encode("utf-8"),
            headers={"Content-Type": "application/json"}
        )
        try:
            with urllib.request.urlopen(req_unlicensed, timeout=5):
                assert False, "Expected HTTP 403 on unlicensed server enrollment"
        except urllib.error.HTTPError as e:
            assert e.code == 403, f"Expected HTTP 403 on unlicensed server, got {e.code}"
            print("[OK] Unlicensed server correctly rejected edge client enrollment.")

    # 2. Fetch CA CSR from Server
    req_csr = urllib.request.Request(
        f"http://{HERMES_HOST}:{HERMES_PORT}/api/license/csr",
        headers={"Authorization": f"Bearer {auth_token}"}
    )
    with urllib.request.urlopen(req_csr, timeout=5) as resp:
        assert resp.status == 200
        csr_data = json.loads(resp.read().decode("utf-8"))
        csr_pem = csr_data["csr"]
    print("[OK] Fetched subordinate CA CSR from server.")

    # 3. Generate test master CA and sign subordinate CA with 10 seats
    master_ca_dir = os.path.join(ROOT_DIR, "test_pipeline_master_ca")
    master_cert, master_key, master_pem, master_key_pem = server_enrollment.generate_ca_if_needed(cert_dir=master_ca_dir, common_name="LOGAR-Master-Root-CA")
    sub_pem = server_enrollment.sign_ca_csr_with_seats(
        csr_pem=csr_pem,
        master_ca_cert_pem=master_pem,
        master_ca_key_pem=master_key_pem,
        license_id="LIC-TEST-PIPELINE-10",
        seats=10,
        days_valid=365
    )

    # 4. Import signed license package to server
    import_payload = {"ca_cert": sub_pem, "master_root_ca": master_pem}
    req_import = urllib.request.Request(
        f"http://{HERMES_HOST}:{HERMES_PORT}/api/license/import",
        data=json.dumps(import_payload).encode("utf-8"),
        headers={"Content-Type": "application/json", "Authorization": f"Bearer {auth_token}"}
    )
    with urllib.request.urlopen(req_import, timeout=5) as resp:
        assert resp.status == 200
        import_res = json.loads(resp.read().decode("utf-8"))
        assert import_res["status"] == "success"
        assert import_res["seats"] == 10
    print("[OK] Successfully imported signed license (.lic). Server activated with 10 seats.")

    # 5. Verify license status
    req_status = urllib.request.Request(f"http://{HERMES_HOST}:{HERMES_PORT}/api/license/status")
    with urllib.request.urlopen(req_status, timeout=5) as resp:
        assert resp.status == 200
        lic_stat = json.loads(resp.read().decode("utf-8"))
        assert lic_stat["status"] == "ACTIVE"
        assert lic_stat["max_seats"] == 10
    print("[OK] Verified server license status is ACTIVE.")

    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
    auth_tok = server_conf.get("auth_token") or client_conf.get("auth_token")
    client_headers = {"Authorization": f"Bearer {auth_tok}"} if auth_tok else {}
    clients_req = urllib.request.Request(f"http://{HERMES_HOST}:{HERMES_PORT}/api/clients", headers=client_headers)
    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": client_id,
                "timestamp": datetime.now(timezone.utc).isoformat(),
                "logs": logs
            }
            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 = 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=== [3] Testing Temporal Persistence & Direct Anomaly Verification over mTLS ===")
    test_signature = f"TestServiceDegraded_{int(time.time())}"
    candidate_log = [{
        "server": client_id,
        "os_type": "linux",
        "signature": test_signature,
        "severity": "WARNING",
        "message": "Resource usage high warning"
    }]

    for run_num in range(1, 5):
        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 via mTLS. Promoted to verified: {promoted}")
        if run_num == 1:
            assert promoted == 1, f"Expected 1 promoted on run 1 for warning, got {promoted}"
        else:
            assert promoted == 0, f"Expected 0 newly promoted on run {run_num} (already verified), got {promoted}"

    # Inspect SQLite database directly
    conn = sqlite3.connect(server_conf.get("db_path", "logar_state.db"))
    cursor = conn.cursor()
    cursor.execute("SELECT run_count, status FROM active_issues WHERE signature = ?", (test_signature,))
    row = cursor.fetchone()
    conn.close()

    assert row is not None, "Issue not found in SQLite"
    run_count, status = row
    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] Direct verification confirmed: Warning immediately verified for Hermes with run count tracking over mTLS!")

    print("\n=== [4] Testing Immediate Pass for Errors over mTLS ===")
    error_signature = f"TestServiceCrashImmediate_{int(time.time())}"
    error_log = [{
        "server": client_id,
        "os_type": "linux",
        "signature": error_signature,
        "severity": "ERROR",
        "message": "Fatal process crash occurred"
    }]
    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, got {err_resp.get('promoted_verified')}"

    conn = sqlite3.connect(server_conf.get("db_path", "logar_state.db"))
    cursor = conn.cursor()
    cursor.execute("SELECT run_count, status FROM active_issues WHERE signature = ?", (error_signature,))
    err_row = cursor.fetchone()
    conn.close()
    assert err_row is not None, "Error issue not found in SQLite"
    err_run_count, err_status = err_row
    print(f"[DB Verification] Issue '{error_signature}' -> run_count: {err_run_count}, status: {err_status}")
    assert err_run_count == 1, f"Expected run_count == 1, got {err_run_count}"
    assert err_status == "VERIFIED", f"Expected status 'VERIFIED', got {err_status}"
    print("[OK] Immediate pass verified: Error promoted to VERIFIED anomaly immediately!")

    print("\n=== [5] Testing Hermes Reporting Endpoint (/api/hermes/report) ===")
    token = server_conf.get("auth_token") or client_conf.get("auth_token")
    headers = {"Authorization": f"Bearer {token}"} if token else {}
    req = urllib.request.Request(f"http://{HERMES_HOST}:{HERMES_PORT}/api/hermes/report", headers=headers)
    with urllib.request.urlopen(req, timeout=5) as response:
        assert response.status == 200, f"Expected 200, got {response.status}"
        hermes_data = json.loads(response.read().decode("utf-8"))

    print(f"[Hermes API] Returned {len(hermes_data)} verified anomalies:")
    found_warning = False
    found_error = False
    for issue in hermes_data:
        print(f"  - Fingerprint: {issue['fingerprint']} | Consecutive Runs: {issue['consecutive_runs']} | Status: {issue['status']}")
        if issue["signature"] == test_signature:
            found_warning = True
            assert issue["verified"] is True
            assert issue["consecutive_runs"] >= 4
        if issue["signature"] == error_signature:
            found_error = True
            assert issue["verified"] is True
            assert issue["consecutive_runs"] == 1

    assert found_warning, f"Warning issue {test_signature} should be in Hermes report"
    assert found_error, f"Error issue {error_signature} should be in Hermes report"
    print("[OK] Hermes reporting validated!")

    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=== [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.")

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