#!/usr/bin/env python3
"""
Enhance VM YAML files with NTP, security hardening, and enhanced final_message
"""

import os
import sys
import yaml
import re
from pathlib import Path

def enhance_vm_yaml(file_path):
    """Enhance a single VM YAML file with all improvements"""
    
    with open(file_path, 'r') as f:
        content = f.read()
    
    # Check if already enhanced
    if 'chrony' in content and 'unattended-upgrades' in content and 'SSH hardening' in content:
        print(f"  ⚠️  {os.path.basename(file_path)} already enhanced, skipping")
        return False
    
    # Create backup
    backup_path = f"{file_path}.backup"
    with open(backup_path, 'w') as f:
        f.write(content)
    
    # Parse YAML
    try:
        data = yaml.safe_load(content)
    except yaml.YAMLError as e:
        print(f"  ❌ Error parsing {file_path}: {e}")
        return False
    
    # Get userData
    user_data = data['spec']['forProvider']['userData']
    
    # Add chrony and unattended-upgrades to packages if not present
    if 'chrony' not in user_data:
        # Find packages section and add new packages
        packages_pattern = r'packages:\s*\n((?:\s+- .+\n?)+)'
        match = re.search(packages_pattern, user_data)
        if match:
            packages_block = match.group(1)
            if 'chrony' not in packages_block:
                # Add after lsb-release
                user_data = re.sub(
                    r'(\s+- lsb-release\n)',
                    r'\1        - chrony\n        - unattended-upgrades\n        - apt-listchanges\n',
                    user_data
                )
    
    # Add NTP configuration if not present
    if 'ntp:' not in user_data:
        ntp_config = '''
      # Time synchronization (NTP)
      ntp:
        enabled: true
        ntp_client: chrony
        servers:
          - 0.pool.ntp.org
          - 1.pool.ntp.org
          - 2.pool.ntp.org
          - 3.pool.ntp.org
      
'''
        # Insert after package_upgrade
        user_data = re.sub(
            r'(package_upgrade: true\n)',
            r'\1' + ntp_config,
            user_data
        )
    
    # Update package verification to include new packages
    user_data = re.sub(
        r'for pkg in qemu-guest-agent curl wget net-tools(?: [^;]+)?;',
        r'for pkg in qemu-guest-agent curl wget net-tools chrony unattended-upgrades',
        user_data
    )
    
    # Add security updates, NTP, and SSH hardening to runcmd if not present
    if 'Configure automatic security updates' not in user_data:
        security_config = '''
        # Configure automatic security updates
        - |
          echo "Configuring automatic security updates..."
          cat > /etc/apt/apt.conf.d/50unattended-upgrades <<'EOF'
          Unattended-Upgrade::Allowed-Origins {
              "${distro_id}:${distro_codename}-security";
              "${distro_id}ESMApps:${distro_codename}-apps-security";
              "${distro_id}ESM:${distro_codename}-infra-security";
          };
          Unattended-Upgrade::AutoFixInterruptedDpkg "true";
          Unattended-Upgrade::MinimalSteps "true";
          Unattended-Upgrade::Remove-Unused-Kernel-Packages "true";
          Unattended-Upgrade::Remove-Unused-Dependencies "true";
          Unattended-Upgrade::Automatic-Reboot "false";
          Unattended-Upgrade::Automatic-Reboot-Time "02:00";
          EOF
          systemctl enable unattended-upgrades
          systemctl start unattended-upgrades
          echo "Automatic security updates configured"
        
        # Configure NTP (Chrony)
        - |
          echo "Configuring NTP (Chrony)..."
          systemctl enable chrony
          systemctl restart chrony
          sleep 3
          if systemctl is-active --quiet chrony; then
            echo "NTP (Chrony) is running"
            chronyc tracking | head -1 || true
          else
            echo "WARNING: NTP (Chrony) may not be running"
          fi
        
        # SSH hardening
        - |
          echo "Hardening SSH configuration..."
          if ! grep -q "^PermitRootLogin no" /etc/ssh/sshd_config; then
            sed -i 's/^#PermitRootLogin.*/PermitRootLogin no/' /etc/ssh/sshd_config
            sed -i 's/^PermitRootLogin yes/PermitRootLogin no/' /etc/ssh/sshd_config
          fi
          if ! grep -q "^PasswordAuthentication no" /etc/ssh/sshd_config; then
            sed -i 's/^#PasswordAuthentication.*/PasswordAuthentication no/' /etc/ssh/sshd_config
            sed -i 's/^PasswordAuthentication yes/PasswordAuthentication no/' /etc/ssh/sshd_config
          fi
          if ! grep -q "^PubkeyAuthentication yes" /etc/ssh/sshd_config; then
            sed -i 's/^#PubkeyAuthentication.*/PubkeyAuthentication yes/' /etc/ssh/sshd_config
          fi
          systemctl restart sshd
          echo "SSH hardening completed"
'''
        
        # Insert before final_message or at end of runcmd
        if 'final_message:' in user_data:
            user_data = re.sub(
                r'(\n      # Final message)',
                security_config + r'\1',
                user_data
            )
        else:
            # Add at end of runcmd
            user_data = re.sub(
                r'(systemctl status qemu-guest-agent --no-pager \|\| true\n)',
                r'\1' + security_config,
                user_data
            )
    
    # Add write_files section if not present
    if 'write_files:' not in user_data:
        write_files = '''
      # Write files for security configuration
      write_files:
        - path: /etc/apt/apt.conf.d/20auto-upgrades
          content: |
            APT::Periodic::Update-Package-Lists "1";
            APT::Periodic::Download-Upgradeable-Packages "1";
            APT::Periodic::AutocleanInterval "7";
            APT::Periodic::Unattended-Upgrade "1";
          permissions: '0644'
          owner: root:root
      
'''
        if 'final_message:' in user_data:
            user_data = re.sub(
                r'(\n      # Final message)',
                write_files + r'\1',
                user_data
            )
        else:
            user_data += write_files
    
    # Enhance final_message
    enhanced_final = '''      # Final message
      final_message: |
        ==========================================
        System Boot Completed Successfully!
        ==========================================
        
        Services Status:
        - QEMU Guest Agent: $(systemctl is-active qemu-guest-agent)
        - NTP (Chrony): $(systemctl is-active chrony)
        - Automatic Security Updates: $(systemctl is-active unattended-upgrades)
        
        System Information:
        - Hostname: $(hostname)
        - IP Address: $(hostname -I | awk '{print $1}')
        - Time: $(date)
        
        Packages Installed:
        - qemu-guest-agent, curl, wget, net-tools
        - chrony (NTP), unattended-upgrades (Security)
        
        Security Configuration:
        - SSH: Root login disabled, Password auth disabled
        - Automatic security updates: Enabled
        - NTP synchronization: Enabled
        
        Next Steps:
        1. Verify all services are running
        2. Check cloud-init logs: /var/log/cloud-init-output.log
        3. Test SSH access
        =========================================='''
    
    # Replace existing final_message
    if 'final_message:' in user_data:
        user_data = re.sub(
            r'      # Final message.*?(?=\n  providerConfigRef|\Z)',
            enhanced_final,
            user_data,
            flags=re.DOTALL
        )
    else:
        # Add final_message before providerConfigRef
        user_data = re.sub(
            r'(\n  providerConfigRef:)',
            '\n' + enhanced_final + r'\1',
            user_data
        )
    
    # Update data structure
    data['spec']['forProvider']['userData'] = user_data
    
    # Write back
    with open(file_path, 'w') as f:
        yaml.dump(data, f, default_flow_style=False, sort_keys=False, allow_unicode=True)
    
    return True

def main():
    if len(sys.argv) < 2:
        print("Usage: python3 enhance-vm-yaml.py <yaml_file> [<yaml_file> ...]")
        sys.exit(1)
    
    files_enhanced = 0
    files_skipped = 0
    
    for file_path in sys.argv[1:]:
        if not os.path.exists(file_path):
            print(f"❌ File not found: {file_path}")
            continue
        
        print(f"Processing {os.path.basename(file_path)}...")
        if enhance_vm_yaml(file_path):
            files_enhanced += 1
            print(f"  ✅ Enhanced")
        else:
            files_skipped += 1
    
    print(f"\n==========================================")
    print(f"Enhanced: {files_enhanced} files")
    print(f"Skipped: {files_skipped} files")
    print(f"==========================================")

if __name__ == '__main__':
    main()