OpenCore_appleid_patcher.py

OpenCore_appleid_patcher.py

#!/usr/bin/env python3
import plistlib
import base64
import os
import sys
import subprocess
import re
import time
import argparse
import threading
from datetime import datetime

# ANSI color codes
COLORS = {
    'RESET': '\033[0m',
    'BLACK': '\033[30m',
    'RED': '\033[31m',
    'GREEN': '\033[32m',
    'YELLOW': '\033[33m',
    'BLUE': '\033[34m',
    'MAGENTA': '\033[35m',
    'CYAN': '\033[36m',
    'WHITE': '\033[37m',
    'BOLD': '\033[1m',
    'UNDERLINE': '\033[4m',
    'BG_GREEN': '\033[42m',
    'BG_BLUE': '\033[44m'
}

# Spinner animation
class Spinner:
    def __init__(self, message="Processing", delay=0.1):
        self.spinner_chars = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"]
        self.delay = delay
        self.message = message
        self.running = False
        self.spinner_thread = None

    def spin(self):
        i = 0
        while self.running:
            sys.stdout.write(f"\r{COLORS['CYAN']}{self.spinner_chars[i]} {self.message}{COLORS['RESET']} ")
            sys.stdout.flush()
            time.sleep(self.delay)
            i = (i + 1) % len(self.spinner_chars)

    def start(self, message=None):
        if message:
            self.message = message
        self.running = True
        self.spinner_thread = threading.Thread(target=self.spin)
        self.spinner_thread.daemon = True
        self.spinner_thread.start()

    def stop(self, message=None):
        self.running = False
        if self.spinner_thread:
            self.spinner_thread.join()
        # FIX: Corrected string multiplication syntax
        sys.stdout.write('\r' + ' ' * (len(self.message) + 10) + '\r')
        if message:
            print(message)

