#!/bin/bash
# Network Security Hardening Script
# Implements proper network segmentation, firewall rules, and security controls

# Import error handling library
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
source "$SCRIPT_DIR/lib/error_handling.sh"

# Configuration
readonly HOSTS=("omv800" "fedora" "surface" "jonathan-2518f5u" "audrey" "raspberrypi")
readonly HOST_IPS=("192.168.50.229" "192.168.50.225" "192.168.50.254" "192.168.50.181" "192.168.50.145" "192.168.50.107")
readonly SECURITY_CONFIG_DIR="/opt/migration/configs/security"
readonly FIREWALL_BACKUP_DIR="/opt/migration/backups/firewall_rules"

# Network zones configuration
readonly DMZ_NETWORK="192.168.51.0/24"
readonly MANAGEMENT_NETWORK="192.168.52.0/24"
readonly INTERNAL_NETWORK="192.168.50.0/24"
readonly DOCKER_SWARM_NETWORK="10.0.0.0/16"

# Service port mappings
declare -A SERVICE_PORTS=(
    ["traefik"]="80,443,8080"
    ["immich"]="3001"
    ["jellyfin"]="8096,8920"
    ["homeassistant"]="8123"
    ["appflowy"]="8000"
    ["paperless"]="8000"
    ["portainer"]="9000,9443"
    ["grafana"]="3000"
    ["prometheus"]="9090"
    ["postgres"]="5432"
    ["redis"]="6379"
    ["ssh"]="22"
)

# Security zones
declare -A SECURITY_ZONES=(
    ["public"]="traefik"
    ["dmz"]="immich,jellyfin,homeassistant,appflowy,paperless"
    ["internal"]="portainer,grafana,prometheus"
    ["data"]="postgres,redis"
    ["management"]="ssh"
)

# Cleanup function
cleanup_security_config() {
    log_info "Cleaning up temporary security configuration..."
    
    # Clean up temporary files
    rm -f /tmp/ufw_rules_*.tmp 2>/dev/null || true
    rm -f /tmp/iptables_rules_*.tmp 2>/dev/null || true
    
    log_info "Security configuration cleanup completed"
}

# Rollback function
rollback_security_config() {
    log_info "Rolling back security configuration changes..."
    
    # Restore original firewall rules from backup
    if [[ -d "$FIREWALL_BACKUP_DIR" ]]; then
        for host in "${HOSTS[@]}"; do
            local backup_file="$FIREWALL_BACKUP_DIR/${host}_ufw_backup.txt"
            if [[ -f "$backup_file" ]]; then
                log_info "Restoring firewall rules for $host from backup"
                ssh -o ConnectTimeout=10 "$host" "sudo ufw --force reset" 2>/dev/null || true
                
                # Restore basic rules to prevent lockout
                ssh "$host" "sudo ufw allow ssh" 2>/dev/null || true
                ssh "$host" "sudo ufw --force enable" 2>/dev/null || true
            fi
        done
    fi
    
    cleanup_security_config
    log_info "Security configuration rollback completed"
}

# Function to backup existing firewall rules
backup_firewall_rules() {
    log_step "Backing up existing firewall rules..."
    
    mkdir -p "$FIREWALL_BACKUP_DIR"
    
    for i in "${!HOSTS[@]}"; do
        local host="${HOSTS[$i]}"
        log_info "Backing up firewall rules from $host..."
        
        # Backup UFW rules
        if ssh -o ConnectTimeout=10 "$host" "sudo ufw status numbered" > "$FIREWALL_BACKUP_DIR/${host}_ufw_backup.txt" 2>/dev/null; then
            log_success "UFW rules backed up for $host"
        else
            log_warn "Could not backup UFW rules for $host (may not be installed)"
            echo "UFW not available" > "$FIREWALL_BACKUP_DIR/${host}_ufw_backup.txt"
        fi
        
        # Backup iptables rules
        if ssh -o ConnectTimeout=10 "$host" "sudo iptables-save" > "$FIREWALL_BACKUP_DIR/${host}_iptables_backup.txt" 2>/dev/null; then
            log_success "iptables rules backed up for $host"
        else
            log_warn "Could not backup iptables rules for $host"
            echo "iptables not available" > "$FIREWALL_BACKUP_DIR/${host}_iptables_backup.txt"
        fi
    done
    
    log_success "Firewall rules backup completed"
}

