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_signature_verification_and_anti_forgery(self):
        """Verifies cryptographic signature validation against master root CA and detection of forged certificates."""
        csr_pem, _ = enrollment.generate_ca_csr(cert_dir=self.server_cert_dir)
        sub_pem, root_pem = _create_mock_signed_license_cert(csr_pem, seats=30, license_id="LIC-VALID-30")

        # 1. Valid signature against genuine root CA
        self.assertTrue(enrollment.verify_certificate_signature(sub_pem, root_pem))
        info_valid = enrollment.extract_license_info_from_cert(sub_pem, upstream_ca_pem=root_pem)
        self.assertTrue(info_valid["signature_verified"])
        self.assertEqual(info_valid["seats"], 30)

        # 2. Fake / Rogue Root CA trying to verify the same cert
        fake_csr, _ = enrollment.generate_ca_csr(cert_dir=self.server_cert_dir)
        _, fake_root_pem = _create_mock_signed_license_cert(fake_csr, seats=999, license_id="LIC-FAKE")
        self.assertFalse(enrollment.verify_certificate_signature(sub_pem, fake_root_pem))
        info_invalid = enrollment.extract_license_info_from_cert(sub_pem, upstream_ca_pem=fake_root_pem)
        self.assertFalse(info_invalid["signature_verified"])

    def test_database_tampering_seat_reconciliation(self):
        """Verifies that direct DB manipulation of max_seats is reconciled against the certified X.509 quota."""
        from license_manager import LicenseManager

        csr_pem, _ = enrollment.generate_ca_csr(cert_dir=self.server_cert_dir)
        sub_pem, root_pem = _create_mock_signed_license_cert(csr_pem, seats=25, license_id="LIC-SEAT-25")

        # Save certified ca.crt and upstream_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)
        with open(os.path.join(self.server_cert_dir, "upstream_ca.crt"), "w", encoding="utf-8") as f:
            f.write(root_pem)

        Server.init_db(self.server_db, max_seats=25)
        server_state = {
            "config": {
                "db_path": self.server_db,
                "cert_dir": self.server_cert_dir,
                "max_seats": 25
            }
        }

        # Attacker attempts direct SQLite manipulation: SET max_seats = 99999
        conn = sqlite3.connect(self.server_db)
        conn.execute("UPDATE license_config SET max_seats = 99999 WHERE id = 1")
        conn.commit()
        conn.close()

        # LicenseManager inspects state and auto-reconciles max_seats to certified 25
        summary = LicenseManager.get_license_summary(server_state=server_state, db_path=self.server_db)
        self.assertEqual(summary["max_seats"], 25)
        self.assertTrue(summary["signature_verified"])

        # Check DB was reconciled
        conn = sqlite3.connect(self.server_db)
        row = conn.execute("SELECT max_seats FROM license_config WHERE id = 1").fetchone()
        conn.close()
        self.assertEqual(row[0], 25)

    def test_import_license_signature_validation(self):
        """Verifies that /api/license/import rejects certificates not signed by the upstream authority."""
        from fastapi import HTTPException
        csr_pem, _ = enrollment.generate_ca_csr(cert_dir=self.server_cert_dir)
        sub_pem, root_pem = _create_mock_signed_license_cert(csr_pem, seats=50, license_id="LIC-MOCK-50")

        # Generate mismatched root CA
        fake_csr, _ = enrollment.generate_ca_csr(cert_dir=self.server_cert_dir)
        _, rogue_root_pem = _create_mock_signed_license_cert(fake_csr, seats=10, license_id="LIC-ROGUE")

        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
        }

        # Attempt to import sub_pem with mismatched upstream_ca -> must raise HTTPException 400
        req_bad = Server.LicenseImportRequest(ca_cert=sub_pem, upstream_ca=rogue_root_pem)
        with self.assertRaises(HTTPException) as ctx:
            Server.import_signed_license(req_bad, auth_user="admin")
        self.assertEqual(ctx.exception.status_code, 400)
        self.assertIn("Cryptographic verification failed", ctx.exception.detail)

        # Import sub_pem with valid matching upstream_ca -> must succeed
        req_good = Server.LicenseImportRequest(ca_cert=sub_pem, upstream_ca=root_pem)
        res = Server.import_signed_license(req_good, auth_user="admin")
        self.assertEqual(res["status"], "success")
        self.assertEqual(res["seats"], 50)

    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"
        }

    def test_auto_update_lock_and_package_upload_protection(self):
        """Verifies that unlicensed servers block package upload with 403 and render locked UI."""
        from fastapi import HTTPException
        Server.init_db(self.server_db, max_seats=0)
        Server.SERVER_STATE["config"] = {
            "db_path": self.server_db,
            "cert_dir": self.server_cert_dir,
            "max_seats": 0
        }

        # Ensure license_config has no auto_update
        conn = sqlite3.connect(self.server_db)
        conn.execute("UPDATE license_config SET license_status = 'unlicensed', features = '[]' WHERE id = 1")
        conn.commit()
        conn.close()

        # WebUI should render locked banner and disabled upload controls
        resp = Server.serve_webui()
        html = resp.body.decode("utf-8")
        self.assertIn("Fleet Auto-Update is Inactive & Locked", html)
        self.assertIn("disabled", html)

    def test_json_lic_package_import(self):
        """Verifies importing an activated .lic JSON package via /api/license/import."""
        csr_pem, _ = enrollment.generate_ca_csr(cert_dir=self.server_cert_dir)
        sub_pem, root_pem = _create_mock_signed_license_cert(
            csr_pem, seats=30, license_id="LIC-TEST-JSON-01"
        )

        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
        }

        # JSON package format
        lic_pkg = {
            "format": "LOGAR_LICENSE_PACKAGE",
            "license_id": "LIC-TEST-JSON-01",
            "seats": 30,
            "features": ["enterprise_db", "auto_update"],
            "ca_cert": sub_pem,
            "master_root_ca": root_pem
        }

        req = Server.LicenseImportRequest(lic=json.dumps(lic_pkg))
        res = Server.import_signed_license(req, auth_user="admin")
        self.assertEqual(res["status"], "success")
        self.assertEqual(res["seats"], 30)


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