# Progress bar
def progress_bar(iteration, total, prefix='', suffix='', length=50, fill='█', empty='░'):
    percent = 100 * (iteration / float(total))
    filled_length = int(length * iteration // total)
    bar = fill * filled_length + empty * (length - filled_length)
    bar_colored = f"{COLORS['GREEN']}{bar}{COLORS['RESET']}"
    sys.stdout.write(f'\r{prefix} |{bar_colored}| {percent:.1f}% {suffix}')
    sys.stdout.flush()
    if iteration == total:
        sys.stdout.write('\n')

# Logger function
def log(message, level="INFO", timestamp=True):
    level_colors = {
        "INFO": COLORS['BLUE'],
        "ERROR": COLORS['RED'],
        "SUCCESS": COLORS['GREEN'],
        "WARNING": COLORS['YELLOW'],
        "TITLE": COLORS['MAGENTA'] + COLORS['BOLD'],
        "HEADER": COLORS['CYAN'] + COLORS['BOLD'],
    }
    
    color = level_colors.get(level, COLORS['RESET'])
    time_str = f"[{datetime.now().strftime('%H:%M:%S')}] " if timestamp else ""
    level_str = f"[{level}] " if level != "TITLE" and level != "HEADER" else ""
    
    if level == "TITLE":
        print("\n" + "="*70)
        print(f"{color}{message.center(70)}{COLORS['RESET']}")
        print("="*70)
    elif level == "HEADER":
        print(f"\n{color}=== {message} ==={COLORS['RESET']}")
    else:
        print(f"{color}{time_str}{level_str}{message}{COLORS['RESET']}")

# Banner for script start
def print_banner():
    banner = """
    ╔════════════════════════════════════════════════════════╗
    ║                                                        ║
    ║   ███████╗ ██████╗ ███╗   ██╗ ██████╗ ███╗   ███╗ █████╗  ║
    ║   ██╔════╝██╔═══██╗████╗  ██║██╔═══██╗████╗ ████║██╔══██╗ ║
    ║   ███████╗██║   ██║██╔██╗ ██║██║   ██║██╔████╔██║███████║ ║
    ║   ╚════██║██║   ██║██║╚██╗██║██║   ██║██║╚██╔╝██║██╔══██║ ║
    ║   ███████║╚██████╔╝██║ ╚████║╚██████╔╝██║ ╚═╝ ██║██║  ██║ ║
    ║   ╚══════╝ ╚═════╝ ╚═╝  ╚═══╝ ╚═════╝ ╚═╝     ╚═╝╚═╝  ╚═╝ ║
    ║                                                        ║
    ║   VM Bluetooth Enabler Patch Tool v1.5                ║
    ║   For MacOS Sonoma OpenCore                           ║
    ║   (Fixed version)                                     ║
    ║                                                        ║
    ╚════════════════════════════════════════════════════════╝
    """
    print(f"{COLORS['CYAN']}{banner}{COLORS['RESET']}")

def restart_system():
    """Restart the system."""
    log("Initiating system restart...", "INFO")
    spinner = Spinner("Preparing to restart")
    spinner.start()
    
    # Countdown animation
    for i in range(5, 0, -1):
        spinner.stop(f"{COLORS['YELLOW']}System will restart in {i} seconds... Press Ctrl+C to cancel{COLORS['RESET']}")
        try:
            time.sleep(1)
        except KeyboardInterrupt:
            print("")  # Add a newline after the ^C
            log("Restart cancelled by user.", "WARNING")
            return False
    
    log("Restarting now...", "INFO")
    try:
        subprocess.run(['shutdown', '-r', 'now'], check=True)
        return True
    except Exception as e:
        log(f"Error restarting system: {e}", "ERROR")
        log("Please restart your system manually to apply changes.", "WARNING")
        return False

def get_disk_list():
    """Get a list of all disks in the system."""
    spinner = Spinner("Scanning disk list")
    spinner.start()
    try:
        result = subprocess.run(['diskutil', 'list'], capture_output=True, text=True)
        spinner.stop()
        return result.stdout
    except Exception as e:
        spinner.stop(f"Error getting disk list: {e}")
        return ""

def get_efi_partitions(disk_info):
    """Extract EFI partition information from diskutil output."""
    log("Analyzing disk information for EFI partitions...", "INFO")
    efi_partitions = []
    disk_lines = disk_info.split('\n')
    current_disk = None
    
    for line in disk_lines:
        # Track which disk we're currently examining
        if line.startswith("/dev/"):
            current_disk = line.split()[0]
        
        # Look for EFI partitions - multiple ways to identify them in different macOS versions
        if ("EFI" in line or "EF00" in line or "C12A7328-F81F-11D2-BA4B-00A0C93EC93B" in line) and current_disk:
            # Parse the partition number - handle multiple patterns
            match = re.search(r'\s+(\d+):\s+.*?(EFI|EF00|C12A7328-F81F-11D2-BA4B-00A0C93EC93B)', line)
            if match:
                partition_num = match.group(1)
                partition_id = f"{current_disk}s{partition_num}"
                efi_partitions.append(partition_id)
    
    if efi_partitions:
        log(f"Found {len(efi_partitions)} EFI partition(s): {', '.join(efi_partitions)}", "SUCCESS")
    else:
        log("No EFI partitions found", "WARNING")
    
    return efi_partitions

def check_if_mounted(partition):
    """Check if the partition is already mounted."""
    try:
        result = subprocess.run(['diskutil', 'info', partition], 
                              capture_output=True, text=True)
        
        # If the partition is mounted, the output will contain "Mounted: Yes"
        if "Mounted: Yes" in result.stdout:
            # Extract the mount point
            match = re.search(r'Mount Point:\s+(.*)', result.stdout)
            if match:
                return match.group(1).strip()
        return None
    except Exception as e:
        log(f"Error checking mount status: {e}", "ERROR")
        return None

def mount_efi(partition):
    """Mount an EFI partition and return the mount point with enhanced error handling."""
    spinner = Spinner(f"Mounting {partition}")
    spinner.start()
    
    # First check if already mounted
    mount_point = check_if_mounted(partition)
    if mount_point:
        spinner.stop(f"{COLORS['GREEN']}✓ Partition {partition} is already mounted at {mount_point}{COLORS['RESET']}")
        return mount_point
    
    # FIX: Use a lock to prevent race conditions between mount methods
    mount_lock = threading.Lock()
    mounted = False
    mount_point = None
    errors = []
    
    # Method 1: Try standard mount
    with mount_lock:
        try:
            result = subprocess.run(['diskutil', 'mount', partition], 
                                capture_output=True, text=True)
            
            # Parse mount point from output
            if result.returncode == 0:
                match = re.search(r'mounted at (.*)', result.stdout)
                if match:
                    mount_point = match.group(1).strip()
                    mounted = True
            else:
                errors.append(f"Standard mount: {result.stderr or result.stdout}")
        except Exception as e:
            errors.append(f"Standard mount exception: {str(e)}")
    
    # If standard mount succeeded, return the mount point
    if mounted and mount_point:
        spinner.stop(f"{COLORS['GREEN']}✓ Successfully mounted {partition} at {mount_point}{COLORS['RESET']}")
        return mount_point
    
    # Method 2: Try mounting by volume name
    with mount_lock:
        if not mounted:
            try:
                # Create a temporary directory for mounting
                efi_mount_dir = '/Volumes/EFI'
                if not os.path.exists(efi_mount_dir):
                    os.makedirs(efi_mount_dir, exist_ok=True)
                
                # Try to mount by volume name
                result = subprocess.run(['diskutil', 'mount', 'EFI'], 
                                    capture_output=True, text=True)
                
                if result.returncode == 0:
                    # Check if the correct partition was mounted
                    info_result = subprocess.run(['diskutil', 'info', '/Volumes/EFI'], 
                                            capture_output=True, text=True)
                    if partition in info_result.stdout:
                        mount_point = '/Volumes/EFI'
                        mounted = True
                    else:
                        # Wrong partition mounted, try to unmount it
                        subprocess.run(['diskutil', 'unmount', '/Volumes/EFI'], 
                                    capture_output=True, text=True)
                else:
                    errors.append(f"Volume name mount: {result.stderr or result.stdout}")
            except Exception as e:
                errors.append(f"Volume name mount exception: {str(e)}")
    
    # If method 2 succeeded, return the mount point
    if mounted and mount_point:
        spinner.stop(f"{COLORS['GREEN']}✓ Successfully mounted {partition} at {mount_point}{COLORS['RESET']}")
        return mount_point
    
    # Method 3: Try direct mount using mount_msdos command
    with mount_lock:
        if not mounted:
            try:
                # Ensure mount point exists
                efi_mount_dir = '/Volumes/EFI'
                if not os.path.exists(efi_mount_dir):
                    os.makedirs(efi_mount_dir, exist_ok=True)
                
                # Try mounting with mount_msdos
                result = subprocess.run(['sudo', 'mount_msdos', partition, efi_mount_dir], 
                                    capture_output=True, text=True)
                
                if result.returncode == 0 or os.path.exists(os.path.join(efi_mount_dir, 'EFI')):
                    mount_point = efi_mount_dir
                    mounted = True
                else:
                    errors.append(f"mount_msdos: {result.stderr or result.stdout}")
            except Exception as e:
                errors.append(f"mount_msdos exception: {str(e)}")
    
    # If method 3 succeeded, return the mount point
    if mounted and mount_point:
        spinner.stop(f"{COLORS['GREEN']}✓ Successfully mounted {partition} at {mount_point} using mount_msdos{COLORS['RESET']}")
        return mount_point
    
    # If all methods failed, log the errors and return None
    spinner.stop(f"{COLORS['RED']}✗ Failed to mount {partition} after trying multiple methods{COLORS['RESET']}")
    
    # Log detailed errors
    log("Mount failure details:", "ERROR")
    for i, error in enumerate(errors):
        log(f"  Method {i+1}: {error}", "ERROR")
    
    # Check for system-level reasons for mounting failures
    check_system_constraints()
    
    return None

def check_system_constraints():
    """Check for system constraints that could prevent mounting."""
    # Check if SIP is enabled (which might restrict mounting)
    try:
        sip_result = subprocess.run(['csrutil', 'status'], 
                                 capture_output=True, text=True)
        
        if "enabled" in sip_result.stdout.lower():
            log("System Integrity Protection (SIP) is enabled, which might restrict mounting operations.", "WARNING")
            log("You may need to temporarily disable SIP to mount EFI partitions. See Apple support for details.", "INFO")
    except Exception:
        pass  # Ignore if csrutil command fails
    
    # Check if running on latest macOS version
    try:
        os_version_result = subprocess.run(['sw_vers', '-productVersion'], 
                                        capture_output=True, text=True)
        
        version = os_version_result.stdout.strip()
        log(f"Running on macOS version: {version}", "INFO")
        
        if version.startswith("14.") or version.startswith("13."):  # Sonoma or Ventura
            log("Recent macOS versions have additional security measures for mounting partitions.", "INFO")
            log("Try using 'Disk Utility' GUI application to mount the EFI partition manually.", "INFO")
    except Exception:
        pass  # Ignore if sw_vers command fails

def unmount_efi(partition):
    """Unmount an EFI partition."""
    spinner = Spinner(f"Unmounting {partition}")
    spinner.start()
    
    # First check if the partition is mounted
    mount_point = check_if_mounted(partition)
    if not mount_point:
        spinner.stop(f"{COLORS['YELLOW']}⚠ Partition {partition} is not mounted{COLORS['RESET']}")
        return True
    
    # Use a lock to prevent race conditions
    unmount_lock = threading.Lock()
    
    # Try unmounting by mount point first (more reliable)
    with unmount_lock:
        try:
            result = subprocess.run(['diskutil', 'unmount', mount_point], 
                                capture_output=True, text=True)
            
            if result.returncode == 0:
                spinner.stop(f"{COLORS['GREEN']}✓ Successfully unmounted {mount_point}{COLORS['RESET']}")
                return True
        except Exception as e:
            log(f"Error unmounting by mount point: {e}", "WARNING")
    
    # If that failed, try unmounting by partition
    with unmount_lock:
        try:
            result = subprocess.run(['diskutil', 'unmount', partition], 
                                capture_output=True, text=True)
            
            if result.returncode == 0:
                spinner.stop(f"{COLORS['GREEN']}✓ Successfully unmounted {partition}{COLORS['RESET']}")
                return True
        except Exception as e:
            log(f"Error unmounting by partition: {e}", "WARNING")
    
    # If both methods failed, try force unmounting
    with unmount_lock:
        try:
            result = subprocess.run(['diskutil', 'unmount', 'force', mount_point], 
                                capture_output=True, text=True)
            
            if result.returncode == 0:
                spinner.stop(f"{COLORS['GREEN']}✓ Successfully force unmounted {mount_point}{COLORS['RESET']}")
                return True
        except Exception as e:
            log(f"Error force unmounting: {e}", "WARNING")
    
    # If we get here, all unmount attempts failed
    spinner.stop(f"{COLORS['RED']}✗ Failed to unmount {partition}{COLORS['RESET']}")
    log("Please unmount the partition manually using Disk Utility.", "WARNING")
    return False

def find_config_plist(mount_point):
    """Find OpenCore config.plist in the mounted EFI partition."""
    # Common paths for OpenCore config.plist
    possible_paths = [
        os.path.join(mount_point, 'EFI', 'OC', 'config.plist'),
        os.path.join(mount_point, 'OC', 'config.plist'),
        # Add additional possible paths
        os.path.join(mount_point, 'config.plist'),
        os.path.join(mount_point, 'EFI', 'CLOVER', 'config.plist')  # Also check for Clover
    ]
    
    spinner = Spinner(f"Searching for OpenCore config.plist")
    spinner.start()
    
    for i, path in enumerate(possible_paths):
        # Add a short pause to show progress
        time.sleep(0.2)
        progress_bar(i+1, len(possible_paths), prefix='Progress:', suffix='Complete', length=30)
        
        if os.path.exists(path):
            if 'CLOVER' in path:
                spinner.stop(f"{COLORS['YELLOW']}⚠ Found Clover config.plist at {path} - may not be compatible{COLORS['RESET']}")
            else:
                spinner.stop(f"{COLORS['GREEN']}✓ Found OpenCore config.plist at {path}{COLORS['RESET']}")
            return path
    
    # Additional deep search for config.plist in case it's in a non-standard location
    try:
        spinner.stop()
        spinner = Spinner(f"Performing deep search for config.plist")
        spinner.start()
        
        # Use find command to search for config.plist files
        find_result = subprocess.run(['find', mount_point, '-name', 'config.plist'], 
                                  capture_output=True, text=True)
        
        if find_result.returncode == 0 and find_result.stdout.strip():
            paths = find_result.stdout.strip().split('\n')
            for path in paths:
                if os.path.exists(path):
                    if 'CLOVER' in path:
                        spinner.stop(f"{COLORS['YELLOW']}⚠ Found Clover config.plist at {path} - may not be compatible{COLORS['RESET']}")
                    else:
                        spinner.stop(f"{COLORS['GREEN']}✓ Found config.plist at {path}{COLORS['RESET']}")
                    return path
    except Exception:
        pass  # Ignore if find command fails
    
    spinner.stop(f"{COLORS['YELLOW']}⚠ No OpenCore config.plist found{COLORS['RESET']}")
    return None

def check_if_patches_exist(config_path):
    """Check if the Sonoma VM BT Enabler patches already exist in the config."""
    spinner = Spinner(f"Checking for existing patches")
    spinner.start()
    try:
        with open(config_path, 'rb') as f:
            config = plistlib.load(f)
        
        if 'Kernel' in config and 'Patch' in config['Kernel']:
            for patch in config['Kernel']['Patch']:
                if isinstance(patch, dict) and 'Comment' in patch:
                    if 'Sonoma VM BT Enabler' in patch['Comment']:
                        spinner.stop(f"{COLORS['YELLOW']}⚠ Found existing patch: {patch['Comment']}{COLORS['RESET']}")
                        return True
        
        spinner.stop(f"{COLORS['BLUE']}ℹ No existing patches found{COLORS['RESET']}")
        return False
    except Exception as e:
        spinner.stop(f"{COLORS['RED']}✗ Error checking for existing patches: {e}{COLORS['RESET']}")
        log("Continuing anyway - will attempt to apply patches", "WARNING")
        return False

def add_kernel_patches(config_path):
    """Add Sonoma VM BT Enabler kernel patches to config.plist."""
    log("Starting patch process...", "INFO")
    
    # Make a backup of the original file
    backup_path = config_path + '.backup'
    
    spinner = Spinner(f"Creating backup at {backup_path}")
    spinner.start()
    
    # FIX: Use shutil instead of os.system for better error handling
    try:
        import shutil
        shutil.copy2(config_path, backup_path)
        spinner.stop(f"{COLORS['GREEN']}✓ Backup created at {backup_path}{COLORS['RESET']}")
    except Exception as e:
        spinner.stop(f"{COLORS['RED']}✗ Error creating backup: {e}{COLORS['RESET']}")
        return "error"
    
    # Read the plist file
    spinner = Spinner(f"Reading config file")
    spinner.start()
    try:
        with open(config_path, 'rb') as f:
            config = plistlib.load(f)
        spinner.stop(f"{COLORS['GREEN']}✓ Config file loaded successfully{COLORS['RESET']}")
    except Exception as e:
        spinner.stop(f"{COLORS['RED']}✗ Error reading config file: {e}{COLORS['RESET']}")
        # Print more detailed error information
        log(f"Error details: {str(e)}", "ERROR")
        log(f"Trying to determine the issue with the plist file...", "INFO")
        
        try:
            # Check if file exists and is readable
            if not os.path.exists(config_path):
                log(f"The file {config_path} does not exist!", "ERROR")
            elif not os.access(config_path, os.R_OK):
                log(f"The file {config_path} is not readable!", "ERROR")
            else:
                # Try to read the file contents
                with open(config_path, 'rb') as f:
                    content = f.read(100)  # Read first 100 bytes
                    log(f"File starts with: {content}", "INFO")
                    log("This might not be a valid plist file", "WARNING")
        except Exception as inner_e:
            log(f"Failed to diagnose file issue: {str(inner_e)}", "ERROR")
            
        return "error"
    
    # Prepare the patch entries
    log("Preparing patch data...", "INFO")
    
    # Show a fancy progress bar for "patch preparation"
    for i in range(10):
        progress_bar(i+1, 10, prefix='Preparing patches:', suffix='Complete', length=30)
        time.sleep(0.05)
    
    # FIX: Updated patches with more reliable patterns
    patch1 = {
        'Arch': 'x86_64',
        'Base': '',
        'Comment': 'Sonoma VM BT Enabler - PART 1 of 2 - Patch kern.hv_vmm_present=0',
        'Count': 1,
        'Enabled': True,
        'Find': base64.b64decode('aGliZXJuYXRlaGlkcmVhZHkAaGliZXJuYXRlY291bnQA'),
        'Identifier': 'kernel',
        'Limit': 0,
        'Mask': b'',
        'MaxKernel': '',
        'MinKernel': '20.4.0',
        'Replace': base64.b64decode('aGliZXJuYXRlaGlkcmVhZHkAaHZfdm1tX3ByZXNlbnQA'),
        'ReplaceMask': b'',
        'Skip': 0,
    }
    
    patch2 = {
        'Arch': 'x86_64',
        'Base': '',
        'Comment': 'Sonoma VM BT Enabler - PART 2 of 2 - Patch kern.hv_vmm_present=0',
        'Count': 1,
        'Enabled': True,
        'Find': base64.b64decode('Ym9vdCBzZXNzaW9uIFVVSUQAaHZfdm1tX3ByZXNlbnQA'),
        'Identifier': 'kernel',
        'Limit': 0,
        'Mask': b'',
        'MaxKernel': '',
        'MinKernel': '22.0.0',
        'Replace': base64.b64decode('Ym9vdCBzZXNzaW9uIFVVSUQAaGliZXJuYXRlY291bnQA'),
        'ReplaceMask': b'',
        'Skip': 0,
    }
    
    # Print the actual binary representation of patch values for debugging
    log("Patch 1 - Find value (hex):", "INFO")
    log(" ".join(f"{b:02x}" for b in patch1['Find']), "INFO")
    log("Patch 1 - Replace value (hex):", "INFO")
    log(" ".join(f"{b:02x}" for b in patch1['Replace']), "INFO")
    
    log("Patch 2 - Find value (hex):", "INFO")
    log(" ".join(f"{b:02x}" for b in patch2['Find']), "INFO")
    log("Patch 2 - Replace value (hex):", "INFO")
    log(" ".join(f"{b:02x}" for b in patch2['Replace']), "INFO")
    
    # Add patches to the kernel patch section
    if 'Kernel' in config and 'Patch' in config['Kernel']:
        # Check if patches already exist
        patch_exists = False
        spinner = Spinner(f"Checking for existing patches in config")
        spinner.start()
        
        for patch in config['Kernel']['Patch']:
            if isinstance(patch, dict) and 'Comment' in patch:
                if 'Sonoma VM BT Enabler' in patch['Comment']:
                    patch_exists = True
                    spinner.stop(f"{COLORS['YELLOW']}⚠ Patch already exists: {patch['Comment']}{COLORS['RESET']}")
        
        if not patch_exists:
            spinner.stop(f"{COLORS['GREEN']}✓ Config is ready for patching{COLORS['RESET']}")
            
            # Adding patches with animation
            spinner = Spinner(f"Adding patches to config")
            spinner.start()
            time.sleep(0.5)  # Add a slight delay for visual effect
            
            # FIX: Ensure Kernel->Patch is a list
            if not isinstance(config['Kernel']['Patch'], list):
                log("Warning: Kernel->Patch is not a list. Converting to list.", "WARNING")
                config['Kernel']['Patch'] = []
            
            config['Kernel']['Patch'].append(patch1)
            config['Kernel']['Patch'].append(patch2)
            
            spinner.stop(f"{COLORS['GREEN']}✓ Added both Sonoma VM BT Enabler patches to config.plist{COLORS['RESET']}")
            
            # Write the updated plist file
            spinner = Spinner(f"Writing updated config to disk")
            spinner.start()
            
            try:
                # FIX: Add more robust error handling for the write operation
                # Create a temporary file first
                temp_path = config_path + '.tmp'
                with open(temp_path, 'wb') as f:
                    plistlib.dump(config, f)
                
                # Verify it can be loaded back
                with open(temp_path, 'rb') as f:
                    test_load = plistlib.load(f)
                
                # If we get here, it's safe to replace the original
                import os
                os.replace(temp_path, config_path)
                
                spinner.stop(f"{COLORS['GREEN']}✓ Successfully updated {config_path}{COLORS['RESET']}")
                return "success"
            except Exception as e:
                spinner.stop(f"{COLORS['RED']}✗ Error writing config file: {e}{COLORS['RESET']}")
                log(f"Detailed error: {str(e)}", "ERROR")
                log(f"Attempting to restore from backup...", "INFO")
                
                try:
                    import shutil
                    shutil.copy2(backup_path, config_path)
                    log(f"Successfully restored from backup.", "SUCCESS")
                except Exception as restore_e:
                    log(f"Failed to restore from backup: {str(restore_e)}", "ERROR")
                
                return "error"
        else:
            return "already_exists"
    else:
        # Create Kernel->Patch section if it doesn't exist
        log("Kernel->Patch section not found in config.plist", "WARNING")
        spinner = Spinner(f"Creating necessary structure for patches")
        spinner.start()
        
        # Initialize Kernel section if it doesn't exist
        if 'Kernel' not in config:
            config['Kernel'] = {}
        
        # Initialize Patch array if it doesn't exist
        if 'Patch' not in config['Kernel']:
            config['Kernel']['Patch'] = []
        
        # Now add the patches
        config['Kernel']['Patch'].append(patch1)
        config['Kernel']['Patch'].append(patch2)
        
        spinner.stop(f"{COLORS['GREEN']}✓ Created Kernel->Patch section and added patches{COLORS['RESET']}")
        
        # Write the updated plist file
        spinner = Spinner(f"Writing updated config to disk")
        spinner.start()
        
        try:
            # Create a temporary file first
            temp_path = config_path + '.tmp'
            with open(temp_path, 'wb') as f:
                plistlib.dump(config, f)
            
            # Verify it can be loaded back
            with open(temp_path, 'rb') as f:
                test_load = plistlib.load(f)
            
            # If we get here, it's safe to replace the original
            import os
            os.replace(temp_path, config_path)
            
            spinner.stop(f"{COLORS['GREEN']}✓ Successfully updated {config_path}{COLORS['RESET']}")
            return "success"
        except Exception as e:
            spinner.stop(f"{COLORS['RED']}✗ Error writing config file: {e}{COLORS['RESET']}")
            log(f"Detailed error: {str(e)}", "ERROR")
            return "error"

def main(auto_confirm=False, auto_restart=False):
    print_banner()
    log("Starting automatic OpenCore config.plist patching...", "TITLE")
    
    # Get list of disks
    disk_info = get_disk_list()
    if not disk_info:
        log("Failed to get disk information.", "ERROR")
        return
    
    # Extract EFI partitions
    efi_partitions = get_efi_partitions(disk_info)
    if not efi_partitions:
        log("No EFI partitions found.", "ERROR")
        log("Checking for alternative methods to identify EFI partitions...", "INFO")
        
        # Try an alternative approach
        try:
            log("Trying to identify EFI partitions using gpt show...", "INFO")
            alt_result = subprocess.run(['sudo', 'gpt', 'show', '/dev/disk0'], 
                                     capture_output=True, text=True)
            
            if alt_result.returncode == 0:
                for line in alt_result.stdout.split('\n'):
                    if 'EFI' in line:
                        log(f"Possible EFI identified: {line}", "INFO")
                        # You could parse this to get partition info
            
            # Check fdisk too
            log("Trying to identify EFI partitions using fdisk...", "INFO")
            fdisk_result = subprocess.run(['sudo', 'fdisk', '/dev/disk0'], 
                                       capture_output=True, text=True)
            
            if fdisk_result.returncode == 0:
                log("Please check the following partitions:", "INFO")
                print(fdisk_result.stdout)
        except Exception:
            pass
        
        log("Consider mounting EFI manually with Disk Utility, then run this script with the specific path.", "INFO")
        return
    
    log(f"Scanning {len(efi_partitions)} EFI partition(s) for OpenCore configuration", "HEADER")
    config_found = False
    
    # Progress counter for partition scanning
    total_partitions = len(efi_partitions)
    
    # Check each EFI partition
    for idx, partition in enumerate(efi_partitions):
        partition_progress = f"[{idx+1}/{total_partitions}]"
        log(f"{partition_progress} Processing partition {partition}", "INFO")
        
        # Mount the EFI partition
        mount_point = mount_efi(partition)
        if not mount_point:
            log(f"{partition_progress} Failed to mount {partition}, skipping...", "WARNING")
            if idx == total_partitions - 1 and idx > 0:
                log("All mount attempts failed. Please try manually mounting with Disk Utility.", "ERROR")
                log("After mounting, run this script with the specific path to config.plist.", "INFO")
                log("Example: sudo python3 sonoma_bt_patcher.py /Volumes/EFI/EFI/OC/config.plist", "INFO")
            continue
        
        try:
            # Look for config.plist
            config_path = find_config_plist(mount_point)
            if config_path:
                # Check if patches already exist before asking user
                if check_if_patches_exist(config_path):
                    log(f"Patches already exist in {config_path}", "TITLE")
                    log("No changes needed. Your system is already patched.", "SUCCESS")
                    unmount_efi(partition)
                    sys.exit(0)  # Exit the script as no changes needed
                
                # Ask for confirmation before applying patch
                log(f"OpenCore configuration found", "TITLE")
                log(f"Path: {config_path}", "INFO")
                
                if auto_confirm:
                    log("Auto-confirm enabled. Applying patches automatically...", "INFO")
                    # Apply patches directly
                else:
                    while True:
                        response = input(f"{COLORS['CYAN']}Do you want to apply the Sonoma VM BT Enabler patch? [y/n/skip]: {COLORS['RESET']}").lower()
                        if response in ['y', 'yes']:
                            break
                        elif response in ['n', 'no']:
                            log("Operation cancelled by user.", "WARNING")
                            sys.exit(0)
                        elif response == 'skip':
                            log(f"Skipping {config_path}. Looking for other configs...", "INFO")
                            break
                        else:
                            log("Please enter 'y' for yes, 'n' for no, or 'skip' to try next partition.", "WARNING")
                    
                    if response == 'skip':
                        continue
                
                # Apply patches
                success = add_kernel_patches(config_path)
                if success == "success":
                    log(f"Successfully patched OpenCore config at {config_path}", "SUCCESS")
                    config_found = True
                elif success == "already_exists":
                    log(f"Patches already exist in {config_path}", "TITLE")
                    log("No changes needed. Your system is already patched.", "SUCCESS")
                    sys.exit(0)  # Exit the script as no changes needed
                else:
                    log(f"Failed to patch config at {config_path}", "ERROR")
            else:
                log(f"{partition_progress} No OpenCore config.plist found on {partition}", "WARNING")
                
                # List contents of mounted EFI for debugging
                try:
                    log(f"Listing contents of {mount_point} for debugging:", "INFO")
                    ls_result = subprocess.run(['ls', '-la', mount_point], 
                                           capture_output=True, text=True)
                    if ls_result.returncode == 0:
                        print(ls_result.stdout)
                    
                    # Check if there's an EFI folder
                    efi_folder = os.path.join(mount_point, 'EFI')
                    if os.path.exists(efi_folder) and os.path.isdir(efi_folder):
                        log(f"EFI folder found. Listing contents:", "INFO")
                        ls_efi_result = subprocess.run(['ls', '-la', efi_folder], 
                                                   capture_output=True, text=True)
                        if ls_efi_result.returncode == 0:
                            print(ls_efi_result.stdout)
                except Exception:
                    pass
        finally:
            # Always unmount the partition when done
            time.sleep(1)  # Brief pause before unmounting
            unmount_efi(partition)
        
        # If we found and patched a config, we can stop
        if config_found:
            break
    
    if config_found:
        log("Patching process completed successfully", "TITLE")
        log("Please reboot your system to apply the changes.", "SUCCESS")
        
        # Offer to restart if auto_restart is enabled
        if auto_restart:
            log("Auto-restart enabled. System will restart automatically.", "INFO")
            restart_system()
        else:
            # Offer restart option
            if not auto_confirm:  # Only ask if not in auto mode
                response = input(f"{COLORS['CYAN']}Would you like to restart now to apply changes? [y/n]: {COLORS['RESET']}").lower()
                if response in ['y', 'yes']:
                    restart_system()
    else:
        log("No OpenCore config.plist found on any EFI partition", "TITLE")
        log("Make sure OpenCore is properly installed.", "WARNING")
        log("If you know the location of your config.plist, try running this script with that path:", "INFO")
        log(f"  sudo python3 {sys.argv[0]} /path/to/config.plist", "INFO")
        log("You may also need to mount your EFI manually using Disk Utility.", "INFO")

if __name__ == "__main__":
    # Check if ANSI colors are supported in the terminal
    if "TERM" in os.environ and os.environ["TERM"] in ["xterm", "xterm-color", "xterm-256color", "screen", "screen-256color"]:
        pass  # Colors are supported
    else:
        # Strip out ANSI color codes for terminals that don't support them
        for key in COLORS:
            COLORS[key] = ""
    
    # Check if running as root
    if os.geteuid() != 0:
        log("This script requires administrative privileges.", "ERROR")
        log(f"Please run with sudo: {COLORS['CYAN']}sudo {sys.argv[0]}{COLORS['RESET']}", "INFO")
        sys.exit(1)
    
    parser = argparse.ArgumentParser(description="Apply Sonoma VM BT Enabler patches to OpenCore config.plist")
    parser.add_argument("config_path", nargs="?", help="Path to specific config.plist (optional)")
    parser.add_argument("--auto", "-a", action="store_true", help="Auto-confirm patches without prompting")
    parser.add_argument("--no-color", action="store_true", help="Disable colored output")
    parser.add_argument("--restart", "-r", action="store_true", help="Automatically restart after patching")
    parser.add_argument("--mount-only", "-m", action="store_true", help="Only mount EFI partitions without patching")
    parser.add_argument("--debug", "-d", action="store_true", help="Enable additional debug output")
    args = parser.parse_args()
    
    # Disable colors if requested
    if args.no_color:
        for key in COLORS:
            COLORS[key] = ""
    
    # Mount-only mode
    if args.mount_only:
        print_banner()
        log("EFI Mount Mode - Will only mount EFI partitions without patching", "TITLE")
        
        disk_info = get_disk_list()
        if not disk_info:
            log("Failed to get disk information.", "ERROR")
            sys.exit(1)
        
        efi_partitions = get_efi_partitions(disk_info)
        if not efi_partitions:
            log("No EFI partitions found.", "ERROR")
            sys.exit(1)
        
        log(f"Found {len(efi_partitions)} EFI partition(s)", "SUCCESS")
        
        for idx, partition in enumerate(efi_partitions):
            log(f"Attempting to mount {partition}...", "INFO")
            mount_point = mount_efi(partition)
            
            if mount_point:
                log(f"Successfully mounted {partition} at {mount_point}", "SUCCESS")
                log(f"When finished, unmount with: diskutil unmount {mount_point}", "INFO")
            else:
                log(f"Failed to mount {partition}", "ERROR")
        
        sys.exit(0)
    
    # If a specific config path is provided, use it directly
    if args.config_path:
        config_path = args.config_path
        if os.path.exists(config_path):
            log(f"Using provided config path: {config_path}", "HEADER")
            
            # Check if patches already exist before asking user
            if check_if_patches_exist(config_path):
                log(f"Patches already exist in {config_path}", "TITLE")
                log("No changes needed. Your system is already patched.", "SUCCESS")
                sys.exit(0)  # Exit the script as no changes needed
            
            # For specific paths, still ask for confirmation unless --auto is specified
            if not args.auto:
                response = input(f"{COLORS['CYAN']}Apply Sonoma VM BT Enabler patch to {config_path}? [y/n]: {COLORS['RESET']}").lower()
                if response not in ['y', 'yes']:
                    log("Operation cancelled.", "WARNING")
                    sys.exit(0)
                    
            success = add_kernel_patches(config_path)
            if success == "success":
                log("Patches applied successfully. Please reboot to apply changes.", "SUCCESS")
                
                # Handle auto-restart if enabled
                if args.restart:
                    log("Auto-restart enabled. System will restart automatically.", "INFO")
                    restart_system()
                elif not args.auto:  # Ask about restart only in interactive mode
                    response = input(f"{COLORS['CYAN']}Would you like to restart now to apply changes? [y/n]: {COLORS['RESET']}").lower()
                    if response in ['y', 'yes']:
                        restart_system()
            elif success == "already_exists":
                log(f"Patches already exist in {config_path}", "TITLE")
                log("No changes needed. Your system is already patched.", "SUCCESS")
                sys.exit(0)
            else:
                log("Failed to apply patches.", "ERROR")
                # Debug option: Print the contents of the plist
                if args.debug:
                    try:
                        with open(config_path, 'rb') as f:
                            test_config = plistlib.load(f)
                            log("Current plist structure:", "INFO")
                            # Print first level keys
                            for key in test_config:
                                log(f"Key: {key}, Type: {type(test_config[key])}", "INFO")
                    except Exception as e:
                        log(f"Error reading config for debug: {e}", "ERROR")
        else:
            log(f"Error: File {config_path} does not exist", "ERROR")
    else:
        # Otherwise, use the automatic EFI partition detection
        main(auto_confirm=args.auto, auto_restart=args.restart)