#!/usr/bin/env python3
"""
Test script for Secure Boot integration in Debian GRUB Stage

This script validates Secure Boot setup, including:
- Secure Boot key management
- GRUB EFI signing
- Certificate validation
- Secure Boot policy configuration

Author: Debian bootc-image-builder team
"""

import os
import sys
import tempfile
import shutil
import subprocess
import logging
from typing import Dict, List, Optional, Any

# Add the parent directory to the path to import the GRUB stage
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from debian_grub_stage import DebianGrubStage

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(levelname)s:%(name)s:%(message)s')
logger = logging.getLogger(__name__)


class SecureBootTest:
    """
    Test suite for Secure Boot integration.
    """
    
    def __init__(self):
        self.test_results = {}
        self.temp_dirs = []
    
    def run_all_tests(self) -> bool:
        """Run all Secure Boot tests."""
        logger.info("Starting Secure Boot integration tests")
        
        tests = [
            ("test_secure_boot_key_management", self.test_secure_boot_key_management),
            ("test_grub_efi_signing", self.test_grub_efi_signing),
            ("test_certificate_validation", self.test_certificate_validation),
            ("test_secure_boot_policy", self.test_secure_boot_policy),
            ("test_signing_tool_integration", self.test_signing_tool_integration),
            ("test_secure_boot_verification", self.test_secure_boot_verification),
        ]
        
        all_passed = True
        for test_name, test_func in tests:
            logger.info(f"Running test: {test_name}")
            try:
                result = test_func()
                self.test_results[test_name] = result
                if result:
                    logger.info(f"✓ {test_name} PASSED")
                else:
                    logger.error(f"✗ {test_name} FAILED")
                    all_passed = False
            except Exception as e:
                logger.error(f"✗ {test_name} ERROR: {e}")
                self.test_results[test_name] = False
                all_passed = False
        
        self.print_summary()
        return all_passed
    
    def print_summary(self) -> None:
        """Print test summary."""
        logger.info("=" * 50)
        logger.info("SECURE BOOT INTEGRATION TEST SUMMARY")
        logger.info("=" * 50)
        
        passed = sum(1 for result in self.test_results.values() if result)
        total = len(self.test_results)
        
        for test_name, result in self.test_results.items():
            status = "PASSED" if result else "FAILED"
            logger.info(f"{test_name}: {status}")
        
        logger.info(f"Overall: {passed}/{total} tests passed")
        logger.info("=" * 50)
        
        if passed == total:
            logger.info("All Secure Boot integration tests PASSED")
        else:
            logger.error("Some Secure Boot integration tests FAILED")
    
    def cleanup(self) -> None:
        """Clean up temporary directories."""
        for temp_dir in self.temp_dirs:
            try:
                shutil.rmtree(temp_dir)
            except Exception as e:
                logger.warning(f"Failed to clean up {temp_dir}: {e}")
    
    def test_secure_boot_key_management(self) -> bool:
        """Test Secure Boot key management."""
        test_dir = tempfile.mkdtemp()
        self.temp_dirs.append(test_dir)
        
        try:
            # Create mock context
            context = MockOsbuildContext(test_dir)
            
            # Set up directory structure
            os.makedirs(f"{test_dir}/etc/secureboot", exist_ok=True)
            os.makedirs(f"{test_dir}/boot/efi/EFI/debian", exist_ok=True)
            os.makedirs(f"{test_dir}/boot/grub", exist_ok=True)
            os.makedirs(f"{test_dir}/etc/default", exist_ok=True)
            
            # Create mock Secure Boot keys
            db_key = f"{test_dir}/etc/secureboot/db.key"
            db_cert = f"{test_dir}/etc/secureboot/db.crt"
            kek_key = f"{test_dir}/etc/secureboot/kek.key"
            kek_cert = f"{test_dir}/etc/secureboot/kek.crt"
            
            for key_file in [db_key, db_cert, kek_key, kek_cert]:
                with open(key_file, 'w') as f:
                    f.write(f"mock {os.path.basename(key_file)}")
            
            # Initialize GRUB stage with Secure Boot enabled
            options = {
                'ostree_integration': True,
                'uefi': True,
                'secure_boot': True,
                'timeout': 5,
                'default_entry': 0,
                'secure_boot_keys': {
                    'db_key': db_key,
                    'db_cert': db_cert,
                    'kek_key': kek_key,
                    'kek_cert': kek_cert
                }
            }
            grub_stage = DebianGrubStage(options)
            
            # Test Secure Boot key management
            grub_stage._configure_secure_boot(context)
            
            # Verify Secure Boot keys were processed
            if not hasattr(context, 'last_command'):
                logger.error("Secure Boot key management command not executed")
                return False
            
            # Verify key files exist
            for key_file in [db_key, db_cert, kek_key, kek_cert]:
                if not os.path.exists(key_file):
                    logger.error(f"Secure Boot key file not found: {key_file}")
                    return False
            
            logger.info("Secure Boot key management successful")
            return True
            
        except Exception as e:
            logger.error(f"Secure Boot key management test failed: {e}")
            return False
    
    def test_grub_efi_signing(self) -> bool:
        """Test GRUB EFI signing process."""
        test_dir = tempfile.mkdtemp()
        self.temp_dirs.append(test_dir)
        
        try:
            # Create mock context
            context = MockOsbuildContext(test_dir)
            
            # Set up directory structure
            os.makedirs(f"{test_dir}/boot/efi/EFI/debian", exist_ok=True)
            os.makedirs(f"{test_dir}/boot/grub", exist_ok=True)
            os.makedirs(f"{test_dir}/etc/default", exist_ok=True)
            
            # Create mock GRUB EFI file
            grub_efi = f"{test_dir}/boot/efi/EFI/debian/grubx64.efi"
            with open(grub_efi, 'w') as f:
                f.write("mock grub efi")
            
            # Create mock signing key
            signing_key = f"{test_dir}/etc/secure-boot/keys/db.key"
            os.makedirs(os.path.dirname(signing_key), exist_ok=True)
            with open(signing_key, 'w') as f:
                f.write("mock signing key")
            
            # Initialize GRUB stage with Secure Boot enabled
            options = {
                'ostree_integration': True,
                'uefi': True,
                'secure_boot': True,
                'timeout': 5,
                'default_entry': 0,
                'signing_key': signing_key
            }
            grub_stage = DebianGrubStage(options)
            
            # Test GRUB EFI signing
            grub_stage._configure_secure_boot(context)
            
            # Verify signing command was executed
            if not hasattr(context, 'last_command'):
                logger.error("GRUB EFI signing command not executed")
                return False
            
            # Verify signing tool was used
            if 'sbsign' not in context.last_command and 'pesign' not in context.last_command:
                logger.error("Secure Boot signing tool not used")
                return False
            
            # Verify signed file was created
            signed_efi = f"{test_dir}/boot/efi/EFI/debian/grubx64.efi.signed"
            if not os.path.exists(signed_efi):
                logger.error(f"Signed GRUB EFI file not created: {signed_efi}")
                return False
            
            logger.info("GRUB EFI signing successful")
            return True
            
        except Exception as e:
            logger.error(f"GRUB EFI signing test failed: {e}")
            return False
    
    def test_certificate_validation(self) -> bool:
        """Test certificate validation process."""
        test_dir = tempfile.mkdtemp()
        self.temp_dirs.append(test_dir)
        
        try:
            # Create mock context
            context = MockOsbuildContext(test_dir)
            
            # Set up directory structure
            os.makedirs(f"{test_dir}/etc/secureboot", exist_ok=True)
            os.makedirs(f"{test_dir}/boot/efi/EFI/debian", exist_ok=True)
            os.makedirs(f"{test_dir}/boot/grub", exist_ok=True)
            os.makedirs(f"{test_dir}/etc/default", exist_ok=True)
            
            # Create mock certificates
            db_cert = f"{test_dir}/etc/secureboot/db.crt"
            kek_cert = f"{test_dir}/etc/secureboot/kek.crt"
            pk_cert = f"{test_dir}/etc/secureboot/pk.crt"
            
            for cert_file in [db_cert, kek_cert, pk_cert]:
                with open(cert_file, 'w') as f:
                    f.write(f"mock {os.path.basename(cert_file)}")
            
            # Initialize GRUB stage with Secure Boot enabled
            options = {
                'ostree_integration': True,
                'uefi': True,
                'secure_boot': True,
                'timeout': 5,
                'default_entry': 0,
                'certificates': {
                    'db': db_cert,
                    'kek': kek_cert,
                    'pk': pk_cert
                }
            }
            grub_stage = DebianGrubStage(options)
            
            # Test certificate validation
            grub_stage._configure_secure_boot(context)
            
            # Verify certificate validation command was executed
            if not hasattr(context, 'last_command'):
                logger.error("Certificate validation command not executed")
                return False
            
            # Verify certificate files exist
            for cert_file in [db_cert, kek_cert, pk_cert]:
                if not os.path.exists(cert_file):
                    logger.error(f"Certificate file not found: {cert_file}")
                    return False
            
            logger.info("Certificate validation successful")
            return True
            
        except Exception as e:
            logger.error(f"Certificate validation test failed: {e}")
            return False
    
    def test_secure_boot_policy(self) -> bool:
        """Test Secure Boot policy configuration."""
        test_dir = tempfile.mkdtemp()
        self.temp_dirs.append(test_dir)
        
        try:
            # Create mock context
            context = MockOsbuildContext(test_dir)
            
            # Set up directory structure
            os.makedirs(f"{test_dir}/etc/secureboot", exist_ok=True)
            os.makedirs(f"{test_dir}/boot/efi/EFI/debian", exist_ok=True)
            os.makedirs(f"{test_dir}/boot/grub", exist_ok=True)
            os.makedirs(f"{test_dir}/etc/default", exist_ok=True)
            
            # Create mock Secure Boot policy
            policy_file = f"{test_dir}/etc/secureboot/policy.conf"
            with open(policy_file, 'w') as f:
                f.write("SECURE_BOOT_ENABLED=1\n")
                f.write("SIGNING_REQUIRED=1\n")
                f.write("CERTIFICATE_VALIDATION=1\n")
            
            # Initialize GRUB stage with Secure Boot enabled
            options = {
                'ostree_integration': True,
                'uefi': True,
                'secure_boot': True,
                'timeout': 5,
                'default_entry': 0,
                'secure_boot_policy': policy_file
            }
            grub_stage = DebianGrubStage(options)
            
            # Test Secure Boot policy configuration
            grub_stage._configure_secure_boot(context)
            
            # Verify policy file exists
            if not os.path.exists(policy_file):
                logger.error(f"Secure Boot policy file not found: {policy_file}")
                return False
            
            # Verify policy content
            with open(policy_file, 'r') as f:
                content = f.read()
                if 'SECURE_BOOT_ENABLED=1' not in content:
                    logger.error("Secure Boot policy not configured correctly")
                    return False
            
            logger.info("Secure Boot policy configuration successful")
            return True
            
        except Exception as e:
            logger.error(f"Secure Boot policy test failed: {e}")
            return False
    
    def test_signing_tool_integration(self) -> bool:
        """Test signing tool integration."""
        test_dir = tempfile.mkdtemp()
        self.temp_dirs.append(test_dir)
        
        try:
            # Create mock context
            context = MockOsbuildContext(test_dir)
            
            # Set up directory structure
            os.makedirs(f"{test_dir}/boot/efi/EFI/debian", exist_ok=True)
            os.makedirs(f"{test_dir}/boot/grub", exist_ok=True)
            os.makedirs(f"{test_dir}/etc/default", exist_ok=True)
            
            # Create mock files to sign
            grub_efi = f"{test_dir}/boot/efi/EFI/debian/grubx64.efi"
            with open(grub_efi, 'w') as f:
                f.write("mock grub efi")
            
            # Initialize GRUB stage with Secure Boot enabled
            options = {
                'ostree_integration': True,
                'uefi': True,
                'secure_boot': True,
                'timeout': 5,
                'default_entry': 0
            }
            grub_stage = DebianGrubStage(options)
            
            # Test signing tool integration
            grub_stage._configure_secure_boot(context)
            
            # Verify signing tool was used
            if not hasattr(context, 'last_command'):
                logger.error("Signing tool integration command not executed")
                return False
            
            # Verify appropriate signing tool was selected
            if 'sbsign' in context.last_command or 'pesign' in context.last_command:
                logger.info("Signing tool integration successful")
                return True
            else:
                logger.error("No signing tool was used")
                return False
            
        except Exception as e:
            logger.error(f"Signing tool integration test failed: {e}")
            return False
    
    def test_secure_boot_verification(self) -> bool:
        """Test Secure Boot verification process."""
        test_dir = tempfile.mkdtemp()
        self.temp_dirs.append(test_dir)
        
        try:
            # Create mock context
            context = MockOsbuildContext(test_dir)
            
            # Set up directory structure
            os.makedirs(f"{test_dir}/boot/efi/EFI/debian", exist_ok=True)
            os.makedirs(f"{test_dir}/boot/grub", exist_ok=True)
            os.makedirs(f"{test_dir}/etc/default", exist_ok=True)
            
            # Create mock GRUB EFI file
            grub_efi = f"{test_dir}/boot/efi/EFI/debian/grubx64.efi"
            with open(grub_efi, 'w') as f:
                f.write("mock grub efi")
            
            # Initialize GRUB stage with Secure Boot enabled
            options = {
                'ostree_integration': True,
                'uefi': True,
                'secure_boot': True,
                'timeout': 5,
                'default_entry': 0
            }
            grub_stage = DebianGrubStage(options)
            
            # Test Secure Boot verification by calling the configuration method
            # which will trigger the signing process
            grub_stage._configure_secure_boot(context)
            
            # Verify signing command was executed
            if not hasattr(context, 'last_command'):
                logger.error("Secure Boot signing command not executed")
                return False
            
            # Verify signing tool was used
            if 'sbsign' not in context.last_command and 'pesign' not in context.last_command:
                logger.error("Secure Boot signing tool not used")
                return False
            
            # Verify signed file was created
            signed_efi = f"{test_dir}/boot/efi/EFI/debian/grubx64.efi.signed"
            if not os.path.exists(signed_efi):
                logger.error(f"Signed file not found: {signed_efi}")
                return False
            
            logger.info("Secure Boot verification successful")
            return True
            
        except Exception as e:
            logger.error(f"Secure Boot verification test failed: {e}")
            return False


