#!/usr/bin/env python3
"""
Test PKI certificate validation for Ansible-generated certificates
Hybrid approach: validates actual certificates when available, else tests templates/config
Based on issues #14755, #14718 - Apple device compatibility
Issues #75, #153 - Security enhancements (name constraints, EKU restrictions)
"""
import glob
import os
import re
import subprocess
import sys
from datetime import UTC

from cryptography import x509
from cryptography.x509.oid import ExtensionOID, NameOID


def find_generated_certificates():
    """Find Ansible-generated certificate files in configs directory"""
    # Look for configs directory structure created by Ansible
    config_patterns = [
        "configs/*/ipsec/.pki/cacert.pem",
        "../configs/*/ipsec/.pki/cacert.pem",  # From tests/unit directory
        "../../configs/*/ipsec/.pki/cacert.pem"  # Alternative path
    ]

    for pattern in config_patterns:
        ca_certs = glob.glob(pattern)
        if ca_certs:
            base_path = os.path.dirname(ca_certs[0])
            return {
                'ca_cert': ca_certs[0],
                'base_path': base_path,
                'server_certs': glob.glob(f"{base_path}/certs/*.crt"),
                'p12_files': glob.glob(f"{base_path.replace('/.pki', '')}/manual/*.p12")
            }

    return None

def test_openssl_version_detection():
    """Test that we can detect OpenSSL version for compatibility checks"""
    result = subprocess.run(
        ['openssl', 'version'],
        capture_output=True,
        text=True
    )

    assert result.returncode == 0, "Failed to get OpenSSL version"

    # Parse version - e.g., "OpenSSL 3.0.2 15 Mar 2022"
    version_match = re.search(r'OpenSSL\s+(\d+)\.(\d+)\.(\d+)', result.stdout)
    assert version_match, f"Can't parse OpenSSL version: {result.stdout}"

    major = int(version_match.group(1))
    minor = int(version_match.group(2))

    print(f"✓ OpenSSL version detected: {major}.{minor}")
    return (major, minor)


def validate_ca_certificate_real(cert_files):
    """Validate actual Ansible-generated CA certificate"""
    # Read the actual CA certificate generated by Ansible
    with open(cert_files['ca_cert'], 'rb') as f:
        cert_data = f.read()

    certificate = x509.load_pem_x509_certificate(cert_data)

    # Check Basic Constraints
    basic_constraints = certificate.extensions.get_extension_for_oid(ExtensionOID.BASIC_CONSTRAINTS).value
    assert basic_constraints.ca is True, "CA certificate should have CA:TRUE"
    assert basic_constraints.path_length == 0, "CA should have pathlen:0 constraint"

    # Check Key Usage
    key_usage = certificate.extensions.get_extension_for_oid(ExtensionOID.KEY_USAGE).value
    assert key_usage.key_cert_sign is True, "CA should have keyCertSign usage"
    assert key_usage.crl_sign is True, "CA should have cRLSign usage"

    # Check Extended Key Usage (Issue #75)
    eku = certificate.extensions.get_extension_for_oid(ExtensionOID.EXTENDED_KEY_USAGE).value
    assert x509.oid.ExtendedKeyUsageOID.SERVER_AUTH in eku, "CA should allow signing server certificates"
    assert x509.oid.ExtendedKeyUsageOID.CLIENT_AUTH in eku, "CA should allow signing client certificates"
    assert x509.ObjectIdentifier("1.3.6.1.5.5.7.3.17") in eku, "CA should have IPsec End Entity EKU"

    # Check Name Constraints (Issue #75) - defense against certificate misuse
    name_constraints = certificate.extensions.get_extension_for_oid(ExtensionOID.NAME_CONSTRAINTS).value
    assert name_constraints.permitted_subtrees is not None, "CA should have permitted name constraints"
    assert name_constraints.excluded_subtrees is not None, "CA should have excluded name constraints"

    # Verify public domains are excluded
    excluded_dns = [constraint.value for constraint in name_constraints.excluded_subtrees
                   if isinstance(constraint, x509.DNSName)]
    public_domains = [".com", ".org", ".net", ".gov", ".edu", ".mil", ".int"]
    for domain in public_domains:
        assert domain in excluded_dns, f"CA should exclude public domain {domain}"

    # Verify private IP ranges are excluded (Issue #75)
    excluded_ips = [constraint.value for constraint in name_constraints.excluded_subtrees
                   if isinstance(constraint, x509.IPAddress)]
    assert len(excluded_ips) > 0, "CA should exclude private IP ranges"

    # Verify email domains are also excluded (Issue #153)
    excluded_emails = [constraint.value for constraint in name_constraints.excluded_subtrees
                      if isinstance(constraint, x509.RFC822Name)]
    email_domains = [".com", ".org", ".net", ".gov", ".edu", ".mil", ".int"]
    for domain in email_domains:
        assert domain in excluded_emails, f"CA should exclude email domain {domain}"

    print(f"✓ Real CA certificate has proper security constraints: {cert_files['ca_cert']}")

