#!/bin/bash
# Foxhunt HFT Trading System - Production Security Hardening Script
# This script implements critical security measures for production deployment

set -euo pipefail

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

# Logging function
log() {
    echo -e "${BLUE}[$(date +'%Y-%m-%d %H:%M:%S')] $1${NC}" | tee -a /var/log/foxhunt-security.log
}

error() {
    echo -e "${RED}[ERROR] $1${NC}" | tee -a /var/log/foxhunt-security.log >&2
}

success() {
    echo -e "${GREEN}[SUCCESS] $1${NC}" | tee -a /var/log/foxhunt-security.log
}

warning() {
    echo -e "${YELLOW}[WARNING] $1${NC}" | tee -a /var/log/foxhunt-security.log
}

# Check if running as root
check_root() {
    if [[ $EUID -ne 0 ]]; then
        error "This script must be run as root for security hardening"
        exit 1
    fi
}

# Generate strong secrets
generate_secrets() {
    log "Generating production secrets..."
    
    local secrets_dir="/opt/foxhunt/secrets"
    mkdir -p "$secrets_dir"
    chmod 700 "$secrets_dir"
    
    # JWT secret (256-bit)
    if [[ ! -f "$secrets_dir/jwt_secret" ]]; then
        openssl rand -hex 32 > "$secrets_dir/jwt_secret"
        chmod 600 "$secrets_dir/jwt_secret"
        success "JWT secret generated"
    fi
    
    # Database password (128-bit)
    if [[ ! -f "$secrets_dir/db_password" ]]; then
        openssl rand -base64 32 > "$secrets_dir/db_password"
        chmod 600 "$secrets_dir/db_password"
        success "Database password generated"
    fi
    
    # Redis password
    if [[ ! -f "$secrets_dir/redis_password" ]]; then
        openssl rand -base64 32 > "$secrets_dir/redis_password"
        chmod 600 "$secrets_dir/redis_password"
        success "Redis password generated"
    fi
    
    # API encryption key
    if [[ ! -f "$secrets_dir/api_key" ]]; then
        openssl rand -hex 32 > "$secrets_dir/api_key"
        chmod 600 "$secrets_dir/api_key"
        success "API encryption key generated"
    fi
    
    chown -R foxhunt:foxhunt "$secrets_dir"
}

# Generate TLS certificates
generate_tls_certificates() {
    log "Generating TLS certificates..."
    
    local certs_dir="/opt/foxhunt/certs"
    mkdir -p "$certs_dir"
    chmod 755 "$certs_dir"
    
    # Generate CA key and certificate
    if [[ ! -f "$certs_dir/ca-key.pem" ]]; then
        openssl genrsa -out "$certs_dir/ca-key.pem" 4096
        chmod 600 "$certs_dir/ca-key.pem"
        
        openssl req -new -x509 -days 365 -key "$certs_dir/ca-key.pem" -out "$certs_dir/ca-cert.pem" \
            -subj "/C=US/ST=NY/L=NYC/O=Foxhunt/OU=Trading/CN=Foxhunt-CA"
        
        success "CA certificate generated"
    fi
    
    # Generate server key and certificate
    if [[ ! -f "$certs_dir/server-key.pem" ]]; then
        openssl genrsa -out "$certs_dir/server-key.pem" 4096
        chmod 600 "$certs_dir/server-key.pem"
        
        openssl req -new -key "$certs_dir/server-key.pem" -out "$certs_dir/server.csr" \
            -subj "/C=US/ST=NY/L=NYC/O=Foxhunt/OU=Trading/CN=trading.foxhunt.internal"
        
        openssl x509 -req -days 365 -in "$certs_dir/server.csr" \
            -CA "$certs_dir/ca-cert.pem" -CAkey "$certs_dir/ca-key.pem" -CAcreateserial \
            -out "$certs_dir/server-cert.pem"
        
        rm "$certs_dir/server.csr"
        success "Server certificate generated"
    fi
    
    # Generate client certificate for mTLS
    if [[ ! -f "$certs_dir/client-key.pem" ]]; then
        openssl genrsa -out "$certs_dir/client-key.pem" 4096
        chmod 600 "$certs_dir/client-key.pem"
        
        openssl req -new -key "$certs_dir/client-key.pem" -out "$certs_dir/client.csr" \
            -subj "/C=US/ST=NY/L=NYC/O=Foxhunt/OU=Trading/CN=client.foxhunt.internal"
        
        openssl x509 -req -days 365 -in "$certs_dir/client.csr" \
            -CA "$certs_dir/ca-cert.pem" -CAkey "$certs_dir/ca-key.pem" -CAcreateserial \
            -out "$certs_dir/client-cert.pem"
        
        rm "$certs_dir/client.csr"
        success "Client certificate generated"
    fi
    
    chown -R foxhunt:foxhunt "$certs_dir"
}

