LOGAR/tests/test_license_authority.py

import os
import sys
import json
import shutil
import sqlite3
import unittest
from datetime import datetime, timezone, timedelta

from cryptography import x509
from cryptography.x509.oid import NameOID
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import rsa

# Ensure src/ is importable
ROOT_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
SRC_DIR = os.path.join(ROOT_DIR, "src")
if ROOT_DIR not in sys.path:
    sys.path.insert(0, ROOT_DIR)
if SRC_DIR not in sys.path:
    sys.path.insert(0, SRC_DIR)

import server_enrollment as enrollment
import Server


def _create_mock_signed_license_cert(csr_pem: str, seats: int = 50, license_id: str = "LIC-TEST-12345") -> tuple:
    """
    Helper to generate a mock master CA and sign a sub-CA CSR with LOGAR license OIDs
    using pure cryptography library (without importing the private License_Authority submodule).
    """
    # 1. Generate Root CA
    ca_key = rsa.generate_private_key(65537, 2048)
    ca_name = x509.Name([
        x509.NameAttribute(NameOID.COUNTRY_NAME, "AT"),
        x509.NameAttribute(NameOID.ORGANIZATION_NAME, "LOGAR Master Authority"),
        x509.NameAttribute(NameOID.COMMON_NAME, "LOGAR-Master-Root-CA")
    ])
    now = datetime.now(timezone.utc)
    ca_cert = (
        x509.CertificateBuilder()
        .subject_name(ca_name)
        .issuer_name(ca_name)
        .public_key(ca_key.public_key())
        .serial_number(x509.random_serial_number())
        .not_valid_before(now - timedelta(days=1))
        .not_valid_after(now + timedelta(days=3650))
        .add_extension(x509.BasicConstraints(ca=True, path_length=1), critical=True)
        .sign(ca_key, hashes.SHA256())
    )
    ca_pem = ca_cert.public_bytes(serialization.Encoding.PEM).decode("utf-8")

    # 2. Parse CSR and sign Subordinate CA with custom OIDs
    csr = x509.load_pem_x509_csr(csr_pem.encode("utf-8"))
    sub_cert = (
        x509.CertificateBuilder()
        .subject_name(csr.subject)
        .issuer_name(ca_name)
        .public_key(csr.public_key())
        .serial_number(x509.random_serial_number())
        .not_valid_before(now - timedelta(hours=1))
        .not_valid_after(now + timedelta(days=365))
        .add_extension(x509.BasicConstraints(ca=True, path_length=0), critical=True)
        .add_extension(x509.UnrecognizedExtension(enrollment.OID_LOGAR_SEATS, str(seats).encode("utf-8")), critical=False)
        .add_extension(x509.UnrecognizedExtension(enrollment.OID_LOGAR_LICENSE_ID, license_id.encode("utf-8")), critical=False)
        .add_extension(x509.UnrecognizedExtension(enrollment.OID_LOGAR_ISSUER_URL, b"https://logar.eibl.tech"), critical=False)
        .sign(ca_key, hashes.SHA256())
    )
    sub_pem = sub_cert.public_bytes(serialization.Encoding.PEM).decode("utf-8")
    return sub_pem, ca_pem