def validate_ca_certificate_config():
    """Validate CA certificate configuration in Ansible files (CI mode)"""
    # Check that the Ansible task file has proper CA certificate configuration
    openssl_task_file = find_ansible_file('roles/strongswan/tasks/openssl.yml')
    if not openssl_task_file:
        print("⚠ Could not find openssl.yml task file")
        return

    with open(openssl_task_file) as f:
        content = f.read()

    # Verify key security configurations are present
    security_checks = [
        ('name_constraints_permitted', 'Name constraints should be configured'),
        ('name_constraints_excluded', 'Excluded name constraints should be configured'),
        ('extended_key_usage', 'Extended Key Usage should be configured'),
        ('1.3.6.1.5.5.7.3.17', 'IPsec End Entity OID should be present'),
        ('serverAuth', 'Server authentication EKU should be present'),
        ('clientAuth', 'Client authentication EKU should be present'),
        ('basic_constraints', 'Basic constraints should be configured'),
        ('CA:TRUE', 'CA certificate should be marked as CA'),
        ('pathlen:0', 'Path length constraint should be set')
    ]

    for check, message in security_checks:
        assert check in content, f"Missing security configuration: {message}"

    # Verify public domains are excluded
    public_domains = [".com", ".org", ".net", ".gov", ".edu", ".mil", ".int"]
    for domain in public_domains:
        # Handle both double quotes and single quotes in YAML
        assert f'"DNS:{domain}"' in content or f"'DNS:{domain}'" in content, f"Public domain {domain} should be excluded"

    # Verify private IP ranges are excluded
    private_ranges = ["10.0.0.0", "172.16.0.0", "192.168.0.0"]
    for ip_range in private_ranges:
        assert ip_range in content, f"Private IP range {ip_range} should be excluded"

    # Verify email domains are excluded (Issue #153)
    email_domains = [".com", ".org", ".net", ".gov", ".edu", ".mil", ".int"]
    for domain in email_domains:
        # Handle both double quotes and single quotes in YAML
        assert f'"email:{domain}"' in content or f"'email:{domain}'" in content, f"Email domain {domain} should be excluded"

    # Verify IPv6 constraints are present (Issue #153)
    assert "IP:::/0" in content, "IPv6 all addresses should be excluded"

    print("✓ CA certificate configuration has proper security constraints")

def test_ca_certificate():
    """Test CA certificate - uses real certs if available, else validates config (Issue #75, #153)"""
    cert_files = find_generated_certificates()
    if cert_files:
        validate_ca_certificate_real(cert_files)
    else:
        validate_ca_certificate_config()


def validate_server_certificates_real(cert_files):
    """Validate actual Ansible-generated server certificates"""
    # Filter to only actual server certificates (not client certs)
    # Server certificates contain IP addresses in the filename
    import re
    server_certs = [f for f in cert_files['server_certs']
                   if not f.endswith('/cacert.pem') and re.search(r'\d+\.\d+\.\d+\.\d+\.crt$', f)]
    if not server_certs:
        print("⚠ No server certificates found")
        return

    for server_cert_path in server_certs:
        with open(server_cert_path, 'rb') as f:
            cert_data = f.read()

        certificate = x509.load_pem_x509_certificate(cert_data)

        # Check it's not a CA certificate
        basic_constraints = certificate.extensions.get_extension_for_oid(ExtensionOID.BASIC_CONSTRAINTS).value
        assert basic_constraints.ca is False, "Server certificate should not be a CA"

        # Check Extended Key Usage (Issue #75)
        eku = certificate.extensions.get_extension_for_oid(ExtensionOID.EXTENDED_KEY_USAGE).value
        assert x509.oid.ExtendedKeyUsageOID.SERVER_AUTH in eku, "Server cert must have serverAuth EKU"
        assert x509.ObjectIdentifier("1.3.6.1.5.5.7.3.17") in eku, "Server cert should have IPsec End Entity EKU"
        # Security check: Server certificates should NOT have clientAuth to prevent role confusion (Issue #153)
        assert x509.oid.ExtendedKeyUsageOID.CLIENT_AUTH not in eku, "Server cert should NOT have clientAuth EKU for role separation"

        # Check SAN extension exists (required for Apple devices)
        try:
            san = certificate.extensions.get_extension_for_oid(ExtensionOID.SUBJECT_ALTERNATIVE_NAME).value
            assert len(san) > 0, "Server certificate must have SAN extension for Apple device compatibility"
        except x509.ExtensionNotFound:
            assert False, "Server certificate missing SAN extension - required for modern clients"

        print(f"✓ Real server certificate valid: {os.path.basename(server_cert_path)}")