# Function to install security tools
install_security_tools() {
    log_step "Installing security tools on all hosts..."
    
    for i in "${!HOSTS[@]}"; do
        local host="${HOSTS[$i]}"
        log_info "Installing security tools on $host..."
        
        # Install UFW, fail2ban, and other security tools
        if ssh -o ConnectTimeout=30 "$host" "sudo apt-get update && sudo apt-get install -y ufw fail2ban iptables-persistent netfilter-persistent" 2>/dev/null; then
            log_success "Security tools installed on $host"
        else
            log_error "Failed to install security tools on $host"
            return 1
        fi
        
        # Install additional monitoring tools
        if ssh -o ConnectTimeout=30 "$host" "sudo apt-get install -y nmap tcpdump htop iotop nethogs" 2>/dev/null; then
            log_success "Monitoring tools installed on $host"
        else
            log_warn "Some monitoring tools may not have installed on $host"
        fi
    done
    
    log_success "Security tools installation completed"
}

# Function to configure network segmentation
configure_network_segmentation() {
    log_step "Configuring network segmentation..."
    
    # Create Docker networks for different security zones
    local manager_host="omv800"
    
    # Public zone (internet-facing)
    if ssh "$manager_host" "docker network create --driver overlay --subnet=10.1.0.0/24 public-zone" 2>/dev/null; then
        log_success "Created public-zone network"
    else
        log_info "Public-zone network may already exist"
    fi
    
    # DMZ zone (web services)
    if ssh "$manager_host" "docker network create --driver overlay --subnet=10.2.0.0/24 dmz-zone" 2>/dev/null; then
        log_success "Created dmz-zone network"
    else
        log_info "DMZ-zone network may already exist"
    fi
    
    # Internal zone (internal services)
    if ssh "$manager_host" "docker network create --driver overlay --subnet=10.3.0.0/24 internal-zone" 2>/dev/null; then
        log_success "Created internal-zone network"
    else
        log_info "Internal-zone network may already exist"
    fi
    
    # Data zone (databases)
    if ssh "$manager_host" "docker network create --driver overlay --subnet=10.4.0.0/24 data-zone" 2>/dev/null; then
        log_success "Created data-zone network"
    else
        log_info "Data-zone network may already exist"
    fi
    
    # Management zone (admin tools)
    if ssh "$manager_host" "docker network create --driver overlay --subnet=10.5.0.0/24 management-zone" 2>/dev/null; then
        log_success "Created management-zone network"
    else
        log_info "Management-zone network may already exist"
    fi
    
    log_success "Network segmentation configuration completed"
}

