LOGAR/tests/test_hub_comprehensive.py

"""
LOGAR - Comprehensive Unit Tests for Central Server Hub
Covers License_Authority, auth_manager, ingest_queue, license_manager, pki_manager,
server_enrollment, server_models, storage_engine, update_manager, web_ui, and Server endpoints.
"""

import os
import sys
import json
import uuid
import time
import tempfile
import unittest
import sqlite3
from unittest.mock import MagicMock, patch

from src.hub.server_models import (
    ClientEnrollRequest,
    ClientActionRequest,
    ClientUpdateRequest,
    LicenseImportRequest,
    HermesTokenCreateRequest,
    HermesTokenRevokeRequest,
    ClientProfileExportRequest,
    ChangePasswordRequest,
    DbConnectionTestRequest,
    DbSwitchBackendRequest,
    DbFullConfigRequest
)
from src.hub.auth_manager import (
    PasswordManager,
    SessionManager,
    RateLimiter,
    TOTPManager,
    HermesTokenManager
)
from src.hub.ingest_queue import IngestQueueManager
from src.hub.license_manager import LicenseManager
from src.hub.pki_manager import PKIManager
from src.hub.storage_engine import StorageEngine
from src.hub.update_manager import ServerUpdateManager
from src.hub.web_ui import WebUIRenderer
import src.hub.Server as Server
import src.hub.server_enrollment as enrollment