# Create secure environment file
create_secure_env() {
    log "Creating secure environment configuration..."
    
    local env_file="/opt/foxhunt/config/production.env"
    local secrets_dir="/opt/foxhunt/secrets"
    
    cat > "$env_file" << EOF
# Foxhunt HFT Production Environment Configuration
# CRITICAL: This file contains sensitive production secrets
# File permissions: 600 (owner read/write only)

# Database Configuration
DATABASE_URL="postgresql://foxhunt:\$(cat $secrets_dir/db_password)@localhost:5432/foxhunt_prod?sslmode=require"
POSTGRES_PASSWORD_FILE="$secrets_dir/db_password"

# JWT Configuration
JWT_SECRET_FILE="$secrets_dir/jwt_secret"
JWT_ISSUER="foxhunt-production"
JWT_AUDIENCE="trading-api-prod"

# Redis Configuration
REDIS_URL="redis://:\$(cat $secrets_dir/redis_password)@localhost:6379/0"
REDIS_PASSWORD_FILE="$secrets_dir/redis_password"

# TLS Configuration
TLS_CERT_FILE="/opt/foxhunt/certs/server-cert.pem"
TLS_KEY_FILE="/opt/foxhunt/certs/server-key.pem"
TLS_CA_FILE="/opt/foxhunt/certs/ca-cert.pem"
TLS_CLIENT_CERT_FILE="/opt/foxhunt/certs/client-cert.pem"
TLS_CLIENT_KEY_FILE="/opt/foxhunt/certs/client-key.pem"

# API Security
API_ENCRYPTION_KEY_FILE="$secrets_dir/api_key"

# Vault Configuration (if using Vault)
VAULT_ADDR="https://vault.foxhunt.internal:8200"
VAULT_ROLE_ID_FILE="/opt/foxhunt/vault/role-id"
VAULT_SECRET_ID_FILE="/opt/foxhunt/vault/secret-id"

# Security Settings
REQUIRE_MTLS="true"
ENABLE_RATE_LIMITING="true"
MAX_REQUESTS_PER_MINUTE="60"
MAX_FAILED_ATTEMPTS="10"
LOCKOUT_DURATION_SECONDS="900"

# Audit and Compliance
ENABLE_AUDIT_LOGGING="true"
AUDIT_LOG_LEVEL="INFO"
SOX_COMPLIANCE_ENABLED="true"
MIFID_COMPLIANCE_ENABLED="true"

# Risk Management
RISK_ENGINE_ENABLED="true"
KILL_SWITCH_ENABLED="true"
POSITION_LIMITS_ENABLED="true"
CIRCUIT_BREAKER_ENABLED="true"
EOF

    chmod 600 "$env_file"
    chown foxhunt:foxhunt "$env_file"
    success "Secure environment configuration created"
}

# Configure firewall rules
configure_firewall() {
    log "Configuring firewall rules..."
    
    # Enable UFW if not already enabled
    ufw --force enable
    
    # Default policies
    ufw default deny incoming
    ufw default allow outgoing
    
    # Allow SSH (modify port as needed)
    ufw allow 22/tcp
    
    # Allow HTTPS for management interfaces
    ufw allow 443/tcp
    
    # Allow gRPC ports (restrict to specific IPs in production)
    ufw allow 50051/tcp comment "Trading Service gRPC"
    ufw allow 50052/tcp comment "ML Training Service gRPC"
    ufw allow 50053/tcp comment "Backtesting Service gRPC"
    
    # Allow database connections (restrict to localhost)
    ufw allow from 127.0.0.1 to any port 5432 comment "PostgreSQL local"
    ufw allow from 127.0.0.1 to any port 6379 comment "Redis local"
    
    # Rate limiting for SSH
    ufw limit ssh
    
    success "Firewall configured"
}

# Set up fail2ban
setup_fail2ban() {
    log "Setting up fail2ban..."
    
    # Install fail2ban if not present
    if ! command -v fail2ban-server &> /dev/null; then
        apt-get update
        apt-get install -y fail2ban
    fi
    
    # Create custom fail2ban configuration for Foxhunt
    cat > /etc/fail2ban/jail.d/foxhunt.conf << EOF
[foxhunt-auth]
enabled = true
port = 50051,50052,50053
filter = foxhunt-auth
logpath = /var/log/foxhunt/auth.log
maxretry = 5
bantime = 1800
findtime = 600

[foxhunt-rate-limit]
enabled = true
port = 50051,50052,50053
filter = foxhunt-rate-limit
logpath = /var/log/foxhunt/security.log
maxretry = 10
bantime = 3600
findtime = 300
EOF

    # Create filter for authentication failures
    cat > /etc/fail2ban/filter.d/foxhunt-auth.conf << EOF
[Definition]
failregex = ^.*AUTH_FAILURE.*client_ip=<HOST>.*$
ignoreregex =
EOF

    # Create filter for rate limiting
    cat > /etc/fail2ban/filter.d/foxhunt-rate-limit.conf << EOF
[Definition]
failregex = ^.*Rate limit exceeded for IP <HOST>.*$
ignoreregex =
EOF

    systemctl enable fail2ban
    systemctl restart fail2ban
    success "Fail2ban configured"
}