# Function to configure host-level firewalls
configure_host_firewalls() {
    log_step "Configuring host-level firewalls..."
    
    for i in "${!HOSTS[@]}"; do
        local host="${HOSTS[$i]}"
        local ip="${HOST_IPS[$i]}"
        
        log_info "Configuring firewall for $host ($ip)..."
        
        # Reset UFW to clean state
        ssh "$host" "sudo ufw --force reset" 2>/dev/null || true
        
        # Set default policies
        ssh "$host" "sudo ufw default deny incoming"
        ssh "$host" "sudo ufw default allow outgoing"
        ssh "$host" "sudo ufw default deny forward"
        
        # Allow SSH from local network only
        ssh "$host" "sudo ufw allow from 192.168.50.0/24 to any port 22"
        
        # Allow Docker Swarm communication between nodes
        if [[ "$host" != "raspberrypi" ]]; then  # raspberrypi is backup only
            for other_ip in "${HOST_IPS[@]}"; do
                if [[ "$other_ip" != "$ip" ]] && [[ "$other_ip" != "192.168.50.107" ]]; then
                    # Docker Swarm ports
                    ssh "$host" "sudo ufw allow from $other_ip to any port 2377"  # Cluster management
                    ssh "$host" "sudo ufw allow from $other_ip to any port 7946"  # Node communication
                    ssh "$host" "sudo ufw allow from $other_ip to any port 4789"  # Overlay network traffic
                fi
            done
        fi
        
        # Configure service-specific rules based on host role
        configure_service_specific_rules "$host" "$ip"
        
        # Enable UFW
        ssh "$host" "sudo ufw --force enable"
        
        # Verify UFW status
        if ssh "$host" "sudo ufw status" | grep -q "Status: active"; then
            log_success "Firewall configured and enabled on $host"
        else
            log_error "Firewall configuration failed on $host"
            return 1
        fi
    done
    
    log_success "Host-level firewall configuration completed"
}

# Function to configure service-specific firewall rules
configure_service_specific_rules() {
    local host=$1
    local ip=$2
    
    case $host in
        "omv800")
            # Primary hub - needs most services accessible
            ssh "$host" "sudo ufw allow from 192.168.50.0/24 to any port 80"      # HTTP
            ssh "$host" "sudo ufw allow from 192.168.50.0/24 to any port 443"     # HTTPS
            ssh "$host" "sudo ufw allow from 192.168.50.0/24 to any port 8080"    # Traefik dashboard
            ssh "$host" "sudo ufw allow from 192.168.50.0/24 to any port 3001"    # Immich
            ssh "$host" "sudo ufw allow from 192.168.50.0/24 to any port 8096"    # Jellyfin
            ssh "$host" "sudo ufw allow from 192.168.50.0/24 to any port 5432"    # PostgreSQL (internal)
            ssh "$host" "sudo ufw allow from 192.168.50.0/24 to any port 6379"    # Redis (internal)
            ssh "$host" "sudo ufw allow from 192.168.50.0/24 to any port 111"     # NFS portmapper
            ssh "$host" "sudo ufw allow from 192.168.50.0/24 to any port 2049"    # NFS
            ;;
        "surface")
            # Development hub
            ssh "$host" "sudo ufw allow from 192.168.50.0/24 to any port 8000"    # AppFlowy
            ssh "$host" "sudo ufw allow from 192.168.50.0/24 to any port 3000"    # Development ports
            ssh "$host" "sudo ufw allow from 192.168.50.0/24 to any port 5000"    # Additional dev ports
            ;;
        "jonathan-2518f5u")
            # IoT hub
            ssh "$host" "sudo ufw allow from 192.168.50.0/24 to any port 8123"    # Home Assistant
            ssh "$host" "sudo ufw allow from 192.168.50.0/24 to any port 1883"    # MQTT
            ssh "$host" "sudo ufw allow from 192.168.50.0/24 to any port 6052"    # ESPHome
            ;;
        "audrey")
            # Monitoring hub
            ssh "$host" "sudo ufw allow from 192.168.50.0/24 to any port 3000"    # Grafana
            ssh "$host" "sudo ufw allow from 192.168.50.0/24 to any port 9090"    # Prometheus
            ssh "$host" "sudo ufw allow from 192.168.50.0/24 to any port 9093"    # Alertmanager
            ;;
        "fedora")
            # Compute hub
            ssh "$host" "sudo ufw allow from 192.168.50.0/24 to any port 8080"    # n8n or other automation
            ssh "$host" "sudo ufw allow from 192.168.50.0/24 to any port 5000"    # General services
            ;;
        "raspberrypi")
            # Backup hub - minimal access
            ssh "$host" "sudo ufw allow from 192.168.50.0/24 to any port 873"     # Rsync
            ssh "$host" "sudo ufw allow from 192.168.50.0/24 to any port 111"     # NFS portmapper
            ssh "$host" "sudo ufw allow from 192.168.50.0/24 to any port 2049"    # NFS
            ;;
    esac
    
    log_debug "Service-specific rules configured for $host"
}