def validate_server_certificates_config():
    """Validate server certificate configuration in Ansible files (CI mode)"""
    openssl_task_file = find_ansible_file('roles/strongswan/tasks/openssl.yml')
    if not openssl_task_file:
        print("⚠ Could not find openssl.yml task file")
        return

    with open(openssl_task_file) as f:
        content = f.read()

    # Look for server certificate CSR section
    server_csr_section = re.search(r'Create CSRs for server certificate.*?register: server_csr', content, re.DOTALL)
    if not server_csr_section:
        print("⚠ Could not find server certificate CSR section")
        return

    server_section = server_csr_section.group(0)

    # Check server certificate CSR configuration
    server_checks = [
        ('subject_alt_name', 'Server certificates should have SAN extension'),
        ('serverAuth', 'Server certificates should have serverAuth EKU'),
        ('1.3.6.1.5.5.7.3.17', 'Server certificates should have IPsec End Entity EKU'),
        ('digitalSignature', 'Server certificates should have digital signature usage'),
        ('keyEncipherment', 'Server certificates should have key encipherment usage')
    ]

    for check, message in server_checks:
        assert check in server_section, f"Missing server certificate configuration: {message}"

    # Security check: Server certificates should NOT have clientAuth (Issue #153)
    # Look for clientAuth in extended_key_usage section, not in comments
    eku_lines = [line for line in server_section.split('\n') if 'extended_key_usage:' in line or (line.strip().startswith('- ') and 'clientAuth' in line)]
    has_client_auth = any('clientAuth' in line for line in eku_lines if line.strip().startswith('- '))
    assert not has_client_auth, "Server certificates should NOT have clientAuth EKU for role separation"

    # Verify SAN extension is configured for Apple compatibility
    assert 'subjectAltName' in server_section, "Server certificates missing SAN configuration for Apple compatibility"

    print("✓ Server certificate configuration has proper EKU and SAN settings")

def test_server_certificates():
    """Test server certificates - uses real certs if available, else validates config"""
    cert_files = find_generated_certificates()
    if cert_files:
        validate_server_certificates_real(cert_files)
    else:
        validate_server_certificates_config()


def validate_client_certificates_real(cert_files):
    """Validate actual Ansible-generated client certificates"""
    # Find client certificates (not CA cert, not server cert with IP/DNS name)
    client_certs = []
    for cert_path in cert_files['server_certs']:
        if 'cacert.pem' in cert_path:
            continue

        with open(cert_path, 'rb') as f:
            cert_data = f.read()
        certificate = x509.load_pem_x509_certificate(cert_data)

        # Check if this looks like a client cert vs server cert
        cn = certificate.subject.get_attributes_for_oid(NameOID.COMMON_NAME)[0].value
        # Server certs typically have IP addresses or domain names as CN
        if not (cn.replace('.', '').isdigit() or '.' in cn and len(cn.split('.')) == 4):
            client_certs.append((cert_path, certificate))

    if not client_certs:
        print("⚠ No client certificates found")
        return

    for cert_path, certificate in client_certs:
        # Check it's not a CA certificate
        basic_constraints = certificate.extensions.get_extension_for_oid(ExtensionOID.BASIC_CONSTRAINTS).value
        assert basic_constraints.ca is False, "Client certificate should not be a CA"

        # Check Extended Key Usage restrictions (Issue #75)
        eku = certificate.extensions.get_extension_for_oid(ExtensionOID.EXTENDED_KEY_USAGE).value
        assert x509.oid.ExtendedKeyUsageOID.CLIENT_AUTH in eku, "Client cert must have clientAuth EKU"
        assert x509.ObjectIdentifier("1.3.6.1.5.5.7.3.17") in eku, "Client cert should have IPsec End Entity EKU"

        # Security check: Client certificates should NOT have serverAuth (prevents impersonation) (Issue #153)
        assert x509.oid.ExtendedKeyUsageOID.SERVER_AUTH not in eku, "Client cert must NOT have serverAuth EKU to prevent server impersonation"

        # Check SAN extension for email
        try:
            san = certificate.extensions.get_extension_for_oid(ExtensionOID.SUBJECT_ALTERNATIVE_NAME).value
            email_sans = [name.value for name in san if isinstance(name, x509.RFC822Name)]
            assert len(email_sans) > 0, "Client certificate should have email SAN"
        except x509.ExtensionNotFound:
            print(f"⚠ Client certificate missing SAN extension: {os.path.basename(cert_path)}")

        print(f"✓ Real client certificate valid: {os.path.basename(cert_path)}")

