#!/bin/bash
# BZZZ v2 Security Hardening Script
# Applies comprehensive security configurations for the cluster

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
LOG_FILE="/var/log/bzzz-security-hardening-$(date +%Y%m%d-%H%M%S).log"

# Colors
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'

log() {
    echo -e "${BLUE}[$(date +'%Y-%m-%d %H:%M:%S')]${NC} $1" | tee -a "$LOG_FILE"
}

error() {
    echo -e "${RED}[ERROR]${NC} $1" | tee -a "$LOG_FILE"
    exit 1
}

warn() {
    echo -e "${YELLOW}[WARN]${NC} $1" | tee -a "$LOG_FILE"
}

success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1" | tee -a "$LOG_FILE"
}

check_root() {
    if [[ $EUID -eq 0 ]]; then
        error "This script should not be run as root. Run as tony user with sudo access."
    fi
}

configure_firewall() {
    log "Configuring UFW firewall for BZZZ v2..."
    
    # Enable UFW if not enabled
    sudo ufw --force enable
    
    # Default policies
    sudo ufw default deny incoming
    sudo ufw default allow outgoing
    
    # SSH access
    sudo ufw allow ssh
    
    # Docker Swarm ports (internal cluster only)
    sudo ufw allow from 192.168.1.0/24 to any port 2376 proto tcp comment "Docker daemon TLS"
    sudo ufw allow from 192.168.1.0/24 to any port 2377 proto tcp comment "Docker Swarm management"
    sudo ufw allow from 192.168.1.0/24 to any port 7946 proto tcp comment "Docker Swarm node communication"
    sudo ufw allow from 192.168.1.0/24 to any port 7946 proto udp comment "Docker Swarm node communication"
    sudo ufw allow from 192.168.1.0/24 to any port 4789 proto udp comment "Docker Swarm overlay networks"
    
    # BZZZ v2 P2P ports (internal cluster only)
    sudo ufw allow from 192.168.1.0/24 to any port 9000:9300 proto tcp comment "BZZZ v2 P2P"
    sudo ufw allow from 192.168.1.0/24 to any port 9000:9300 proto udp comment "BZZZ v2 P2P"
    
    # DHT bootstrap ports
    sudo ufw allow from 192.168.1.0/24 to any port 9101:9103 proto tcp comment "BZZZ DHT Bootstrap"
    
    # mDNS discovery (local network only)
    sudo ufw allow from 192.168.1.0/24 to any port 5353 proto udp comment "mDNS discovery"
    
    # HTTP/HTTPS through Traefik (external access)
    sudo ufw allow 80/tcp comment "HTTP"
    sudo ufw allow 443/tcp comment "HTTPS"
    
    # Internal service ports (cluster only)
    sudo ufw allow from 192.168.1.0/24 to any port 3000:3100 proto tcp comment "BZZZ v2 services"
    sudo ufw allow from 192.168.1.0/24 to any port 5432 proto tcp comment "PostgreSQL"
    sudo ufw allow from 192.168.1.0/24 to any port 6379 proto tcp comment "Redis"
    
    # Monitoring ports (cluster only)
    sudo ufw allow from 192.168.1.0/24 to any port 9090:9203 proto tcp comment "Monitoring"
    
    # Rate limiting rules
    sudo ufw limit ssh comment "Rate limit SSH"
    
    # Log denied connections
    sudo ufw logging on
    
    success "Firewall configured successfully"
}