# Function to configure fail2ban
configure_fail2ban() {
    log_step "Configuring fail2ban intrusion detection..."
    
    mkdir -p "$SECURITY_CONFIG_DIR/fail2ban"
    
    # Create custom jail configuration
    cat > "$SECURITY_CONFIG_DIR/fail2ban/jail.local" << 'EOF'
[DEFAULT]
# Ban settings
bantime = 3600
findtime = 600
maxretry = 3
backend = auto

# Email settings (configure SMTP if needed)
destemail = admin@homelab.local
sender = fail2ban@homelab.local
mta = sendmail

# Action
action = %(action_mwl)s

[sshd]
enabled = true
port = ssh
filter = sshd
logpath = /var/log/auth.log
maxretry = 3
bantime = 3600

[docker-auth]
enabled = true
port = 2376,2377
filter = docker-auth
logpath = /var/log/daemon.log
maxretry = 3
bantime = 1800

[traefik-auth]
enabled = true
port = http,https
filter = traefik-auth
logpath = /var/log/traefik/access.log
maxretry = 5
bantime = 1800

[nginx-http-auth]
enabled = true
port = http,https
filter = nginx-http-auth
logpath = /var/log/nginx/error.log
maxretry = 5
bantime = 600
EOF

    # Create custom filter for Docker authentication
    cat > "$SECURITY_CONFIG_DIR/fail2ban/filter.d/docker-auth.conf" << 'EOF'
[Definition]
failregex = ^.*authentication failure.*rhost=<HOST>.*$
            ^.*authentication error.*rhost=<HOST>.*$
            ^.*invalid user.*from <HOST>.*$
ignoreregex =
EOF

    # Create custom filter for Traefik authentication
    cat > "$SECURITY_CONFIG_DIR/fail2ban/filter.d/traefik-auth.conf" << 'EOF'
[Definition]
failregex = ^.*"GET.*HTTP/1\.[01]" 401 .*".*" ".*" .*"<HOST>".*$
            ^.*"POST.*HTTP/1\.[01]" 401 .*".*" ".*" .*"<HOST>".*$
            ^.*"GET.*HTTP/1\.[01]" 403 .*".*" ".*" .*"<HOST>".*$
ignoreregex =
EOF

    # Deploy fail2ban configuration to all hosts
    for host in "${HOSTS[@]}"; do
        log_info "Configuring fail2ban on $host..."
        
        # Copy configuration files
        scp "$SECURITY_CONFIG_DIR/fail2ban/jail.local" "$host:/tmp/"
        ssh "$host" "sudo mv /tmp/jail.local /etc/fail2ban/"
        
        # Create filter directories and copy filters
        ssh "$host" "sudo mkdir -p /etc/fail2ban/filter.d"
        scp "$SECURITY_CONFIG_DIR/fail2ban/filter.d/"* "$host:/tmp/"
        ssh "$host" "sudo mv /tmp/*.conf /etc/fail2ban/filter.d/"
        
        # Restart fail2ban
        if ssh "$host" "sudo systemctl restart fail2ban && sudo systemctl enable fail2ban"; then
            log_success "fail2ban configured on $host"
        else
            log_warn "fail2ban configuration may have issues on $host"
        fi
    done
    
    log_success "fail2ban configuration completed"
}

