#!/bin/bash
################################################################################
# Script Name: install-fail2ban.sh
# Version: 1.0
# Description: Automated Fail2ban installation with multi-jail configuration,
#              custom ban actions, allowlists, Prometheus metrics integration,
#              and service auto-detection for Debian/Ubuntu and RHEL/Rocky/Alma
#
# Author:  Phil Connor
# Contact: contact@mylinux.work
# Website: https://mylinux.work
# License: MIT
#
# Usage:
#   sudo ./install-fail2ban.sh
#   sudo ./install-fail2ban.sh --jails "sshd,nginx-http-auth,postfix"
#   sudo ./install-fail2ban.sh --allowlist "10.0.0.0/8,192.168.1.0/24"
#   sudo ./install-fail2ban.sh --ban-action nftables --bantime 3600
#   sudo ./install-fail2ban.sh --dry-run
#   sudo ./install-fail2ban.sh --uninstall
#
################################################################################

set -euo pipefail

# ============================================================================
# DEFAULTS
# ============================================================================

readonly VERSION="1.0"
readonly SCRIPT_NAME="${0##*/}"
readonly LOG_FILE="/var/log/fail2ban-install.log"

JAILS=""
ALLOWLIST="127.0.0.1/8 ::1"
BAN_ACTION="auto"
BANTIME="3600"
FINDTIME="600"
MAXRETRY="5"
BACKEND="auto"
DRY_RUN=false
UNINSTALL=false
PROMETHEUS=false
RECIDIVE=true

# OS detection
OS_ID=""
OS_VERSION=""
PKG_MGR=""

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

# ============================================================================
# HELPER FUNCTIONS
# ============================================================================

log_info()  { echo -e "${GREEN}[INFO]${NC} $*"; echo "[INFO] $(date '+%Y-%m-%d %H:%M:%S') $*" >> "$LOG_FILE" 2>/dev/null || true; }
log_warn()  { echo -e "${YELLOW}[WARN]${NC} $*"; echo "[WARN] $(date '+%Y-%m-%d %H:%M:%S') $*" >> "$LOG_FILE" 2>/dev/null || true; }
log_error() { echo -e "${RED}[ERROR]${NC} $*" >&2; echo "[ERROR] $(date '+%Y-%m-%d %H:%M:%S') $*" >> "$LOG_FILE" 2>/dev/null || true; }
log_step()  { echo -e "${CYAN}[STEP]${NC} $*"; echo "[STEP] $(date '+%Y-%m-%d %H:%M:%S') $*" >> "$LOG_FILE" 2>/dev/null || true; }

show_usage() {
    cat <<EOF
Usage: sudo $SCRIPT_NAME [OPTIONS]

Installs and configures Fail2ban with multi-jail support.

OPTIONS:
    --jails LIST        Comma-separated jails (default: auto-detect services)
    --allowlist IPS     Space or comma-separated trusted IPs/CIDRs to ignore
    --ban-action TYPE   Ban action: iptables, nftables, or auto (default: auto)
    --bantime SECS      Ban duration in seconds (default: 3600)
    --findtime SECS     Window for counting failures (default: 600)
    --maxretry NUM      Max failures before ban (default: 5)
    --no-recidive       Disable the recidive jail for repeat offenders
    --prometheus        Enable Prometheus metrics via fail2ban-exporter
    --dry-run           Show what would be done without making changes
    --uninstall         Remove Fail2ban and configuration
    -h, --help          Show this help message

SUPPORTED JAILS:
    sshd                SSH brute force (always enabled)
    nginx-http-auth     Nginx HTTP basic auth failures
    nginx-botsearch     Nginx bot/scanner blocking
    nginx-limit-req     Nginx rate limit violations
    apache-auth         Apache authentication failures
    apache-badbots      Apache bad bot blocking
    postfix             Postfix SMTP auth failures
    dovecot             Dovecot IMAP/POP3 auth failures
    mysqld-auth         MySQL/MariaDB auth failures
    named-refused       BIND DNS refused queries
    recidive            Repeat offender escalation

EXAMPLES:
    # Auto-detect and install
    sudo $SCRIPT_NAME

    # Specific jails with custom settings
    sudo $SCRIPT_NAME --jails "sshd,nginx-http-auth,postfix" --bantime 7200

    # With allowlist and nftables
    sudo $SCRIPT_NAME --allowlist "10.0.0.0/8" --ban-action nftables

EOF
    exit 0
}

