LOGAR/tests/test_server.py

import os
import sys
import json
import socket
import struct
import sqlite3
import unittest
import urllib.request
import uuid
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 = f"test_server_state_{uuid.uuid4().hex[:8]}.db"
        self.test_config = f"test_server_config_{uuid.uuid4().hex[:8]}.json"

    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_direct_hermes_anomaly_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]
        }

        # Run 1: WARNING is immediately VERIFIED for Hermes AI reasoning
        res1 = Server.process_ingested_logs(payload, self.test_db, window_hours=12, min_runs=1)
        self.assertEqual(res1["status"], "success")
        self.assertEqual(res1["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], 1)
        self.assertEqual(row[1], "VERIFIED")

        # Runs 2 to 4: Increments consecutive run count while remaining VERIFIED
        for run_idx in range(2, 5):
            res = Server.process_ingested_logs(payload, self.test_db, window_hours=12, min_runs=1)
            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], 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=1)
        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=1)
        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)

        # All actionable issues (ERROR, WARNING, INFO) are promoted to VERIFIED immediately for Hermes reasoning
        self.assertEqual(rows["SigErr"], "VERIFIED")
        self.assertEqual(rows["SigWarn"], "VERIFIED")
        self.assertEqual(rows["SigInfo"], "VERIFIED")

    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, "cert_dir": test_cert_dir}
            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.assertTrue(resp1["ca_cert"].strip().startswith(ca_pem.strip()))

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

    def test_cert_validity_and_hub_pki_renewal(self):
        test_cert_dir = "test_certs_renew"
        try:
            ca_cert, ca_key, ca_pem, _ = Server.enrollment.generate_ca_if_needed(cert_dir=test_cert_dir)
            srv_cert, srv_key, srv_pem, _ = Server.enrollment.generate_server_cert_if_needed(
                ca_cert, ca_key, hostnames=["127.0.0.1"], cert_dir=test_cert_dir
            )

            # 1. Freshly generated certificates should NOT be expiring soon with standard 30-day threshold
            self.assertFalse(Server.enrollment.is_cert_expiring_soon(ca_pem, threshold_days=30))
            self.assertFalse(Server.enrollment.is_cert_expiring_soon(srv_pem, threshold_days=30))

            # 2. Huge threshold (e.g. 5000 days) should flag expiration
            self.assertTrue(Server.enrollment.is_cert_expiring_soon(srv_pem, threshold_days=5000))

            # 3. check_and_renew_hub_pki with standard threshold should report no renewal needed
            ca_renewed, srv_renewed = Server.enrollment.check_and_renew_hub_pki(cert_dir=test_cert_dir, threshold_days=30)
            self.assertFalse(ca_renewed)
            self.assertFalse(srv_renewed)

            # 4. In-flight reload of SSLContext
            ssl_ctx = Server.init_mtls_server_context(cert_dir=test_cert_dir)
            Server.SERVER_STATE["ssl_ctx"] = ssl_ctx
            Server.SERVER_STATE["config"] = {"db_path": self.test_db, "cert_dir": test_cert_dir, "tcp_host": "127.0.0.1"}

            # Trigger rotation using high threshold
            rotated = Server.check_and_rotate_server_certs(cert_dir=test_cert_dir, hostnames=["127.0.0.1"], threshold_days=5000)
            self.assertTrue(rotated)

            # Check that backup files were generated
            bak_files = [f for f in os.listdir(test_cert_dir) if f.endswith(".bak")]
            self.assertGreater(len(bak_files), 0)
        finally:
            import shutil
            if os.path.exists(test_cert_dir):
                shutil.rmtree(test_cert_dir, ignore_errors=True)

    def test_user_auth_and_session(self):
        """Tests admin password hashing, verification, session tokens, and password change."""
        Server.init_db(self.test_db)
        
        # Verify default admin credentials
        self.assertTrue(Server.verify_user_credentials("admin", "admin_logar_2026!", db_path=self.test_db))
        self.assertFalse(Server.verify_user_credentials("admin", "wrong_password", db_path=self.test_db))
        self.assertFalse(Server.verify_user_credentials("nonexistent", "password", db_path=self.test_db))

        # Test session token generation & validation
        token = Server.generate_session_token("admin")
        self.assertEqual(Server.verify_session_token(token), "admin")
        self.assertIsNone(Server.verify_session_token("invalid.token.structure"))

        # Test changing password
        Server.change_user_password("admin", "new_secure_pass123", db_path=self.test_db)
        self.assertFalse(Server.verify_user_credentials("admin", "admin_logar_2026!", db_path=self.test_db))
        self.assertTrue(Server.verify_user_credentials("admin", "new_secure_pass123", db_path=self.test_db))

    def test_hermes_tokens_and_auth(self):
        """Tests Hermes Bearer token creation, verification, revocation, and deletion."""
        Server.init_db(self.test_db)
        
        # Test default Hermes token was created on init
        tokens = Server.list_hermes_tokens(db_path=self.test_db)
        self.assertEqual(len(tokens), 1)
        self.assertEqual(tokens[0]["status"], "active")
        default_token = tokens[0]["token"]
        self.assertTrue(Server.verify_hermes_token(default_token, db_path=self.test_db))

        # Test creating new token
        new_tok = Server.create_hermes_token(name="Hermes Autonomous Agent 01", db_path=self.test_db)
        self.assertIn("token", new_tok)
        self.assertEqual(new_tok["name"], "Hermes Autonomous Agent 01")
        self.assertTrue(Server.verify_hermes_token(new_tok["token"], db_path=self.test_db))

        # Test revoking token
        self.assertTrue(Server.revoke_hermes_token(new_tok["token"], db_path=self.test_db))
        self.assertFalse(Server.verify_hermes_token(new_tok["token"], db_path=self.test_db))

        # Test deleting token
        self.assertTrue(Server.delete_hermes_token(new_tok["token"], db_path=self.test_db))
        tokens_after = Server.list_hermes_tokens(db_path=self.test_db)
        self.assertEqual(len(tokens_after), 1)

    def test_client_profile_export(self):
        """Tests client profile dictionary generation and export formatting."""
        cfg = Server.load_or_init_config(self.test_config)
        profile = Server.generate_client_profile_dict(
            server_host="logar-hub.corp.internal",
            server_port=9443,
            client_id="worker-node-101",
            config_dict=cfg
        )
        self.assertEqual(profile["server_host"], "logar-hub.corp.internal")
        self.assertEqual(profile["server_port"], 9443)
        self.assertEqual(profile["client_id"], "worker-node-101")
        self.assertEqual(profile["server_fingerprint"], cfg["server_fingerprint"])
        self.assertEqual(profile["server_public_key"], cfg["public_key"])
        self.assertEqual(profile["enrollment_secret"], cfg["enrollment_secret"])
        self.assertEqual(profile["format"], "LOGAR_CLIENT_PROFILE")

    def test_client_seat_management_and_relative_time(self):
        """Tests client listing, seat approval, revocation, deletion, and relative timestamp helper."""
        Server.init_db(self.test_db, enrollment_secret="test-sec-123", max_seats=5)
        Server.SERVER_STATE["config"] = {"db_path": self.test_db}

        # Populate test clients
        conn = sqlite3.connect(self.test_db)
        now = datetime.now(timezone.utc)
        five_min_ago = (now - timedelta(minutes=5)).strftime("%Y-%m-%d %H:%M:%S")
        two_hours_ago = (now - timedelta(hours=2)).strftime("%Y-%m-%d %H:%M:%S")
        
        conn.execute("INSERT INTO clients (client_id, hostname, os_type, cert_fingerprint, status, last_seen) VALUES ('c-win', 'WIN-PROD-01', 'windows', 'FP_W1', 'active', ?)", (five_min_ago,))
        conn.execute("INSERT INTO clients (client_id, hostname, os_type, cert_fingerprint, status, last_seen) VALUES ('c-lin', 'LIN-PROD-01', 'linux', 'FP_L1', 'active', ?)", (two_hours_ago,))
        conn.commit()
        conn.close()

        # 1. Test listing
        client_data = Server.list_clients()
        self.assertEqual(client_data["active_seats"], 2)
        self.assertEqual(client_data["max_seats"], 5)
        self.assertEqual(len(client_data["clients"]), 2)
        self.assertEqual(client_data["clients"][0]["last_activity_relative"], "5m ago")

        # 2. Test revoking
        rev_req = Server.ClientActionRequest(client_id="c-win")
        rev_res = Server.revoke_client(rev_req)
        self.assertEqual(rev_res["status"], "revoked")
        client_data = Server.list_clients()
        self.assertEqual(client_data["active_seats"], 1)

        # 3. Test approving / reactivating
        app_res = Server.approve_client(rev_req)
        self.assertEqual(app_res["status"], "active")
        client_data = Server.list_clients()
        self.assertEqual(client_data["active_seats"], 2)

        # 4. Test deleting client record
        del_res = Server.delete_client(rev_req)
        self.assertEqual(del_res["status"], "deleted")
        client_data = Server.list_clients()
        self.assertEqual(len(client_data["clients"]), 1)
        self.assertEqual(client_data["active_seats"], 1)

        # 5. Test format_relative_time edge cases
        self.assertEqual(Server.format_relative_time(None), "Never")
        self.assertEqual(Server.format_relative_time(""), "Never")
        self.assertEqual(Server.format_relative_time("invalid-date-string"), "invalid-date-str")
        now_fresh = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S")
        rel = Server.format_relative_time(now_fresh)
        self.assertTrue(rel.endswith("s ago") or rel == "Just now", f"Unexpected relative time: {rel}")

    def test_integrated_local_log_collector_zero_seats(self):
        """Tests that local hub log collection runs without license, registering a 0-seat built-in client."""
        Server.init_db(self.test_db, enrollment_secret="test-sec-456", max_seats=0)
        Server.SERVER_STATE["config"] = {
            "db_path": self.test_db,
            "evaluation_window_hours": 12,
            "min_persistence_runs": 1
        }

        # Run cycle
        state = {}
        processed = Server.run_local_log_collection_cycle(self.test_db, window_hours=12, min_runs=1, state=state)
        
        # Verify local hub client entry exists in clients table
        machine_id = Server.get_machine_identifier()
        conn = sqlite3.connect(self.test_db)
        c = conn.cursor()
        c.execute("SELECT client_id, cert_fingerprint, status FROM clients WHERE client_id = ?", (machine_id,))
        row = c.fetchone()
        conn.close()

        self.assertIsNotNone(row)
        self.assertEqual(row[0], machine_id)
        self.assertEqual(row[1], "LOCAL-HOST-BUILTIN")
        self.assertEqual(row[2], "active")

        # Verify active seat count is STILL 0 (exempt from seat accounting)
        client_data = Server.list_clients()
        self.assertEqual(client_data["active_seats"], 0)
        self.assertEqual(len(client_data["clients"]), 1)

    def test_two_factor_auth_and_totp_flow(self):
        """Tests TOTP secret generation, token verification, URI generation, and enable/disable lifecycle."""
        Server.init_db(self.test_db)
        
        # 1. Test TOTP helpers
        secret = Server.generate_totp_secret()
        self.assertTrue(len(secret) >= 16)
        token = Server.get_totp_token(secret)
        self.assertEqual(len(token), 6)
        self.assertTrue(token.isdigit())
        self.assertTrue(Server.verify_totp_code(secret, token))
        self.assertFalse(Server.verify_totp_code(secret, "000000" if token != "000000" else "999999"))
        self.assertFalse(Server.verify_totp_code(secret, "invalid"))

        # 2. Test OTPAuth URI
        uri = Server.get_totp_uri("admin", secret, issuer="LOGAR Fleet Hub")
        self.assertTrue(uri.startswith("otpauth://totp/"))
        self.assertIn("secret=" + secret, uri)
        self.assertIn("LOGAR", uri)

        # 3. Test initial user 2FA status
        info = Server.get_user_totp_info("admin", db_path=self.test_db)
        self.assertFalse(info["enabled"])
        self.assertEqual(info["secret"], "")

        # 4. Test enabling 2FA
        enabled = Server.enable_user_totp("admin", secret, db_path=self.test_db)
        self.assertTrue(enabled)

        info_active = Server.get_user_totp_info("admin", db_path=self.test_db)
        self.assertTrue(info_active["enabled"])
        self.assertEqual(info_active["secret"], secret)

        # 5. Test disabling 2FA
        disabled = Server.disable_user_totp("admin", db_path=self.test_db)
        self.assertTrue(disabled)

        info_disabled = Server.get_user_totp_info("admin", db_path=self.test_db)
        self.assertFalse(info_disabled["enabled"])


if __name__ == "__main__":
    unittest.main()