# Function to enhance SSL/TLS configuration
enhance_ssl_configuration() {
    log_step "Enhancing SSL/TLS configuration..."
    
    mkdir -p "$SECURITY_CONFIG_DIR/tls"
    
    # Create enhanced TLS configuration for Traefik
    cat > "$SECURITY_CONFIG_DIR/tls/tls-security.yml" << 'EOF'
# Enhanced TLS Configuration for Traefik
tls:
  options:
    default:
      minVersion: "VersionTLS12"
      maxVersion: "VersionTLS13"
      cipherSuites:
        - "TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384"
        - "TLS_ECDHE_RSA_WITH_CHACHA20_POLY1305"
        - "TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256"
        - "TLS_RSA_WITH_AES_256_GCM_SHA384"
        - "TLS_RSA_WITH_AES_128_GCM_SHA256"
      curvePreferences:
        - "CurveP521"
        - "CurveP384"
      sniStrict: true
      
    strict:
      minVersion: "VersionTLS12"
      maxVersion: "VersionTLS13"
      cipherSuites:
        - "TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384"
        - "TLS_ECDHE_RSA_WITH_CHACHA20_POLY1305"
      curvePreferences:
        - "CurveP521"
        - "CurveP384"
      sniStrict: true
      clientAuth:
        caFiles:
          - "/etc/traefik/ca-cert.pem"
        clientAuthType: "RequireAndVerifyClientCert"
        
  certificates:
    - certFile: "/etc/traefik/certs/homelab.crt"
      keyFile: "/etc/traefik/certs/homelab.key"
      stores:
        - "default"
EOF

    # Create security headers configuration
    cat > "$SECURITY_CONFIG_DIR/tls/security-headers-enhanced.yml" << 'EOF'
# Enhanced Security Headers
http:
  middlewares:
    security-headers-enhanced:
      headers:
        # HSTS headers
        forceSTSHeader: true
        stsIncludeSubdomains: true
        stsPreload: true
        stsSeconds: 63072000  # 2 years
        
        # XSS Protection
        browserXssFilter: true
        customResponseHeaders:
          X-XSS-Protection: "1; mode=block"
          
        # Content Type Options
        contentTypeNosniff: true
        
        # Frame Options
        frameDeny: true
        customFrameOptionsValue: "SAMEORIGIN"
        
        # Referrer Policy
        referrerPolicy: "strict-origin-when-cross-origin"
        
        # Permissions Policy
        permissionsPolicy: "camera=(), microphone=(), geolocation=(), payment=(), usb=()"
        
        # Content Security Policy
        contentSecurityPolicy: "default-src 'self'; script-src 'self' 'unsafe-inline' 'unsafe-eval'; style-src 'self' 'unsafe-inline'; img-src 'self' data: https:; font-src 'self' data:; connect-src 'self'; frame-ancestors 'none'"
        
        # Additional security headers
        customResponseHeaders:
          X-Content-Type-Options: "nosniff"
          X-Frame-Options: "SAMEORIGIN"
          X-Permitted-Cross-Domain-Policies: "none"
          Cross-Origin-Embedder-Policy: "require-corp"
          Cross-Origin-Opener-Policy: "same-origin"
          Cross-Origin-Resource-Policy: "same-origin"
          
        # Remove server information
        customRequestHeaders:
          X-Forwarded-Proto: "https"
        
        # SSL redirect
        sslRedirect: true
        sslForceHost: true
EOF

    log_success "Enhanced SSL/TLS configuration created"
}