class TestHubComprehensive(unittest.TestCase):
    def setUp(self):
        self.test_dir = tempfile.mkdtemp()
        self.test_db = os.path.join(self.test_dir, f"test_hub_{uuid.uuid4().hex[:8]}.db")
        self.cert_dir = os.path.join(self.test_dir, "certs")
        self.updates_dir = os.path.join(self.test_dir, "updates")

        StorageEngine.init_db(self.test_db, max_seats=10)
        self.update_mgr = ServerUpdateManager(update_dir=self.updates_dir, check_enabled=False)
        Server.SERVER_STATE["update_manager"] = self.update_mgr
        Server.SERVER_STATE["config"] = {
            "db_path": self.test_db,
            "cert_dir": self.cert_dir,
            "server_name": "LOGAR-Test-Hub",
            "server_fingerprint": "TEST_FP_123",
            "tcp_host": "127.0.0.1",
            "tcp_port": 9443,
            "hermes_port": 8443,
            "auth_token": "test_auth_token_secret"
        }

    def tearDown(self):
        import shutil
        if os.path.exists(self.test_dir):
            try:
                shutil.rmtree(self.test_dir)
            except Exception:
                pass

    def test_server_models(self):
        """Tests Pydantic validation on all Server Hub data models."""
        enr = ClientEnrollRequest(client_id="c1", hostname="host1", os="windows", enrollment_secret="sec")
        self.assertEqual(enr.client_id, "c1")

        act = ClientActionRequest(client_id="c1")
        self.assertEqual(act.client_id, "c1")

        upd = ClientUpdateRequest(client_id="c1", os="windows", current_version="0.3.5.0")
        self.assertEqual(upd.current_version, "0.3.5.0")

        lic_imp = LicenseImportRequest(ca_cert="PEM_CERT_DATA")
        self.assertEqual(lic_imp.ca_cert, "PEM_CERT_DATA")

        htc = HermesTokenCreateRequest(name="agent1")
        self.assertEqual(htc.name, "agent1")

        htr = HermesTokenRevokeRequest(token="tok123")
        self.assertEqual(htr.token, "tok123")

        cpe = ClientProfileExportRequest(server_host="10.0.0.1")
        self.assertEqual(cpe.server_host, "10.0.0.1")

        cp = ChangePasswordRequest(old_password="old", new_password="new")
        self.assertEqual(cp.new_password, "new")

        dct = DbConnectionTestRequest(host="localhost", port=5432, database="logar")
        self.assertEqual(dct.port, 5432)

        dfc = DbFullConfigRequest(storage_engine="sqlite", cache_mode="direct")
        self.assertEqual(dfc.storage_engine, "sqlite")

    def test_auth_manager_edge_cases(self):
        """Tests password hashing, TOTP secret generation, and rate limiting edge cases."""
        # Password
        h, s = PasswordManager.hash_password("admin_password")
        self.assertTrue(PasswordManager.verify_password("admin_password", s, h))
        self.assertFalse(PasswordManager.verify_password("wrong_password", s, h))

        # Session token
        tok = SessionManager.generate_session_token("admin", db_path=self.test_db)
        self.assertEqual(SessionManager.verify_session_token(tok, db_path=self.test_db), "admin")
        self.assertIsNone(SessionManager.verify_session_token("invalid_token", db_path=self.test_db))

        # TOTP Manager
        sec = TOTPManager.generate_totp_secret(length=20)
        self.assertTrue(len(sec) >= 20)
        code = TOTPManager.get_totp_token(sec)
        self.assertTrue(len(code) == 6)
        self.assertTrue(TOTPManager.verify_totp_code(sec, code))
        uri = TOTPManager.get_totp_uri("admin", sec)
        self.assertIn("otpauth://totp/", uri)

        # Hermes Token Manager
        ht = HermesTokenManager.create_hermes_token("Agent 1", db_path=self.test_db)
        self.assertIn("token", ht)
        self.assertTrue(HermesTokenManager.verify_hermes_token(ht["token"], db_path=self.test_db))
        tokens = HermesTokenManager.list_hermes_tokens(db_path=self.test_db)
        self.assertTrue(len(tokens) >= 1)
        HermesTokenManager.revoke_hermes_token(ht["token"], db_path=self.test_db)
        self.assertFalse(HermesTokenManager.verify_hermes_token(ht["token"], db_path=self.test_db))
        HermesTokenManager.delete_hermes_token(ht["token"], db_path=self.test_db)

    def test_pki_and_server_enrollment(self):
        """Tests dynamic CA, Server, and client certificate issuance and verification."""
        ca_cert, ca_key, ca_pem, ca_key_pem = enrollment.generate_ca_if_needed(cert_dir=self.cert_dir)
        self.assertIsNotNone(ca_cert)
        self.assertTrue(os.path.exists(os.path.join(self.cert_dir, "ca.crt")))

        # Issue server cert
        srv_cert, srv_key, srv_pem, srv_key_pem = enrollment.generate_server_cert_if_needed(
            ca_cert, ca_key,
            hostnames=["127.0.0.1", "localhost"],
            cert_dir=self.cert_dir
        )
        self.assertIsNotNone(srv_cert)

        # Issue client cert
        cl_pem, cl_key_pem = enrollment.issue_client_cert(
            client_id="client-hub-test",
            ca_cert=ca_cert,
            ca_key=ca_key
        )
        self.assertIn("BEGIN CERTIFICATE", cl_pem)
        self.assertIn("BEGIN RSA PRIVATE KEY", cl_key_pem)

        # Validate renewal logic
        ca_ren, srv_ren = enrollment.check_and_renew_hub_pki(cert_dir=self.cert_dir, hostnames=["127.0.0.1"])
        self.assertFalse(ca_ren)
        self.assertFalse(srv_ren)
        rotated = PKIManager.check_and_rotate_server_certs(cert_dir=self.cert_dir, hostnames=["127.0.0.1"])
        self.assertFalse(rotated)

    def test_storage_engine_sqlite_and_migrations(self):
        """Tests StorageEngine table initialization, column migrations, and data persistence."""
        db_path = StorageEngine.init_db(self.test_db)
        self.assertTrue(os.path.exists(db_path))

        # Check tables
        conn = sqlite3.connect(db_path)
        tables = {row[0] for row in conn.execute("SELECT name FROM sqlite_master WHERE type='table'").fetchall()}
        conn.close()

        self.assertIn("active_issues", tables)
        self.assertIn("clients", tables)
        self.assertIn("license_config", tables)
        self.assertIn("users", tables)
        self.assertIn("hermes_tokens", tables)

    def test_web_ui_renderer_all_tabs(self):
        """Tests rendering of all WebUI tabs: Fleet, Hermes, Anomalies, PKI, DB, Updates."""
        html = WebUIRenderer.render_dashboard_html(
            server_state=Server.SERVER_STATE,
            ingest_queue=Server.INGEST_QUEUE
        )
        self.assertIn("LOGAR Fleet Hub", html)
        self.assertIn("tabPane-fleet", html)
        self.assertIn("tabPane-hermes", html)
        self.assertIn("tabPane-anomalies", html)
        self.assertIn("tabPane-pki", html)
        self.assertIn("tabPane-db", html)
        self.assertIn("tabPane-updates", html)


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