configure_docker_security() {
    log "Configuring Docker security..."
    
    # Create Docker daemon configuration
    sudo mkdir -p /etc/docker
    
    cat << 'EOF' | sudo tee /etc/docker/daemon.json > /dev/null
{
    "log-driver": "json-file",
    "log-opts": {
        "max-size": "100m",
        "max-file": "3"
    },
    "live-restore": true,
    "userland-proxy": false,
    "icc": false,
    "userns-remap": "default",
    "no-new-privileges": true,
    "seccomp-profile": "/etc/docker/seccomp-default.json",
    "apparmor-profile": "docker-default",
    "storage-driver": "overlay2",
    "storage-opts": [
        "overlay2.override_kernel_check=true"
    ],
    "default-ulimits": {
        "nofile": {
            "Name": "nofile",
            "Hard": 65536,
            "Soft": 65536
        }
    },
    "registry-mirrors": ["https://registry.home.deepblack.cloud"],
    "insecure-registries": ["registry.home.deepblack.cloud:5000"],
    "features": {
        "buildkit": true
    }
}
EOF
    
    # Create custom seccomp profile
    cat << 'EOF' | sudo tee /etc/docker/seccomp-default.json > /dev/null
{
    "defaultAction": "SCMP_ACT_ERRNO",
    "architectures": [
        "SCMP_ARCH_X86_64",
        "SCMP_ARCH_X86",
        "SCMP_ARCH_X32"
    ],
    "syscalls": [
        {
            "names": [
                "accept",
                "access",
                "arch_prctl",
                "bind",
                "brk",
                "chdir",
                "chmod",
                "chown",
                "clone",
                "close",
                "connect",
                "dup",
                "dup2",
                "epoll_create",
                "epoll_ctl",
                "epoll_wait",
                "execve",
                "exit",
                "exit_group",
                "fcntl",
                "fstat",
                "futex",
                "getcwd",
                "getdents",
                "getgid",
                "getpid",
                "getppid",
                "gettid",
                "getuid",
                "listen",
                "lstat",
                "mmap",
                "mprotect",
                "munmap",
                "nanosleep",
                "open",
                "openat",
                "pipe",
                "poll",
                "prctl",
                "read",
                "readlink",
                "recv",
                "recvfrom",
                "rt_sigaction",
                "rt_sigprocmask",
                "rt_sigreturn",
                "sched_yield",
                "send",
                "sendto",
                "set_robust_list",
                "setsockopt",
                "socket",
                "stat",
                "write"
            ],
            "action": "SCMP_ACT_ALLOW"
        }
    ]
}
EOF
    
    # Restart Docker to apply changes
    sudo systemctl daemon-reload
    sudo systemctl restart docker
    
    success "Docker security configuration applied"
}