# Function to create network security monitoring
setup_network_monitoring() {
    log_step "Setting up network security monitoring..."
    
    mkdir -p "$SECURITY_CONFIG_DIR/monitoring"
    
    # Create network monitoring script
    cat > "$SECURITY_CONFIG_DIR/monitoring/network_monitor.sh" << 'EOF'
#!/bin/bash
# Network Security Monitor
# Monitors for suspicious network activity

LOG_FILE="/var/log/network_monitor.log"
ALERT_THRESHOLD=100  # connections per minute

log_alert() {
    echo "$(date): ALERT - $1" >> "$LOG_FILE"
    # Send alert (configure notification method)
    logger "NETWORK_SECURITY_ALERT: $1"
}

# Monitor connection attempts
monitor_connections() {
    local connections=$(ss -tn | grep :22 | wc -l)
    if [[ $connections -gt $ALERT_THRESHOLD ]]; then
        log_alert "High SSH connection count: $connections"
    fi
    
    # Monitor failed authentication attempts
    local failed_auth=$(tail -100 /var/log/auth.log | grep "authentication failure" | wc -l)
    if [[ $failed_auth -gt 10 ]]; then
        log_alert "High failed authentication count: $failed_auth"
    fi
}

# Monitor Docker security
monitor_docker_security() {
    # Check for privileged containers
    local privileged_containers=$(docker ps --filter "privileged=true" -q | wc -l)
    if [[ $privileged_containers -gt 0 ]]; then
        log_alert "Privileged containers detected: $privileged_containers"
    fi
    
    # Check for containers with host network
    local host_network_containers=$(docker ps --format "{{.Names}} {{.NetworkMode}}" | grep host | wc -l)
    if [[ $host_network_containers -gt 1 ]]; then  # Allow one for monitoring
        log_alert "Multiple containers using host network: $host_network_containers"
    fi
}

# Main monitoring loop
while true; do
    monitor_connections
    monitor_docker_security
    sleep 60
done
EOF

    chmod +x "$SECURITY_CONFIG_DIR/monitoring/network_monitor.sh"
    
    # Deploy monitoring to all hosts
    for host in "${HOSTS[@]}"; do
        log_info "Setting up network monitoring on $host..."
        
        scp "$SECURITY_CONFIG_DIR/monitoring/network_monitor.sh" "$host:/tmp/"
        ssh "$host" "sudo mv /tmp/network_monitor.sh /usr/local/bin/ && sudo chmod +x /usr/local/bin/network_monitor.sh"
        
        # Create systemd service for monitoring
        ssh "$host" "cat > /tmp/network-monitor.service << 'SERVICE_EOF'
[Unit]
Description=Network Security Monitor
After=network.target

[Service]
ExecStart=/usr/local/bin/network_monitor.sh
Restart=always
RestartSec=10
User=root

[Install]
WantedBy=multi-user.target
SERVICE_EOF"
        
        ssh "$host" "sudo mv /tmp/network-monitor.service /etc/systemd/system/"
        ssh "$host" "sudo systemctl daemon-reload && sudo systemctl enable network-monitor.service"
        
        if ssh "$host" "sudo systemctl start network-monitor.service"; then
            log_success "Network monitoring started on $host"
        else
            log_warn "Network monitoring may have issues on $host"
        fi
    done
    
    log_success "Network security monitoring setup completed"
}