def validate_client_certificates_config():
    """Validate client certificate configuration in Ansible files (CI mode)"""
    openssl_task_file = find_ansible_file('roles/strongswan/tasks/openssl.yml')
    if not openssl_task_file:
        print("⚠ Could not find openssl.yml task file")
        return

    with open(openssl_task_file) as f:
        content = f.read()

    # Look for client certificate CSR section
    client_csr_section = re.search(r'Create CSRs for client certificates.*?register: client_csr_jobs', content, re.DOTALL)
    if not client_csr_section:
        print("⚠ Could not find client certificate CSR section")
        return

    client_section = client_csr_section.group(0)

    # Check client certificate configuration
    client_checks = [
        ('clientAuth', 'Client certificates should have clientAuth EKU'),
        ('1.3.6.1.5.5.7.3.17', 'Client certificates should have IPsec End Entity EKU'),
        ('digitalSignature', 'Client certificates should have digital signature usage'),
        ('keyEncipherment', 'Client certificates should have key encipherment usage'),
        ('email:', 'Client certificates should have email SAN')
    ]

    for check, message in client_checks:
        assert check in client_section, f"Missing client certificate configuration: {message}"

    # Security check: Client certificates should NOT have serverAuth (Issue #153)
    # Look for serverAuth in extended_key_usage section, not in comments
    eku_lines = [line for line in client_section.split('\n') if 'extended_key_usage:' in line or (line.strip().startswith('- ') and 'serverAuth' in line)]
    has_server_auth = any('serverAuth' in line for line in eku_lines if line.strip().startswith('- '))
    assert not has_server_auth, "Client certificates must NOT have serverAuth EKU to prevent server impersonation"

    # Verify client certificates use unique email domains (Issue #153)
    assert 'openssl_constraint_random_id' in client_section, "Client certificates should use unique email domain per deployment"

    print("✓ Client certificate configuration has proper EKU restrictions (no serverAuth)")

def test_client_certificates():
    """Test client certificates - uses real certs if available, else validates config (Issue #75, #153)"""
    cert_files = find_generated_certificates()
    if cert_files:
        validate_client_certificates_real(cert_files)
    else:
        validate_client_certificates_config()


def validate_pkcs12_files_real(cert_files):
    """Validate actual Ansible-generated PKCS#12 files"""
    if not cert_files.get('p12_files'):
        print("⚠ No PKCS#12 files found")
        return

    major, minor = test_openssl_version_detection()

    for p12_file in cert_files['p12_files']:
        assert os.path.exists(p12_file), f"PKCS#12 file should exist: {p12_file}"

        # Test that PKCS#12 file can be read (validates format)
        legacy_flag = ['-legacy'] if major >= 3 else []

        result = subprocess.run([
            'openssl', 'pkcs12', '-info',
            '-in', p12_file,
            '-passin', 'pass:',  # Try empty password first
            '-noout'
        ] + legacy_flag, capture_output=True, text=True)

        # PKCS#12 files should be readable (even if password-protected)
        # We're just testing format validity, not trying to extract contents
        if result.returncode != 0:
            # Try with common password patterns if empty password fails
            print(f"⚠ PKCS#12 file may require password: {os.path.basename(p12_file)}")

        print(f"✓ Real PKCS#12 file exists: {os.path.basename(p12_file)}")

def validate_pkcs12_files_config():
    """Validate PKCS#12 file configuration in Ansible files (CI mode)"""
    openssl_task_file = find_ansible_file('roles/strongswan/tasks/openssl.yml')
    if not openssl_task_file:
        print("⚠ Could not find openssl.yml task file")
        return

    with open(openssl_task_file) as f:
        content = f.read()

    # Check PKCS#12 generation configuration
    p12_checks = [
        ('openssl_pkcs12', 'PKCS#12 generation should be configured'),
        ('encryption_level', 'PKCS#12 encryption level should be configured'),
        ('compatibility2022', 'PKCS#12 should use Apple-compatible encryption'),
        ('friendly_name', 'PKCS#12 should have friendly names'),
        ('other_certificates', 'PKCS#12 should include CA certificate for full chain'),
        ('passphrase', 'PKCS#12 files should be password protected'),
        ('mode: "0600"', 'PKCS#12 files should have secure permissions')
    ]

    for check, message in p12_checks:
        assert check in content, f"Missing PKCS#12 configuration: {message}"

    print("✓ PKCS#12 configuration has proper Apple device compatibility settings")