setup_tls_certificates() {
    log "Setting up TLS certificates..."
    
    # Create certificates directory
    mkdir -p /rust/bzzz-v2/config/tls/{ca,server,client}
    
    # Generate CA key and certificate
    if [[ ! -f /rust/bzzz-v2/config/tls/ca/ca-key.pem ]]; then
        openssl genrsa -out /rust/bzzz-v2/config/tls/ca/ca-key.pem 4096
        openssl req -new -x509 -days 3650 -key /rust/bzzz-v2/config/tls/ca/ca-key.pem \
            -out /rust/bzzz-v2/config/tls/ca/ca.pem \
            -subj "/C=US/ST=Cloud/L=DeepBlack/O=BZZZ/CN=bzzz-ca"
        
        log "Generated new CA certificate"
    fi
    
    # Generate server certificates for each node
    local nodes=("walnut" "ironwood" "acacia")
    for node in "${nodes[@]}"; do
        if [[ ! -f "/rust/bzzz-v2/config/tls/server/${node}-key.pem" ]]; then
            # Generate server key
            openssl genrsa -out "/rust/bzzz-v2/config/tls/server/${node}-key.pem" 4096
            
            # Generate server certificate request
            openssl req -new -key "/rust/bzzz-v2/config/tls/server/${node}-key.pem" \
                -out "/rust/bzzz-v2/config/tls/server/${node}.csr" \
                -subj "/C=US/ST=Cloud/L=DeepBlack/O=BZZZ/CN=${node}.deepblack.cloud"
            
            # Create extensions file
            cat > "/rust/bzzz-v2/config/tls/server/${node}-ext.cnf" << EOF
subjectAltName = DNS:${node}.deepblack.cloud,DNS:${node},DNS:localhost,IP:127.0.0.1,IP:192.168.1.27
extendedKeyUsage = serverAuth,clientAuth
EOF
            
            # Generate server certificate
            openssl x509 -req -days 365 -in "/rust/bzzz-v2/config/tls/server/${node}.csr" \
                -CA /rust/bzzz-v2/config/tls/ca/ca.pem \
                -CAkey /rust/bzzz-v2/config/tls/ca/ca-key.pem \
                -out "/rust/bzzz-v2/config/tls/server/${node}.pem" \
                -extensions v3_req -extfile "/rust/bzzz-v2/config/tls/server/${node}-ext.cnf" \
                -CAcreateserial
            
            # Clean up CSR and extensions file
            rm "/rust/bzzz-v2/config/tls/server/${node}.csr" "/rust/bzzz-v2/config/tls/server/${node}-ext.cnf"
            
            log "Generated TLS certificate for $node"
        fi
    done
    
    # Generate client certificates for inter-service communication
    if [[ ! -f /rust/bzzz-v2/config/tls/client/client-key.pem ]]; then
        openssl genrsa -out /rust/bzzz-v2/config/tls/client/client-key.pem 4096
        openssl req -new -key /rust/bzzz-v2/config/tls/client/client-key.pem \
            -out /rust/bzzz-v2/config/tls/client/client.csr \
            -subj "/C=US/ST=Cloud/L=DeepBlack/O=BZZZ/CN=bzzz-client"
        
        openssl x509 -req -days 365 -in /rust/bzzz-v2/config/tls/client/client.csr \
            -CA /rust/bzzz-v2/config/tls/ca/ca.pem \
            -CAkey /rust/bzzz-v2/config/tls/ca/ca-key.pem \
            -out /rust/bzzz-v2/config/tls/client/client.pem \
            -CAcreateserial
        
        rm /rust/bzzz-v2/config/tls/client/client.csr
        
        log "Generated client certificate"
    fi
    
    # Set appropriate permissions
    chmod -R 600 /rust/bzzz-v2/config/tls
    chmod 755 /rust/bzzz-v2/config/tls /rust/bzzz-v2/config/tls/{ca,server,client}
    
    success "TLS certificates configured"
}

configure_secrets_management() {
    log "Configuring secrets management..."
    
    # Create secrets directory with restricted permissions
    mkdir -p /rust/bzzz-v2/config/secrets
    chmod 700 /rust/bzzz-v2/config/secrets
    
    # Generate random secrets if they don't exist
    local secrets=(
        "postgres_password"
        "redis_password"
        "grafana_admin_password"
        "prometheus_web_password"
        "alertmanager_web_password"
    )
    
    for secret in "${secrets[@]}"; do
        local secret_file="/rust/bzzz-v2/config/secrets/${secret}"
        if [[ ! -f "$secret_file" ]]; then
            openssl rand -base64 32 > "$secret_file"
            chmod 600 "$secret_file"
            log "Generated secret: $secret"
        fi
    done
    
    # Create Docker secrets
    for secret in "${secrets[@]}"; do
        local secret_file="/rust/bzzz-v2/config/secrets/${secret}"
        if docker secret inspect "bzzz_${secret}" >/dev/null 2>&1; then
            log "Docker secret bzzz_${secret} already exists"
        else
            docker secret create "bzzz_${secret}" "$secret_file"
            log "Created Docker secret: bzzz_${secret}"
        fi
    done
    
    # Handle OpenAI API key if it exists
    local openai_key_file="/home/tony/chorus/business/secrets/openai-api-key"
    if [[ -f "$openai_key_file" ]]; then
        if ! docker secret inspect bzzz_openai_api_key >/dev/null 2>&1; then
            docker secret create bzzz_openai_api_key "$openai_key_file"
            log "Created OpenAI API key secret"
        fi
    else
        warn "OpenAI API key not found at $openai_key_file"
    fi
    
    success "Secrets management configured"
}