# Function to create security audit report
create_security_audit() {
    log_step "Creating security audit report..."
    
    local audit_file="/opt/migration/security_audit_$(date +%Y%m%d_%H%M%S).md"
    
    cat > "$audit_file" << EOF
# Network Security Audit Report

**Generated:** $(date)  
**Configuration:** Enhanced network segmentation and security hardening  

## Security Zones Implemented

### Network Segmentation
- **Public Zone:** 10.1.0.0/24 (Traefik reverse proxy)
- **DMZ Zone:** 10.2.0.0/24 (Web services - Immich, Jellyfin, Home Assistant)
- **Internal Zone:** 10.3.0.0/24 (Management tools - Portainer, Grafana)
- **Data Zone:** 10.4.0.0/24 (Databases - PostgreSQL, Redis)
- **Management Zone:** 10.5.0.0/24 (Admin tools)

### Host Firewall Status
EOF

    # Check firewall status on each host
    for i in "${!HOSTS[@]}"; do
        local host="${HOSTS[$i]}"
        local ip="${HOST_IPS[$i]}"
        
        echo "#### $host ($ip)" >> "$audit_file"
        
        # Check UFW status
        local ufw_status=$(ssh -o ConnectTimeout=10 "$host" "sudo ufw status" 2>/dev/null || echo "Error getting status")
        echo "\`\`\`" >> "$audit_file"
        echo "$ufw_status" >> "$audit_file"
        echo "\`\`\`" >> "$audit_file"
        echo "" >> "$audit_file"
    done
    
    cat >> "$audit_file" << EOF

### Security Tools Status
EOF

    # Check fail2ban status
    for host in "${HOSTS[@]}"; do
        echo "#### fail2ban on $host" >> "$audit_file"
        local fail2ban_status=$(ssh -o ConnectTimeout=10 "$host" "sudo fail2ban-client status" 2>/dev/null || echo "Error getting status")
        echo "\`\`\`" >> "$audit_file"
        echo "$fail2ban_status" >> "$audit_file"
        echo "\`\`\`" >> "$audit_file"
        echo "" >> "$audit_file"
    done
    
    cat >> "$audit_file" << EOF

### Recommendations
1. **Regular Updates:** Ensure all security tools are regularly updated
2. **Log Monitoring:** Implement centralized log monitoring and alerting
3. **Certificate Management:** Set up automated certificate renewal monitoring
4. **Penetration Testing:** Schedule regular security assessments
5. **Backup Security:** Verify backup encryption and off-site storage

### Next Steps
- [ ] Test all firewall rules and service accessibility
- [ ] Configure centralized logging for security events
- [ ] Set up automated security scanning
- [ ] Implement network intrusion detection system (IDS)
- [ ] Create incident response procedures
EOF

    log_success "Security audit report created: $audit_file"
    echo "$audit_file"
}

# Main execution function
main() {
    local action=${1:-"full"}
    
    # Register cleanup and rollback functions
    register_cleanup cleanup_security_config
    register_rollback rollback_security_config
    
    case $action in
        "full")
            log_step "Starting full network security hardening..."
            
            # Validate prerequisites
            validate_prerequisites ssh scp
            
            # Validate network connectivity
            validate_network_connectivity "${HOST_IPS[@]}"
            
            # Create checkpoint
            create_checkpoint "security_hardening_start"
            
            # Backup existing configurations
            backup_firewall_rules
            create_checkpoint "firewall_backup_complete"
            
            # Install security tools
            install_security_tools
            create_checkpoint "security_tools_installed"
            
            # Configure network segmentation
            configure_network_segmentation
            create_checkpoint "network_segmentation_complete"
            
            # Configure host firewalls
            configure_host_firewalls
            create_checkpoint "host_firewalls_complete"
            
            # Configure fail2ban
            configure_fail2ban
            create_checkpoint "fail2ban_complete"
            
            # Enhance SSL configuration
            enhance_ssl_configuration
            create_checkpoint "ssl_enhancement_complete"
            
            # Setup network monitoring
            setup_network_monitoring
            create_checkpoint "network_monitoring_complete"
            
            # Create security audit
            local audit_report=$(create_security_audit)
            
            log_success "✅ Network security hardening completed successfully!"
            log_info "🔒 Security audit report: $audit_report"
            ;;
            
        "backup-only")
            backup_firewall_rules
            ;;
            
        "firewall-only")
            configure_host_firewalls
            ;;
            
        "fail2ban-only")
            configure_fail2ban
            ;;
            
        "audit-only")
            create_security_audit
            ;;
            
        "help"|*)
            cat << EOF
Network Security Hardening Script

Usage: $0 <action>

Actions:
  full         - Complete security hardening (default)
  backup-only  - Only backup existing firewall rules
  firewall-only - Only configure host firewalls
  fail2ban-only - Only configure fail2ban
  audit-only   - Only create security audit report
  help         - Show this help

Examples:
  $0 full
  $0 firewall-only
  $0 audit-only
EOF
            ;;
    esac
}

# Execute main function
main "$@"