class TestServerLicensingAndWebUI(unittest.TestCase):
    def setUp(self):
        self.test_dir = os.path.join(ROOT_DIR, "test_lic_tmp")
        os.makedirs(self.test_dir, exist_ok=True)
        self.server_cert_dir = os.path.join(self.test_dir, "server_certs")
        self.server_db = os.path.join(self.test_dir, "server_state.db")

    def tearDown(self):
        if os.path.exists(self.test_dir):
            shutil.rmtree(self.test_dir, ignore_errors=True)

    def test_ca_csr_and_lic_request_generation(self):
        """Verifies subordinate CA CSR and .lic request file generation on the server."""
        lic_data, lic_json = enrollment.create_license_request_lic(
            cert_dir=self.server_cert_dir,
            license_server="https://logar.eibl.tech",
            server_fingerprint="AABBCCDDEEFF00112233",
            hostname="test-hub.corp.internal",
            requested_seats=30
        )
        self.assertEqual(lic_data["format"], "LOGAR_LICENSE_REQUEST")
        self.assertEqual(lic_data["license_server"], "https://logar.eibl.tech")
        self.assertEqual(lic_data["requested_seats"], 30)
        self.assertIn("BEGIN CERTIFICATE REQUEST", lic_data["csr"])
        self.assertTrue(os.path.exists(os.path.join(self.server_cert_dir, "ca.csr")))
        self.assertTrue(os.path.exists(os.path.join(self.server_cert_dir, "logar_server.lic")))

    def test_license_info_extraction_and_custom_oids(self):
        """Verifies extraction of embedded custom seat OIDs and license attributes from certificates."""
        csr_pem, _ = enrollment.generate_ca_csr(cert_dir=self.server_cert_dir)
        sub_pem, _ = _create_mock_signed_license_cert(csr_pem, seats=42, license_id="LIC-ABC-4242")

        info = enrollment.extract_license_info_from_cert(sub_pem)
        self.assertTrue(info["valid"])
        self.assertEqual(info["seats"], 42)
        self.assertEqual(info["license_id"], "LIC-ABC-4242")
        self.assertEqual(info["issuer_url"], "https://logar.eibl.tech")
        self.assertFalse(info["is_expired"])
        self.assertTrue(info["is_ca"])

    def test_server_license_import_and_seat_update(self):
        """Verifies that importing a signed CA cert updates the server's seat quota in DB and memory."""
        # 1. Generate CSR on server side
        csr_pem, _ = enrollment.generate_ca_csr(cert_dir=self.server_cert_dir)

        # 2. Mock sign it with 50 seats
        sub_pem, _ = _create_mock_signed_license_cert(csr_pem, seats=50, license_id="LIC-IMPORT-50")

        # 3. Setup mock server state and DB
        Server.init_db(self.server_db, max_seats=5)
        Server.SERVER_STATE["config"] = {
            "db_path": self.server_db,
            "cert_dir": self.server_cert_dir,
            "max_seats": 5,
            "tcp_host": "127.0.0.1",
            "tcp_port": 9443,
            "hermes_host": "127.0.0.1",
            "hermes_port": 8443,
            "server_fingerprint": "MOCK_FP",
            "server_name": "Test-Hub",
            "license_server_url": "https://logar.eibl.tech"
        }

        # 4. Extract summary before import
        summary_before = Server.get_license_summary()
        self.assertEqual(summary_before["max_seats"], 5)

        # 5. Save the signed CA certificate to cert_dir/ca.crt
        ca_path = os.path.join(self.server_cert_dir, "ca.crt")
        with open(ca_path, "w", encoding="utf-8") as f:
            f.write(sub_pem)

        # 6. Extract summary after import
        summary_after = Server.get_license_summary()
        self.assertEqual(summary_after["max_seats"], 50)
        self.assertEqual(summary_after["license_id"], "LIC-IMPORT-50")
        self.assertEqual(summary_after["status"], "ACTIVE")

    def test_server_client_revocation_and_reactivation(self):
        """Verifies revocation and reactivation of edge client seats on the server hub."""
        Server.init_db(self.server_db, max_seats=2)
        Server.SERVER_STATE["config"] = {
            "db_path": self.server_db,
            "cert_dir": self.server_cert_dir,
            "max_seats": 2
        }

        # Insert active clients
        conn = sqlite3.connect(self.server_db)
        conn.execute("INSERT INTO clients (client_id, hostname, os_type, cert_fingerprint, status) VALUES ('c1', 'win1', 'windows', 'FP1', 'active')")
        conn.execute("INSERT INTO clients (client_id, hostname, os_type, cert_fingerprint, status) VALUES ('c2', 'lin1', 'linux', 'FP2', 'active')")
        conn.commit()
        conn.close()

        # Revoke c1
        req = Server.ClientActionRequest(client_id="c1")
        res = Server.revoke_client(req)
        self.assertEqual(res["status"], "revoked")

        summary = Server.get_license_summary()
        self.assertEqual(summary["active_seats"], 1)
        self.assertEqual(summary["revoked_seats"], 1)
        self.assertEqual(summary["available_seats"], 1)

        # Reactivate c1
        res_react = Server.reactivate_client(req)
        self.assertEqual(res_react["status"], "active")

        summary2 = Server.get_license_summary()
        self.assertEqual(summary2["active_seats"], 2)
        self.assertEqual(summary2["available_seats"], 0)

    def test_activated_lic_package_handling(self):
        """Verifies that an activated .lic package structure can be processed and verified."""
        csr_pem, _ = enrollment.generate_ca_csr(cert_dir=self.server_cert_dir)
        sub_pem, root_pem = _create_mock_signed_license_cert(csr_pem, seats=20, license_id="LIC-ACT-20")

        lic_package = {
            "format": "LOGAR_ACTIVATED_LICENSE",
            "version": "2.0.1",
            "license_id": "LIC-ACT-20",
            "seats": 20,
            "server_fingerprint": "FP_TEST_HUB",
            "ca_cert": sub_pem,
            "master_root_ca": root_pem,
            "days_valid": 365
        }
        self.assertEqual(lic_package["seats"], 20)
        self.assertIn("BEGIN CERTIFICATE", lic_package["ca_cert"])
        self.assertIn("BEGIN CERTIFICATE", lic_package["master_root_ca"])

        # Verify certificate parses
        info = enrollment.extract_license_info_from_cert(lic_package["ca_cert"])
        self.assertEqual(info["seats"], 20)
        self.assertEqual(info["license_id"], "LIC-ACT-20")

    def test_webui_rendering(self):
        """Verifies that the WebUI HTML contains key components, stats, and download triggers."""
        Server.init_db(self.server_db, max_seats=10)
        Server.SERVER_STATE["config"] = {
            "db_path": self.server_db,
            "cert_dir": self.server_cert_dir,
            "max_seats": 10,
            "tcp_port": 9443,
            "hermes_port": 8443,
            "server_name": "Test-Cloud-Hub",
            "server_fingerprint": "1234567890ABCDEF",
            "license_server_url": "https://logar.eibl.tech"
        }

        response = Server.serve_webui()
        html = response.body.decode("utf-8")
        self.assertIn("LOGAR Fleet Hub", html)
        self.assertNotIn("LOGAR Fleet Hub & License Management", html)
        self.assertNotIn("Export .lic Request", html)
        self.assertNotIn("Import License (.lic / ca.crt)", html)
        self.assertIn("Sync & Renew via mTLS", html)
        self.assertIn("logar.eibl.tech", html)


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