#!/bin/bash
#
# Docker Compose AppArmor Fix Script
#
# Adds 'security_opt: ["apparmor=unconfined"]' to all services in docker-compose.yml files
# This is required for Docker containers running inside LXC containers.
#
# Usage: ./fix-docker-apparmor.sh <LXC_IP> [COMPOSE_DIR]
#
# Example: ./fix-docker-apparmor.sh 10.10.0.214
# Example: ./fix-docker-apparmor.sh 10.10.0.214 /home/cal/container-data
#
# Arguments:
#   LXC_IP      - IP address of the LXC container to SSH into
#   COMPOSE_DIR - Optional directory containing docker-compose files (default: /home/cal/container-data)
#
# What this script does:
# 1. SSHs into the LXC container
# 2. Finds all docker-compose.yml files
# 3. Adds security_opt configuration to each service
# 4. Creates backups of original files
#
# Why this is needed:
# Docker containers in LXC need AppArmor disabled to function properly.
# Without this fix, containers may fail to start or have permission issues.
#

set -euo pipefail

# Color codes for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Function to print colored messages
log_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

log_debug() {
    echo -e "${BLUE}[DEBUG]${NC} $1"
}

# Parse arguments
if [[ $# -lt 1 ]]; then
    log_error "Insufficient arguments"
    echo "Usage: $0 <LXC_IP> [COMPOSE_DIR]"
    echo ""
    echo "Example: $0 10.10.0.214"
    echo "Example: $0 10.10.0.214 /home/cal/container-data"
    exit 1
fi

LXC_IP=$1
COMPOSE_DIR=${2:-/home/cal/container-data}

log_info "Starting AppArmor fix for Docker Compose files"
log_info "Target: root@$LXC_IP"
log_info "Directory: $COMPOSE_DIR"
echo ""

# Check SSH connectivity
log_info "Testing SSH connection to $LXC_IP..."
if ! ssh -o ConnectTimeout=5 -o BatchMode=yes root@"$LXC_IP" "echo 'SSH OK'" &>/dev/null; then
    log_error "Cannot connect to root@$LXC_IP via SSH"
    log_error "Please ensure:"
    echo "  1. SSH key is copied to the LXC container"
    echo "  2. Container is running"
    echo "  3. IP address is correct"
    exit 1
fi
log_info "✅ SSH connection successful"
echo ""

# Create Python script on remote host
log_info "Creating AppArmor fix script on remote host..."
ssh root@"$LXC_IP" "cat > /tmp/fix_apparmor.py" <<'PYTHON_SCRIPT'
#!/usr/bin/env python3
"""
Fix Docker Compose files to work in LXC by adding AppArmor unconfined security option.
"""
import yaml
import glob
import sys
import os
from pathlib import Path

def add_apparmor_fix(compose_file):
    """Add security_opt to all services in a docker-compose file."""
    print(f"\n📄 Processing: {compose_file}")

    # Create backup
    backup_file = f"{compose_file}.backup"
    if not os.path.exists(backup_file):
        os.system(f"cp '{compose_file}' '{backup_file}'")
        print(f"   ✅ Backup created: {backup_file}")
    else:
        print(f"   ⏭️  Backup already exists: {backup_file}")

    # Load compose file
    try:
        with open(compose_file, 'r') as f:
            compose_data = yaml.safe_load(f)
    except yaml.YAMLError as e:
        print(f"   ❌ Error parsing YAML: {e}")
        return False
    except Exception as e:
        print(f"   ❌ Error reading file: {e}")
        return False

    if not compose_data or 'services' not in compose_data:
        print(f"   ⚠️  No services found in compose file")
        return False

    # Track changes
    services_modified = 0
    services_skipped = 0

    # Add security_opt to each service
    for service_name, service_config in compose_data['services'].items():
        if service_config is None:
            service_config = {}
            compose_data['services'][service_name] = service_config

        # Check if security_opt already exists
        existing_security = service_config.get('security_opt', [])

        if 'apparmor=unconfined' in existing_security or 'apparmor:unconfined' in existing_security:
            print(f"   ⏭️  {service_name}: Already has AppArmor unconfined")
            services_skipped += 1
        else:
            # Add apparmor=unconfined
            if not existing_security:
                service_config['security_opt'] = ['apparmor=unconfined']
            else:
                if 'apparmor=unconfined' not in existing_security:
                    existing_security.append('apparmor=unconfined')
                    service_config['security_opt'] = existing_security

            print(f"   ✅ {service_name}: Added AppArmor unconfined")
            services_modified += 1

    # Write updated compose file
    try:
        with open(compose_file, 'w') as f:
            yaml.dump(compose_data, f, default_flow_style=False, sort_keys=False, indent=2)

        if services_modified > 0:
            print(f"   💾 Saved changes ({services_modified} services modified)")

        return services_modified > 0
    except Exception as e:
        print(f"   ❌ Error writing file: {e}")
        # Restore backup
        os.system(f"cp '{backup_file}' '{compose_file}'")
        print(f"   🔄 Restored from backup")
        return False

def main():
    """Main function to process all docker-compose files."""
    compose_dir = sys.argv[1] if len(sys.argv) > 1 else "/home/cal/container-data"

    print(f"🔍 Searching for docker-compose.yml files in {compose_dir}")

    # Find all docker-compose files
    patterns = [
        f"{compose_dir}/**/docker-compose.yml",
        f"{compose_dir}/**/docker-compose.yaml",
    ]

    compose_files = []
    for pattern in patterns:
        compose_files.extend(glob.glob(pattern, recursive=True))

    # Remove duplicates and sort
    compose_files = sorted(set(compose_files))

    if not compose_files:
        print(f"⚠️  No docker-compose files found in {compose_dir}")
        return 1

    print(f"📋 Found {len(compose_files)} docker-compose file(s)")

    # Process each file
    total_modified = 0
    total_errors = 0

    for compose_file in compose_files:
        try:
            if add_apparmor_fix(compose_file):
                total_modified += 1
        except Exception as e:
            print(f"   ❌ Unexpected error: {e}")
            total_errors += 1

    # Summary
    print("\n" + "="*60)
    print("📊 SUMMARY")
    print("="*60)
    print(f"Total files found:    {len(compose_files)}")
    print(f"Files modified:       {total_modified}")
    print(f"Files with errors:    {total_errors}")
    print(f"Files unchanged:      {len(compose_files) - total_modified - total_errors}")
    print("="*60)

    if total_modified > 0:
        print("\n✅ AppArmor fix applied successfully!")
        print("\n💡 Next steps:")
        print("   1. Review changes in modified files")
        print("   2. Start containers: docker compose up -d")
        print("   3. Check container status: docker compose ps")
        print("\n📝 Note: Backups created with .backup extension")

    return 0 if total_errors == 0 else 1

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

log_info "✅ Script uploaded to LXC container"
echo ""

# Install PyYAML if needed
log_info "Ensuring Python and PyYAML are installed..."
ssh root@"$LXC_IP" "apt-get update -qq && apt-get install -y -qq python3 python3-yaml > /dev/null 2>&1" || true
log_info "✅ Dependencies ready"
echo ""

# Run the fix script
log_info "Running AppArmor fix script..."
echo ""
ssh root@"$LXC_IP" "python3 /tmp/fix_apparmor.py '$COMPOSE_DIR'"
EXIT_CODE=$?
echo ""

# Cleanup
log_info "Cleaning up temporary files..."
ssh root@"$LXC_IP" "rm /tmp/fix_apparmor.py"
log_info "✅ Cleanup complete"
echo ""

if [[ $EXIT_CODE -eq 0 ]]; then
    log_info "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    log_info "🎉 AppArmor Fix Complete!"
    log_info "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo ""
    echo "Your docker-compose files have been updated to work in LXC."
    echo ""
    echo "Next steps:"
    echo "  1. SSH into container:"
    echo "     ssh root@$LXC_IP"
    echo ""
    echo "  2. Navigate to a service directory:"
    echo "     cd $COMPOSE_DIR/[service-name]"
    echo ""
    echo "  3. Start containers:"
    echo "     docker compose up -d"
    echo ""
    echo "  4. Check status:"
    echo "     docker compose ps"
    echo ""
else
    log_error "AppArmor fix encountered errors. Please review output above."
    exit 1
fi