#!/bin/bash
# enhance-all-vm-cloudinit.sh
# Enhances all VM YAML files with NTP, security hardening, and final_message

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"
VM_DIR="$PROJECT_ROOT/examples/production"

log() {
    echo -e "\033[0;34m[$(date +'%Y-%m-%d %H:%M:%S')]\033[0m $*"
}

log_success() {
    echo -e "\033[0;32m[$(date +'%Y-%m-%d %H:%M:%S')] ✅\033[0m $*"
}

log_warning() {
    echo -e "\033[1;33m[$(date +'%Y-%m-%d %H:%M:%S')] ⚠️\033[0m $*"
}

log_error() {
    echo -e "\033[0;31m[$(date +'%Y-%m-%d %H:%M:%S')] ❌\033[0m $*"
}

# Function to enhance a single VM YAML file
enhance_vm_file() {
    local file=$1
    local vm_name=$(basename "$file" .yaml)
    
    log "Enhancing $vm_name..."
    
    # Check if file already has chrony (already enhanced)
    if grep -q "chrony" "$file" && grep -q "unattended-upgrades" "$file"; then
        log_warning "  $vm_name already enhanced, skipping"
        return 0
    fi
    
    # Create backup
    cp "$file" "${file}.backup"
    
    # Add chrony and unattended-upgrades to packages list
    if ! grep -q "chrony" "$file"; then
        sed -i '/- lsb-release$/a\        - chrony\n        - unattended-upgrades\n        - apt-listchanges' "$file"
    fi
    
    # Add NTP configuration after package_upgrade
    if ! grep -q "ntp:" "$file"; then
        sed -i '/package_upgrade: true/a\      \n      # Time synchronization (NTP)\n      ntp:\n        enabled: true\n        ntp_client: chrony\n        servers:\n          - 0.pool.ntp.org\n          - 1.pool.ntp.org\n          - 2.pool.ntp.org\n          - 3.pool.ntp.org' "$file"
    fi
    
    # Add security updates and NTP configuration to runcmd
    # This is complex, so we'll use a Python script for this
    python3 <<EOF
import yaml
import sys

with open("$file", 'r') as f:
    content = f.read()

# Add security updates configuration before final_message
if "Configure automatic security updates" not in content:
    security_updates = '''        # 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 content:
        content = content.replace("      # Final message", security_updates + "\n      # Final message")
    elif "systemctl status qemu-guest-agent --no-pager || true" in content:
        content = content.replace("          systemctl status qemu-guest-agent --no-pager || true", 
                                  "          systemctl status qemu-guest-agent --no-pager || true" + security_updates)
    
    # Add write_files section if not present
    if "write_files:" not in content:
        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 content:
            content = content.replace("      # Final message", write_files + "      # Final message")
        else:
            content = content + "\n" + write_files
    
    # Enhance final_message
    if "final_message:" in content:
        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
        import re
        content = re.sub(r'      # Final message.*?(?=\n  providerConfigRef|\Z)', enhanced_final, content, flags=re.DOTALL)
    
    with open("$file", 'w') as f:
        f.write(content)
EOF
    
    # Update package verification to include new packages
    sed -i 's/for pkg in qemu-guest-agent curl wget net-tools; do/for pkg in qemu-guest-agent curl wget net-tools chrony unattended-upgrades; do/' "$file"
    
    log_success "  $vm_name enhanced"
}

echo "=========================================="
echo "Enhancing All VM Cloud-Init Configurations"
echo "=========================================="
echo ""

# Find all VM YAML files
VM_FILES=$(find "$VM_DIR" -name "*.yaml" -type f | grep -v ".backup" | sort)

TOTAL=$(echo "$VM_FILES" | wc -l)
COUNT=0

for file in $VM_FILES; do
    COUNT=$((COUNT + 1))
    enhance_vm_file "$file"
done

echo ""
echo "=========================================="
log_success "Enhanced $COUNT VM files"
echo "=========================================="
echo ""
log "Backup files created with .backup extension"
log "Review changes and remove backups when satisfied"