def test_pkcs12_files():
    """Test PKCS#12 files - uses real files if available, else validates config (Issue #14755, #14718)"""
    cert_files = find_generated_certificates()
    if cert_files:
        validate_pkcs12_files_real(cert_files)
    else:
        validate_pkcs12_files_config()


def validate_certificate_chain_real(cert_files):
    """Validate actual Ansible-generated certificate chain"""
    # Load CA certificate
    with open(cert_files['ca_cert'], 'rb') as f:
        ca_cert_data = f.read()
    ca_certificate = x509.load_pem_x509_certificate(ca_cert_data)

    # Test that all other certificates are signed by the CA
    other_certs = [f for f in cert_files['server_certs'] if f != cert_files['ca_cert']]

    if not other_certs:
        print("⚠ No client/server certificates found to validate")
        return

    for cert_path in other_certs:
        with open(cert_path, 'rb') as f:
            cert_data = f.read()
        certificate = x509.load_pem_x509_certificate(cert_data)

        # Verify the certificate was signed by our CA
        assert certificate.issuer == ca_certificate.subject, f"Certificate {cert_path} not signed by CA"

        # Verify certificate is currently valid (not expired)
        from datetime import datetime
        now = datetime.now(UTC)
        assert certificate.not_valid_before_utc <= now, f"Certificate {cert_path} not yet valid"
        assert certificate.not_valid_after_utc >= now, f"Certificate {cert_path} has expired"

        print(f"✓ Real certificate chain valid: {os.path.basename(cert_path)}")

    print("✓ All real certificates properly signed by CA")

def validate_certificate_chain_config():
    """Validate certificate chain configuration in Ansible files (CI mode)"""
    openssl_task_file = find_ansible_file('roles/strongswan/tasks/openssl.yml')
    if not openssl_task_file:
        print("⚠ Could not find openssl.yml task file")
        return

    with open(openssl_task_file) as f:
        content = f.read()

    # Check certificate signing configuration
    chain_checks = [
        ('provider: ownca', 'Certificates should be signed by own CA'),
        ('ownca_path', 'CA certificate path should be specified'),
        ('ownca_privatekey_path', 'CA private key path should be specified'),
        ('ownca_privatekey_passphrase', 'CA private key should be password protected'),
        ('certificate_validity_days: 3650', 'Certificate validity should be configurable (default 10 years)'),
        ('ownca_not_after: "+{{ certificate_validity_days }}d"', 'Certificates should use configurable validity period'),
        ('ownca_not_before: "-1d"', 'Certificates should have backdated start time'),
        ('curve: secp384r1', 'Should use strong elliptic curve cryptography'),
        ('type: ECC', 'Should use elliptic curve keys for better security')
    ]

    for check, message in chain_checks:
        assert check in content, f"Missing certificate chain configuration: {message}"

    print("✓ Certificate chain configuration properly set up for CA signing")

def test_certificate_chain():
    """Test certificate chain - uses real certs if available, else validates config"""
    cert_files = find_generated_certificates()
    if cert_files:
        validate_certificate_chain_real(cert_files)
    else:
        validate_certificate_chain_config()


def find_ansible_file(relative_path):
    """Find Ansible file from various possible locations"""
    # Try different base paths
    possible_bases = [
        ".",  # Current directory
        "..",  # Parent directory (from tests/unit)
        "../..",  # Grandparent (from tests/unit to project root)
        "../../..",  # Alternative deep path
    ]

    for base in possible_bases:
        full_path = os.path.join(base, relative_path)
        if os.path.exists(full_path):
            return full_path

    return None

if __name__ == "__main__":
    tests = [
        test_openssl_version_detection,
        test_ca_certificate,
        test_server_certificates,
        test_client_certificates,
        test_pkcs12_files,
        test_certificate_chain,
    ]

    failed = 0
    for test in tests:
        try:
            test()
        except AssertionError as e:
            print(f"✗ {test.__name__} failed: {e}")
            failed += 1
        except Exception as e:
            print(f"✗ {test.__name__} error: {e}")
            failed += 1

    if failed > 0:
        print(f"\n{failed} tests failed")
        sys.exit(1)
    else:
        print(f"\nAll {len(tests)} tests passed!")