parse_args() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --jails)           JAILS="$2"; shift 2 ;;
            --allowlist)       ALLOWLIST="$ALLOWLIST ${2//,/ }"; shift 2 ;;
            --ban-action)      BAN_ACTION="$2"; shift 2 ;;
            --bantime)         BANTIME="$2"; shift 2 ;;
            --findtime)        FINDTIME="$2"; shift 2 ;;
            --maxretry)        MAXRETRY="$2"; shift 2 ;;
            --no-recidive)     RECIDIVE=false; shift ;;
            --prometheus)      PROMETHEUS=true; shift ;;
            --dry-run)         DRY_RUN=true; shift ;;
            --uninstall)       UNINSTALL=true; shift ;;
            -h|--help)         show_usage ;;
            *)                 log_error "Unknown option: $1"; show_usage ;;
        esac
    done
}

check_root() {
    if [[ $EUID -ne 0 ]]; then
        log_error "This script must be run as root (sudo)"
        exit 1
    fi
}

# ============================================================================
# OS DETECTION
# ============================================================================

detect_os() {
    if [ -f /etc/os-release ]; then
        . /etc/os-release
        OS_ID="${ID}"
        OS_VERSION="${VERSION_ID}"
    else
        log_error "Cannot detect OS — /etc/os-release not found"
        exit 1
    fi

    case "$OS_ID" in
        ubuntu|debian)
            PKG_MGR="apt"
            ;;
        rhel|rocky|almalinux|centos|fedora)
            PKG_MGR="dnf"
            if ! command -v dnf &>/dev/null; then
                PKG_MGR="yum"
            fi
            ;;
        *)
            log_error "Unsupported OS: $OS_ID"
            exit 1
            ;;
    esac

    log_info "Detected OS: $OS_ID $OS_VERSION (package manager: $PKG_MGR)"
}

# ============================================================================
# SERVICE DETECTION
# ============================================================================

detect_services() {
    if [ -n "$JAILS" ]; then
        log_info "Using specified jails: $JAILS"
        return
    fi

    log_step "Auto-detecting services for jail configuration..."
    local detected=()

    # SSH is always enabled
    detected+=("sshd")

    # Nginx
    if systemctl is-active --quiet nginx 2>/dev/null || [ -f /etc/nginx/nginx.conf ]; then
        detected+=("nginx-http-auth")
        detected+=("nginx-botsearch")
        log_info "  Detected: nginx"
    fi

    # Apache
    if systemctl is-active --quiet httpd 2>/dev/null || systemctl is-active --quiet apache2 2>/dev/null; then
        detected+=("apache-auth")
        detected+=("apache-badbots")
        log_info "  Detected: apache"
    fi

    # Postfix
    if systemctl is-active --quiet postfix 2>/dev/null; then
        detected+=("postfix")
        log_info "  Detected: postfix"
    fi

    # Dovecot
    if systemctl is-active --quiet dovecot 2>/dev/null; then
        detected+=("dovecot")
        log_info "  Detected: dovecot"
    fi

    # MySQL/MariaDB
    if systemctl is-active --quiet mysqld 2>/dev/null || systemctl is-active --quiet mariadb 2>/dev/null; then
        detected+=("mysqld-auth")
        log_info "  Detected: mysql/mariadb"
    fi

    # BIND
    if systemctl is-active --quiet named 2>/dev/null || systemctl is-active --quiet bind9 2>/dev/null; then
        detected+=("named-refused")
        log_info "  Detected: bind/named"
    fi

    JAILS=$(IFS=','; echo "${detected[*]}")
    log_info "Auto-detected jails: $JAILS"
}

# ============================================================================
# BAN ACTION DETECTION
# ============================================================================

detect_ban_action() {
    if [ "$BAN_ACTION" != "auto" ]; then
        log_info "Using specified ban action: $BAN_ACTION"
        return
    fi

    if command -v nft &>/dev/null && nft list ruleset &>/dev/null; then
        BAN_ACTION="nftables-multiport"
        log_info "Detected nftables — using nftables-multiport action"
    elif command -v iptables &>/dev/null; then
        BAN_ACTION="iptables-multiport"
        log_info "Detected iptables — using iptables-multiport action"
    else
        BAN_ACTION="iptables-multiport"
        log_warn "No firewall detected — defaulting to iptables-multiport"
    fi
}

# ============================================================================
# INSTALLATION
# ============================================================================

install_fail2ban() {
    log_step "Installing Fail2ban..."

    if $DRY_RUN; then
        log_info "[DRY RUN] Would install fail2ban via $PKG_MGR"
        return
    fi

    case "$PKG_MGR" in
        apt)
            apt-get update -qq
            apt-get install -y -qq fail2ban
            ;;
        dnf|yum)
            $PKG_MGR install -y -q epel-release 2>/dev/null || true
            $PKG_MGR install -y -q fail2ban fail2ban-firewalld 2>/dev/null || \
            $PKG_MGR install -y -q fail2ban
            ;;
    esac

    log_info "Fail2ban installed successfully"
}