# Configure system security settings
configure_system_security() {
    log "Configuring system security settings..."
    
    # Kernel parameters for security
    cat >> /etc/sysctl.d/99-foxhunt-security.conf << EOF
# Network security
net.ipv4.conf.default.rp_filter=1
net.ipv4.conf.all.rp_filter=1
net.ipv4.conf.all.accept_redirects=0
net.ipv4.conf.default.accept_redirects=0
net.ipv4.conf.all.secure_redirects=0
net.ipv4.conf.default.secure_redirects=0
net.ipv6.conf.all.accept_redirects=0
net.ipv6.conf.default.accept_redirects=0
net.ipv4.conf.all.send_redirects=0
net.ipv4.conf.default.send_redirects=0
net.ipv4.ip_forward=0
net.ipv6.conf.all.forwarding=0
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

# Protect against SYN flood attacks
net.ipv4.tcp_syncookies=1
net.ipv4.tcp_max_syn_backlog=2048
net.ipv4.tcp_synack_retries=2
net.ipv4.tcp_syn_retries=5

# IP Spoofing protection
net.ipv4.conf.all.log_martians=1
net.ipv4.conf.default.log_martians=1

# Ignore ping requests
net.ipv4.icmp_echo_ignore_all=1

# Memory protection
kernel.dmesg_restrict=1
kernel.kptr_restrict=2
kernel.yama.ptrace_scope=1

# File system security
fs.protected_hardlinks=1
fs.protected_symlinks=1
fs.suid_dumpable=0
EOF

    sysctl -p /etc/sysctl.d/99-foxhunt-security.conf
    success "System security settings configured"
}

# Set up log monitoring
setup_log_monitoring() {
    log "Setting up log monitoring..."
    
    local log_dir="/var/log/foxhunt"
    mkdir -p "$log_dir"
    chown foxhunt:foxhunt "$log_dir"
    chmod 750 "$log_dir"
    
    # Logrotate configuration
    cat > /etc/logrotate.d/foxhunt << EOF
/var/log/foxhunt/*.log {
    daily
    missingok
    rotate 90
    compress
    delaycompress
    notifempty
    create 0640 foxhunt foxhunt
    postrotate
        /bin/systemctl reload foxhunt-* 2>/dev/null || true
    endscript
}
EOF

    success "Log monitoring configured"
}

# Verify security configuration
verify_security() {
    log "Verifying security configuration..."
    
    local errors=0
    
    # Check file permissions
    if [[ $(stat -c %a /opt/foxhunt/secrets) != "700" ]]; then
        error "Secrets directory has incorrect permissions"
        ((errors++))
    fi
    
    if [[ $(stat -c %a /opt/foxhunt/config/production.env) != "600" ]]; then
        error "Production environment file has incorrect permissions"
        ((errors++))
    fi
    
    # Check certificate validity
    if ! openssl verify -CAfile /opt/foxhunt/certs/ca-cert.pem /opt/foxhunt/certs/server-cert.pem &>/dev/null; then
        error "Server certificate verification failed"
        ((errors++))
    fi
    
    # Check services
    if ! systemctl is-active --quiet ufw; then
        error "UFW firewall is not active"
        ((errors++))
    fi
    
    if ! systemctl is-active --quiet fail2ban; then
        error "Fail2ban is not active"
        ((errors++))
    fi
    
    if [[ $errors -eq 0 ]]; then
        success "Security configuration verified successfully"
    else
        error "Security verification failed with $errors errors"
        return 1
    fi
}

# Main execution
main() {
    log "Starting Foxhunt HFT security hardening..."
    
    check_root
    generate_secrets
    generate_tls_certificates
    create_secure_env
    configure_firewall
    setup_fail2ban
    configure_system_security
    setup_log_monitoring
    verify_security
    
    success "Security hardening completed successfully!"
    warning "Please review the generated certificates and secrets before production deployment"
    warning "Update your application configuration to use the new environment variables"
}

# Run main function
main "$@"