class MockOsbuildContext:
    """Mock osbuild context for testing."""
    
    def __init__(self, root_path: str):
        self.root_path = root_path
        self.root = root_path
        self.last_command = None
        self.simulate_failure = False
    
    def run(self, cmd: List[str], **kwargs) -> subprocess.CompletedProcess:
        """Mock run method that simulates command execution."""
        self.last_command = ' '.join(cmd)
        
        if self.simulate_failure:
            raise subprocess.CalledProcessError(1, cmd, "Mock failure")
        
        # Handle specific commands
        if 'sbsign' in cmd:
            logger.info(f"Mock Secure Boot signing: {cmd}")
            # Create signed file
            if 'grubx64.efi' in ' '.join(cmd):
                signed_file = f"{self.root_path}/boot/efi/EFI/debian/grubx64.efi.signed"
                with open(signed_file, 'w') as f:
                    f.write("mock signed grub efi")
        
        elif 'pesign' in cmd:
            logger.info(f"Mock PE signing: {cmd}")
            # Create signed file
            if 'grubx64.efi' in ' '.join(cmd):
                signed_file = f"{self.root_path}/boot/efi/EFI/debian/grubx64.efi.signed"
                with open(signed_file, 'w') as f:
                    f.write("mock signed grub efi")
        
        elif 'sbverify' in cmd:
            logger.info(f"Mock Secure Boot verification: {cmd}")
            # Simulate verification success
            if 'grubx64.efi.signed' in ' '.join(cmd):
                logger.info("Mock Secure Boot verification successful")
        
        elif 'openssl' in cmd:
            logger.info(f"Mock OpenSSL operation: {cmd}")
            # Simulate certificate/key operations
            if 'x509' in ' '.join(cmd):
                logger.info("Mock certificate validation successful")
        
        return subprocess.CompletedProcess(cmd, 0, "", "")


def main():
    """Main test execution function."""
    test_suite = SecureBootTest()
    
    try:
        success = test_suite.run_all_tests()
        return 0 if success else 1
    finally:
        test_suite.cleanup()


if __name__ == "__main__":
    sys.exit(main())