# ============================================================================
# CONFIGURATION
# ============================================================================

configure_fail2ban() {
    log_step "Configuring Fail2ban..."

    local jail_local="/etc/fail2ban/jail.local"

    if $DRY_RUN; then
        log_info "[DRY RUN] Would write configuration to $jail_local"
        show_config
        return
    fi

    # Backup existing config
    if [ -f "$jail_local" ]; then
        cp "$jail_local" "${jail_local}.bak.$(date +%s)"
        log_info "Backed up existing $jail_local"
    fi

    # Write jail.local
    cat > "$jail_local" <<JAILEOF
# Fail2ban configuration — generated by $SCRIPT_NAME v$VERSION
# $(date)

[DEFAULT]
bantime  = $BANTIME
findtime = $FINDTIME
maxretry = $MAXRETRY
banaction = $BAN_ACTION
banaction_allports = $BAN_ACTION
backend  = $BACKEND
ignoreip = $ALLOWLIST

# Send alerts (uncomment and configure)
# destemail = admin@example.com
# sender    = fail2ban@example.com
# mta       = sendmail
# action    = %(action_mwl)s

JAILEOF

    # Generate jail sections
    IFS=',' read -ra JAIL_ARRAY <<< "$JAILS"
    for jail in "${JAIL_ARRAY[@]}"; do
        jail=$(echo "$jail" | xargs)  # trim whitespace
        write_jail_config "$jail" >> "$jail_local"
    done

    # Recidive jail
    if $RECIDIVE; then
        cat >> "$jail_local" <<'RECIDIVEEOF'

[recidive]
enabled  = true
logpath  = /var/log/fail2ban.log
bantime  = 604800
findtime = 86400
maxretry = 3
RECIDIVEEOF
        log_info "Recidive jail enabled (repeat offenders banned for 7 days)"
    fi

    log_info "Configuration written to $jail_local"
}

write_jail_config() {
    local jail="$1"

    case "$jail" in
        sshd)
            cat <<'EOF'

[sshd]
enabled  = true
port     = ssh
logpath  = %(sshd_log)s
maxretry = 5
EOF
            ;;
        nginx-http-auth)
            cat <<'EOF'

[nginx-http-auth]
enabled  = true
port     = http,https
logpath  = /var/log/nginx/error.log
maxretry = 5
EOF
            ;;
        nginx-botsearch)
            cat <<'EOF'

[nginx-botsearch]
enabled  = true
port     = http,https
logpath  = /var/log/nginx/access.log
maxretry = 2
EOF
            ;;
        nginx-limit-req)
            cat <<'EOF'

[nginx-limit-req]
enabled  = true
port     = http,https
logpath  = /var/log/nginx/error.log
maxretry = 5
EOF
            ;;
        apache-auth)
            cat <<'EOF'

[apache-auth]
enabled  = true
port     = http,https
logpath  = %(apache_error_log)s
maxretry = 5
EOF
            ;;
        apache-badbots)
            cat <<'EOF'

[apache-badbots]
enabled  = true
port     = http,https
logpath  = %(apache_access_log)s
maxretry = 2
bantime  = 86400
EOF
            ;;
        postfix)
            cat <<'EOF'

[postfix]
enabled  = true
port     = smtp,465,submission
logpath  = /var/log/mail.log
maxretry = 5
EOF
            ;;
        dovecot)
            cat <<'EOF'

[dovecot]
enabled  = true
port     = pop3,pop3s,imap,imaps
logpath  = /var/log/mail.log
maxretry = 5
EOF
            ;;
        mysqld-auth)
            cat <<'EOF'

[mysqld-auth]
enabled  = true
port     = 3306
logpath  = /var/log/mysql/error.log
maxretry = 5
EOF
            ;;
        named-refused)
            cat <<'EOF'

[named-refused]
enabled  = true
port     = domain,953
logpath  = /var/log/named/security.log
maxretry = 5
EOF
            ;;
        *)
            log_warn "Unknown jail: $jail — skipping"
            ;;
    esac
}

show_config() {
    echo ""
    echo "=== Generated Configuration ==="
    echo "[DEFAULT]"
    echo "bantime  = $BANTIME"
    echo "findtime = $FINDTIME"
    echo "maxretry = $MAXRETRY"
    echo "banaction = $BAN_ACTION"
    echo "ignoreip = $ALLOWLIST"
    echo ""
    IFS=',' read -ra JAIL_ARRAY <<< "$JAILS"
    for jail in "${JAIL_ARRAY[@]}"; do
        jail=$(echo "$jail" | xargs)
        echo "[$jail] = enabled"
    done
    if $RECIDIVE; then
        echo "[recidive] = enabled (7-day ban for repeat offenders)"
    fi
    echo "================================"
}