setup_network_security() {
    log "Setting up network security..."
    
    # Configure iptables rules for container isolation
    cat << 'EOF' | sudo tee /etc/iptables/rules.v4 > /dev/null
*filter
:INPUT ACCEPT [0:0]
:FORWARD DROP [0:0]
:OUTPUT ACCEPT [0:0]
:DOCKER-USER - [0:0]

# Allow established connections
-A INPUT -m state --state ESTABLISHED,RELATED -j ACCEPT

# Allow loopback
-A INPUT -i lo -j ACCEPT

# Allow SSH (with rate limiting)
-A INPUT -p tcp --dport 22 -m state --state NEW -m recent --set
-A INPUT -p tcp --dport 22 -m state --state NEW -m recent --update --seconds 60 --hitcount 4 -j DROP
-A INPUT -p tcp --dport 22 -j ACCEPT

# Allow HTTP/HTTPS
-A INPUT -p tcp --dport 80 -j ACCEPT
-A INPUT -p tcp --dport 443 -j ACCEPT

# Allow Docker Swarm (internal network only)
-A INPUT -s 192.168.1.0/24 -p tcp --dport 2376 -j ACCEPT
-A INPUT -s 192.168.1.0/24 -p tcp --dport 2377 -j ACCEPT
-A INPUT -s 192.168.1.0/24 -p tcp --dport 7946 -j ACCEPT
-A INPUT -s 192.168.1.0/24 -p udp --dport 7946 -j ACCEPT
-A INPUT -s 192.168.1.0/24 -p udp --dport 4789 -j ACCEPT

# Allow BZZZ P2P (internal network only)
-A INPUT -s 192.168.1.0/24 -p tcp --dport 9000:9300 -j ACCEPT
-A INPUT -s 192.168.1.0/24 -p udp --dport 9000:9300 -j ACCEPT

# Block container-to-host access except for specific services
-A DOCKER-USER -i docker_gwbridge -j ACCEPT
-A DOCKER-USER -i docker0 -j ACCEPT
-A DOCKER-USER -j DROP

# Drop everything else
-A INPUT -j DROP

COMMIT
EOF
    
    # Apply iptables rules
    sudo iptables-restore < /etc/iptables/rules.v4
    
    # Enable IP forwarding for Docker
    echo 'net.ipv4.ip_forward=1' | sudo tee -a /etc/sysctl.conf
    echo 'net.ipv6.conf.all.forwarding=1' | sudo tee -a /etc/sysctl.conf
    
    # Kernel security parameters
    cat << 'EOF' | sudo tee -a /etc/sysctl.conf > /dev/null

# BZZZ v2 Security Parameters
net.ipv4.conf.all.rp_filter=1
net.ipv4.conf.default.rp_filter=1
net.ipv4.icmp_echo_ignore_broadcasts=1
net.ipv4.icmp_ignore_bogus_error_responses=1
net.ipv4.tcp_syncookies=1
net.ipv4.conf.all.log_martians=1
net.ipv4.conf.default.log_martians=1
net.ipv4.conf.all.accept_source_route=0
net.ipv4.conf.default.accept_source_route=0
net.ipv6.conf.all.accept_source_route=0
net.ipv6.conf.default.accept_source_route=0
net.ipv4.conf.all.accept_redirects=0
net.ipv4.conf.default.accept_redirects=0
net.ipv6.conf.all.accept_redirects=0
net.ipv6.conf.default.accept_redirects=0
net.ipv4.conf.all.secure_redirects=0
net.ipv4.conf.default.secure_redirects=0
net.ipv4.conf.all.send_redirects=0
net.ipv4.conf.default.send_redirects=0

# Kernel hardening
kernel.dmesg_restrict=1
kernel.kptr_restrict=2
kernel.yama.ptrace_scope=1
fs.suid_dumpable=0
kernel.core_uses_pid=1
EOF
    
    # Apply sysctl settings
    sudo sysctl -p
    
    success "Network security configured"
}

