import os import sys import json import socket import struct import sqlite3 import unittest import urllib.request import warnings from datetime import datetime, timezone, timedelta warnings.filterwarnings("ignore") # Ensure parent directory and src directory are in path to import Server sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "src"))) import Server import pgpy class TestServerComponent(unittest.TestCase): def setUp(self): self.test_db = "test_server_state.db" self.test_config = "test_server_config.json" if os.path.exists(self.test_db): os.remove(self.test_db) if os.path.exists(self.test_config): os.remove(self.test_config) def tearDown(self): if os.path.exists(self.test_db): try: os.remove(self.test_db) except Exception: pass if os.path.exists(self.test_config): try: os.remove(self.test_config) except Exception: pass def test_first_run_config_and_keypair_generation(self): config = Server.load_or_init_config(self.test_config) self.assertTrue(os.path.exists(self.test_config)) self.assertIn("server_fingerprint", config) self.assertIn("public_key", config) self.assertIn("private_key", config) self.assertIn("auth_token", config) self.assertNotIn("site_name", config) # Verify keypair priv_key, _ = pgpy.PGPKey.from_blob(config["private_key"]) pub_key, _ = pgpy.PGPKey.from_blob(config["public_key"]) self.assertEqual(str(pub_key.fingerprint), config["server_fingerprint"]) def test_create_client_config(self): Server.load_or_init_config(self.test_config) client_out = "test_client_out.json" try: client_conf = Server.create_client_config( server_host="10.0.0.1", server_port=9443, output_path=client_out, config_path=self.test_config ) self.assertTrue(os.path.exists(client_out)) self.assertEqual(client_conf["server_host"], "10.0.0.1") self.assertEqual(client_conf["server_port"], 9443) # Verify no machine name or site_name is included self.assertNotIn("server_name", client_conf) self.assertNotIn("name", client_conf) self.assertNotIn("site_name", client_conf) finally: if os.path.exists(client_out): os.remove(client_out) def test_4_run_rule_and_12h_window(self): Server.init_db(self.test_db) log_entry = { "server": "app-worker-01.corp.local", "signature": "PostgresConnWarning", "severity": "WARNING", "message": "Connection to database pool near capacity: 85%", "os_type": "linux" } payload = { "server": "app-worker-01.corp.local", "logs": [log_entry] } # Runs 1 to 3: WARNING should remain TRANSIENT for run_idx in range(1, 4): res = Server.process_ingested_logs(payload, self.test_db, window_hours=12, min_runs=4) self.assertEqual(res["status"], "success") self.assertEqual(res["promoted_verified"], 0) conn = sqlite3.connect(self.test_db) c = conn.cursor() c.execute("SELECT run_count, status FROM active_issues WHERE signature = ?", ("PostgresConnWarning",)) row = c.fetchone() conn.close() self.assertEqual(row[0], 3) self.assertEqual(row[1], "TRANSIENT") # Run 4: promotes WARNING to VERIFIED! res4 = Server.process_ingested_logs(payload, self.test_db, window_hours=12, min_runs=4) self.assertEqual(res4["promoted_verified"], 1) conn = sqlite3.connect(self.test_db) c = conn.cursor() c.execute("SELECT run_count, status FROM active_issues WHERE signature = ?", ("PostgresConnWarning",)) row = c.fetchone() conn.close() self.assertEqual(row[0], 4) self.assertEqual(row[1], "VERIFIED") def test_error_immediate_pass(self): Server.init_db(self.test_db) log_entry = { "server": "app-worker-01.corp.local", "signature": "KernelPanicCritical", "severity": "ERROR", "message": "Kernel panic - not syncing: Fatal hardware error", "os_type": "linux" } payload = { "server": "app-worker-01.corp.local", "logs": [log_entry] } # Run 1: ERROR must immediately promote to VERIFIED res = Server.process_ingested_logs(payload, self.test_db, window_hours=12, min_runs=4) self.assertEqual(res["status"], "success") self.assertEqual(res["promoted_verified"], 1) conn = sqlite3.connect(self.test_db) c = conn.cursor() c.execute("SELECT run_count, status, severity FROM active_issues WHERE signature = ?", ("KernelPanicCritical",)) row = c.fetchone() conn.close() self.assertIsNotNone(row) self.assertEqual(row[0], 1) self.assertEqual(row[1], "VERIFIED") self.assertEqual(row[2], "ERROR") def test_server_severity_filtering(self): Server.init_db(self.test_db) payload = { "server": "app-worker-01.corp.local", "logs": [ {"server": "app-worker-01", "signature": "SigInfo", "severity": "INFO", "message": "Info msg", "os_type": "linux"}, {"server": "app-worker-01", "signature": "SigWarn", "severity": "WARNING", "message": "Warn msg", "os_type": "linux"}, {"server": "app-worker-01", "signature": "SigErr", "severity": "ERROR", "message": "Err msg", "os_type": "linux"}, {"server": "app-worker-01", "signature": "SigDebug", "severity": "DEBUG", "message": "Debug msg", "os_type": "linux"}, {"server": "app-worker-01", "signature": "SigTrace", "severity": "TRACE", "message": "Trace msg", "os_type": "linux"} ] } res = Server.process_ingested_logs(payload, self.test_db, window_hours=12, min_runs=4) self.assertEqual(res["status"], "success") conn = sqlite3.connect(self.test_db) c = conn.cursor() c.execute("SELECT signature, status FROM active_issues ORDER BY signature") rows = dict(c.fetchall()) conn.close() self.assertIn("SigInfo", rows) self.assertIn("SigWarn", rows) self.assertIn("SigErr", rows) self.assertNotIn("SigDebug", rows) self.assertNotIn("SigTrace", rows) # SigErr is immediately VERIFIED; SigWarn and SigInfo are TRANSIENT on run 1 self.assertEqual(rows["SigErr"], "VERIFIED") self.assertEqual(rows["SigWarn"], "TRANSIENT") self.assertEqual(rows["SigInfo"], "TRANSIENT") def test_license_schema_and_pki_generation(self): secret = "test-secret-12345" Server.init_db(self.test_db, enrollment_secret=secret, max_seats=5) conn = sqlite3.connect(self.test_db) c = conn.cursor() c.execute("SELECT max_seats, enrollment_secret FROM license_config WHERE id = 1") row = c.fetchone() conn.close() self.assertIsNotNone(row) self.assertEqual(row[0], 5) self.assertEqual(row[1], secret) # Test Dynamic PKI test_cert_dir = "test_certs_pki" try: ca_cert, ca_key, ca_pem, ca_key_pem = Server.enrollment.generate_ca_if_needed(cert_dir=test_cert_dir) self.assertIn("BEGIN CERTIFICATE", ca_pem) self.assertIn("BEGIN RSA PRIVATE KEY", ca_key_pem) srv_cert, srv_key, srv_pem, srv_key_pem = Server.enrollment.generate_server_cert_if_needed( ca_cert, ca_key, hostnames=["127.0.0.1", "localhost"], cert_dir=test_cert_dir ) self.assertIn("BEGIN CERTIFICATE", srv_pem) client_cert_pem, client_key_pem = Server.enrollment.issue_client_cert("node-test-1", ca_cert, ca_key) self.assertIn("BEGIN CERTIFICATE", client_cert_pem) self.assertIn("BEGIN RSA PRIVATE KEY", client_key_pem) fp = Server.enrollment.calculate_cert_fingerprint(client_cert_pem) self.assertEqual(len(fp), 64) finally: import shutil if os.path.exists(test_cert_dir): shutil.rmtree(test_cert_dir, ignore_errors=True) def test_enrollment_endpoint_and_seat_quota(self): from fastapi import HTTPException secret = "super-secret-enrollment" max_seats = 2 Server.init_db(self.test_db, enrollment_secret=secret, max_seats=max_seats) test_cert_dir = "test_certs_enroll" try: ca_cert, ca_key, ca_pem, _ = Server.enrollment.generate_ca_if_needed(cert_dir=test_cert_dir) Server.SERVER_STATE["config"] = {"db_path": self.test_db} Server.SERVER_STATE["ca_cert"] = ca_cert Server.SERVER_STATE["ca_key"] = ca_key Server.SERVER_STATE["ca_cert_pem"] = ca_pem # 1. Invalid secret should raise 403 bad_req = Server.ClientEnrollRequest( client_id="client-1", hostname="host-1", os="linux", enrollment_secret="wrong-secret" ) with self.assertRaises(HTTPException) as cm: Server.enroll_client(bad_req) self.assertEqual(cm.exception.status_code, 403) # 2. Valid enrollment for client 1 req1 = Server.ClientEnrollRequest( client_id="client-1", hostname="host-1", os="linux", enrollment_secret=secret ) resp1 = Server.enroll_client(req1) self.assertIn("client_cert", resp1) self.assertIn("client_key", resp1) self.assertEqual(resp1["ca_cert"], ca_pem) # 3. Valid enrollment for client 2 req2 = Server.ClientEnrollRequest( client_id="client-2", hostname="host-2", os="windows", enrollment_secret=secret ) resp2 = Server.enroll_client(req2) self.assertIn("client_cert", resp2) # 4. Seat quota exhausted: client 3 should raise 403 req3 = Server.ClientEnrollRequest( client_id="client-3", hostname="host-3", os="linux", enrollment_secret=secret ) with self.assertRaises(HTTPException) as cm: Server.enroll_client(req3) self.assertEqual(cm.exception.status_code, 403) self.assertIn("License seat limit reached", cm.exception.detail) # 5. Re-enrollment for existing client 1 should succeed resp1_re = Server.enroll_client(req1) self.assertIn("client_cert", resp1_re) # 6. Revoked client should be rejected conn = sqlite3.connect(self.test_db) conn.execute("UPDATE clients SET status = 'revoked' WHERE client_id = 'client-1'") conn.commit() conn.close() with self.assertRaises(HTTPException) as cm: Server.enroll_client(req1) self.assertEqual(cm.exception.status_code, 403) self.assertIn("revoked", cm.exception.detail) finally: import shutil if os.path.exists(test_cert_dir): shutil.rmtree(test_cert_dir, ignore_errors=True) if __name__ == "__main__": unittest.main()