# ============================================================================
# PROMETHEUS INTEGRATION
# ============================================================================

setup_prometheus() {
    if ! $PROMETHEUS; then
        return
    fi

    log_step "Setting up Prometheus metrics integration..."

    if $DRY_RUN; then
        log_info "[DRY RUN] Would download fail2ban-exporter.sh"
        return
    fi

    if [ -f /usr/local/bin/fail2ban-exporter.sh ]; then
        log_info "fail2ban-exporter.sh already exists"
        return
    fi

    local exporter_url="https://mylinux.work/downloads/fail2ban-exporter.sh"
    if curl -fsSL "$exporter_url" -o /usr/local/bin/fail2ban-exporter.sh 2>/dev/null; then
        chmod +x /usr/local/bin/fail2ban-exporter.sh
        log_info "Downloaded fail2ban-exporter.sh"

        # Add cron job
        if ! crontab -l 2>/dev/null | grep -q "fail2ban-exporter"; then
            (crontab -l 2>/dev/null; echo "*/2 * * * * /usr/local/bin/fail2ban-exporter.sh --textfile 2>/dev/null") | crontab -
            log_info "Added cron job for fail2ban-exporter"
        fi
    else
        log_warn "Could not download fail2ban-exporter.sh — configure manually"
    fi
}

# ============================================================================
# SERVICE MANAGEMENT
# ============================================================================

start_fail2ban() {
    log_step "Starting Fail2ban service..."

    if $DRY_RUN; then
        log_info "[DRY RUN] Would enable and start fail2ban"
        return
    fi

    systemctl enable fail2ban
    systemctl restart fail2ban
    sleep 2

    if systemctl is-active --quiet fail2ban; then
        log_info "Fail2ban is running"
    else
        log_error "Fail2ban failed to start — check journalctl -u fail2ban"
        exit 1
    fi
}

# ============================================================================
# VERIFICATION
# ============================================================================

verify_installation() {
    log_step "Verifying installation..."

    if $DRY_RUN; then
        log_info "[DRY RUN] Would verify installation"
        return
    fi

    echo ""
    echo "=== Fail2ban Status ==="
    fail2ban-client status
    echo ""

    IFS=',' read -ra JAIL_ARRAY <<< "$JAILS"
    for jail in "${JAIL_ARRAY[@]}"; do
        jail=$(echo "$jail" | xargs)
        echo "--- $jail ---"
        fail2ban-client status "$jail" 2>/dev/null || echo "  (jail not active — check logs)"
        echo ""
    done

    if $RECIDIVE; then
        echo "--- recidive ---"
        fail2ban-client status recidive 2>/dev/null || echo "  (recidive not active)"
        echo ""
    fi

    echo "=== Configuration ==="
    echo "Ban action:  $BAN_ACTION"
    echo "Ban time:    ${BANTIME}s ($(( BANTIME / 60 )) minutes)"
    echo "Find time:   ${FINDTIME}s ($(( FINDTIME / 60 )) minutes)"
    echo "Max retry:   $MAXRETRY"
    echo "Ignore list: $ALLOWLIST"
    echo ""
    log_info "Installation complete. Logs: $LOG_FILE"
}

# ============================================================================
# UNINSTALL
# ============================================================================

uninstall_fail2ban() {
    log_step "Uninstalling Fail2ban..."

    if $DRY_RUN; then
        log_info "[DRY RUN] Would stop and remove fail2ban"
        return
    fi

    systemctl stop fail2ban 2>/dev/null || true
    systemctl disable fail2ban 2>/dev/null || true

    case "$PKG_MGR" in
        apt)
            apt-get purge -y -qq fail2ban
            apt-get autoremove -y -qq
            ;;
        dnf|yum)
            $PKG_MGR remove -y -q fail2ban
            ;;
    esac

    # Clean up config
    if [ -d /etc/fail2ban ]; then
        log_info "Configuration directory /etc/fail2ban preserved — remove manually if desired"
    fi

    log_info "Fail2ban uninstalled"
    exit 0
}

# ============================================================================
# MAIN
# ============================================================================

main() {
    parse_args "$@"

    echo ""
    echo "============================================"
    echo "  Fail2ban Install Script v${VERSION}"
    echo "  https://mylinux.work"
    echo "============================================"
    echo ""

    check_root
    detect_os

    if $UNINSTALL; then
        uninstall_fail2ban
    fi

    detect_services
    detect_ban_action
    install_fail2ban
    configure_fail2ban
    setup_prometheus
    start_fail2ban
    verify_installation
}

main "$@"