configure_audit_logging() {
    log "Configuring audit logging..."
    
    # Install auditd if not present
    if ! command -v auditctl &> /dev/null; then
        sudo apt-get update
        sudo apt-get install -y auditd audispd-plugins
    fi
    
    # Configure audit rules
    cat << 'EOF' | sudo tee /etc/audit/rules.d/bzzz-v2.rules > /dev/null
# BZZZ v2 Audit Rules

# Monitor file changes in sensitive directories
-w /etc/docker/ -p wa -k docker-config
-w /rust/bzzz-v2/config/secrets/ -p wa -k bzzz-secrets
-w /rust/bzzz-v2/config/tls/ -p wa -k bzzz-tls
-w /etc/ssl/ -p wa -k ssl-config

# Monitor process execution
-a always,exit -F arch=b64 -S execve -k process-execution
-a always,exit -F arch=b32 -S execve -k process-execution

# Monitor network connections
-a always,exit -F arch=b64 -S socket -k network-socket
-a always,exit -F arch=b32 -S socket -k network-socket

# Monitor file permission changes
-a always,exit -F arch=b64 -S chmod,fchmod,fchmodat -k file-permissions
-a always,exit -F arch=b32 -S chmod,fchmod,fchmodat -k file-permissions

# Monitor privilege escalation
-w /usr/bin/sudo -p x -k privilege-escalation
-w /bin/su -p x -k privilege-escalation

# Monitor Docker daemon
-w /var/lib/docker/ -p wa -k docker-data
-w /usr/bin/docker -p x -k docker-exec
-w /usr/bin/dockerd -p x -k docker-daemon

# Make rules immutable
-e 2
EOF
    
    # Restart auditd to apply rules
    sudo systemctl restart auditd
    
    # Configure log rotation for audit logs
    cat << 'EOF' | sudo tee /etc/logrotate.d/bzzz-audit > /dev/null
/var/log/audit/*.log {
    daily
    rotate 30
    compress
    delaycompress
    missingok
    notifempty
    create 640 root adm
    postrotate
        /sbin/service auditd restart > /dev/null 2>&1 || true
    endscript
}
EOF
    
    success "Audit logging configured"
}

setup_intrusion_detection() {
    log "Setting up intrusion detection..."
    
    # Install fail2ban if not present
    if ! command -v fail2ban-server &> /dev/null; then
        sudo apt-get update
        sudo apt-get install -y fail2ban
    fi
    
    # Configure fail2ban for BZZZ v2
    cat << 'EOF' | sudo tee /etc/fail2ban/jail.d/bzzz-v2.conf > /dev/null
[DEFAULT]
bantime = 3600
findtime = 600
maxretry = 5
backend = systemd

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

[docker-auth]
enabled = true
port = 2376
filter = docker-auth
logpath = /var/log/audit/audit.log
maxretry = 3
bantime = 3600

[bzzz-p2p]
enabled = true
port = 9000:9300
filter = bzzz-p2p
logpath = /rust/bzzz-v2/logs/application/bzzz-agent.log
maxretry = 10
bantime = 1800

[traefik-auth]
enabled = true
port = http,https
filter = traefik-auth
logpath = /var/log/traefik/access.log
maxretry = 5
bantime = 3600
EOF
    
    # Create custom fail2ban filters
    cat << 'EOF' | sudo tee /etc/fail2ban/filter.d/docker-auth.conf > /dev/null
[Definition]
failregex = ^.*type=SYSCALL.*comm="dockerd".*res=failed.*$
ignoreregex =
EOF
    
    cat << 'EOF' | sudo tee /etc/fail2ban/filter.d/bzzz-p2p.conf > /dev/null
[Definition]
failregex = ^.*level=error.*msg="unauthorized connection attempt".*peer=<HOST>.*$
            ^.*level=warn.*msg="rate limit exceeded".*source=<HOST>.*$
ignoreregex =
EOF
    
    cat << 'EOF' | sudo tee /etc/fail2ban/filter.d/traefik-auth.conf > /dev/null
[Definition]
failregex = ^<HOST>.*"(GET|POST|PUT|DELETE).*" (401|403) .*$
ignoreregex =
EOF
    
    # Start and enable fail2ban
    sudo systemctl enable fail2ban
    sudo systemctl start fail2ban
    
    success "Intrusion detection configured"
}

configure_container_security() {
    log "Configuring container security policies..."
    
    # Create AppArmor profile for BZZZ containers
    cat << 'EOF' | sudo tee /etc/apparmor.d/bzzz-container > /dev/null
#include <tunables/global>

profile bzzz-container flags=(attach_disconnected,mediate_deleted) {
    #include <abstractions/base>
    
    capability,
    file,
    network,
    
    deny @{PROC}/* w,
    deny @{PROC}/sys/fs/** w,
    deny @{PROC}/sysrq-trigger rwklx,
    deny @{PROC}/mem rwklx,
    deny @{PROC}/kmem rwklx,
    deny @{PROC}/sys/kernel/[^s][^h][^m]* w,
    deny mount,
    deny /sys/[^f]** wklx,
    deny /sys/f[^s]** wklx,
    deny /sys/fs/[^c]** wklx,
    deny /sys/fs/c[^g]** wklx,
    deny /sys/fs/cg[^r]** wklx,
    deny /sys/firmware/** rwklx,
    deny /sys/kernel/security/** rwklx,
    
    # Allow access to application directories
    /app/** r,
    /app/bzzz rix,
    /data/** rw,
    /config/** r,
    
    # Allow temporary files
    /tmp/** rw,
    
    # Network access
    network inet,
    network inet6,
    network unix,
}
EOF
    
    # Load AppArmor profile
    sudo apparmor_parser -r /etc/apparmor.d/bzzz-container
    
    # Create seccomp profile for BZZZ containers
    mkdir -p /rust/bzzz-v2/config/security
    cat << 'EOF' > /rust/bzzz-v2/config/security/bzzz-seccomp.json
{
    "defaultAction": "SCMP_ACT_ERRNO",
    "architectures": [
        "SCMP_ARCH_X86_64",
        "SCMP_ARCH_X86",
        "SCMP_ARCH_X32"
    ],
    "syscalls": [
        {
            "names": [
                "accept", "access", "arch_prctl", "bind", "brk",
                "chdir", "chmod", "chown", "clone", "close",
                "connect", "dup", "dup2", "epoll_create", "epoll_ctl",
                "epoll_wait", "execve", "exit", "exit_group", "fcntl",
                "fstat", "futex", "getcwd", "getdents", "getgid",
                "getpid", "getppid", "gettid", "getuid", "listen",
                "lstat", "mmap", "mprotect", "munmap", "nanosleep",
                "open", "openat", "pipe", "poll", "prctl",
                "read", "readlink", "recv", "recvfrom", "rt_sigaction",
                "rt_sigprocmask", "rt_sigreturn", "sched_yield", "send",
                "sendto", "set_robust_list", "setsockopt", "socket",
                "stat", "write"
            ],
            "action": "SCMP_ACT_ALLOW"
        }
    ]
}
EOF
    
    success "Container security policies configured"
}

main() {
    log "Starting BZZZ v2 security hardening..."
    
    check_root
    configure_firewall
    configure_docker_security
    setup_tls_certificates
    configure_secrets_management
    setup_network_security
    configure_audit_logging
    setup_intrusion_detection
    configure_container_security
    
    success "BZZZ v2 security hardening completed successfully!"
    log "Security configuration saved to: $LOG_FILE"
    log "Review firewall rules: sudo ufw status verbose"
    log "Check fail2ban status: sudo fail2ban-client status"
    log "Verify audit rules: sudo auditctl -l"
}

# Execute main function
main "$@"