#!/bin/bash
################################################################################
# Script Name: crowdsec-install.sh
# Version: 1.0
# Description: Automated CrowdSec installation with firewall bouncer,
#              collection selection, allowlists, and Prometheus integration
#              on Debian/Ubuntu and RHEL/Rocky/AlmaLinux
#
# Author:  Phil Connor
# Contact: contact@mylinux.work
# Website: https://mylinux.work
# License: MIT
#
# Usage:
#   sudo ./crowdsec-install.sh
#   sudo ./crowdsec-install.sh --collections "sshd,nginx"
#   sudo ./crowdsec-install.sh --allowlist "10.0.0.0/8" --prometheus
#   sudo ./crowdsec-install.sh --dry-run
#
################################################################################

set -euo pipefail

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

COLLECTIONS=""
ALLOWLIST=""
BOUNCER_TYPE="iptables"
PROMETHEUS=false
ENROLL_KEY=""
NO_BOUNCER=false
DRY_RUN=false

# 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} $*"; }
log_warn()  { echo -e "${YELLOW}[WARN]${NC} $*"; }
log_error() { echo -e "${RED}[ERROR]${NC} $*" >&2; }
log_step()  { echo -e "${CYAN}[STEP]${NC} $*"; }

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

Installs and configures CrowdSec with firewall bouncer and collections.

OPTIONS:
    --collections LIST  Comma-separated collections (default: auto-detect)
    --allowlist IPS     Comma-separated trusted IP ranges to whitelist
    --bouncer TYPE      Bouncer type: iptables or nftables (default: iptables)
    --prometheus        Enable Prometheus metrics on :6060
    --enroll KEY        CrowdSec console enrollment key
    --no-bouncer        Install agent only, skip bouncer
    --dry-run           Show what would be done
    -h, --help          Show this help

EXAMPLES:
    sudo $0
    sudo $0 --collections "sshd,nginx,postfix"
    sudo $0 --allowlist "10.0.0.0/8,192.168.1.0/24" --prometheus
    sudo $0 --bouncer nftables --enroll "abc123"

EOF
    exit 0
}

parse_args() {
    while [[ $# -gt 0 ]]; do
        case $1 in
            -h|--help) show_usage ;;
            --collections) COLLECTIONS="$2"; shift 2 ;;
            --allowlist) ALLOWLIST="$2"; shift 2 ;;
            --bouncer) BOUNCER_TYPE="$2"; shift 2 ;;
            --prometheus) PROMETHEUS=true; shift ;;
            --enroll) ENROLL_KEY="$2"; shift 2 ;;
            --no-bouncer) NO_BOUNCER=true; shift ;;
            --dry-run) DRY_RUN=true; shift ;;
            *) log_error "Unknown option: $1"; exit 1 ;;
        esac
    done
}

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

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

detect_os() {
    if [ -f /etc/os-release ]; then
        . /etc/os-release
        OS_FAMILY=""
        case "$ID" in
            debian|ubuntu) OS_FAMILY="debian" ;;
            rhel|centos|rocky|almalinux|fedora) OS_FAMILY="rhel" ;;
            *)
                if [ -n "${ID_LIKE:-}" ]; then
                    case "$ID_LIKE" in
                        *debian*|*ubuntu*) OS_FAMILY="debian" ;;
                        *rhel*|*centos*|*fedora*) OS_FAMILY="rhel" ;;
                    esac
                fi
                ;;
        esac

        if [ -z "$OS_FAMILY" ]; then
            log_error "Unsupported OS: $ID"
            exit 1
        fi

        log_info "Detected OS: $PRETTY_NAME (family: $OS_FAMILY)"
    else
        log_error "Cannot detect OS — /etc/os-release not found"
        exit 1
    fi
}

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

detect_services() {
    log_step "Detecting running services..."
    local detected=()

    # Always include base Linux collection
    detected+=("crowdsecurity/linux")

    # SSH
    if systemctl is-active --quiet sshd 2>/dev/null || systemctl is-active --quiet ssh 2>/dev/null; then
        detected+=("crowdsecurity/sshd")
        log_info "  Detected: SSH"
    fi

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

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

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

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

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

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

    DETECTED_COLLECTIONS="${detected[*]}"
}

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

add_repo_debian() {
    log_step "Adding CrowdSec repository (Debian/Ubuntu)..."

    if [ "$DRY_RUN" = true ]; then
        log_info "[DRY RUN] Would add CrowdSec apt repository"
        return
    fi

    apt-get update -qq
    apt-get install -y -qq curl gnupg apt-transport-https >/dev/null 2>&1

    curl -s https://packagecloud.io/install/repositories/crowdsec/crowdsec/script.deb.sh | bash >/dev/null 2>&1
}

add_repo_rhel() {
    log_step "Adding CrowdSec repository (RHEL/Rocky)..."

    if [ "$DRY_RUN" = true ]; then
        log_info "[DRY RUN] Would add CrowdSec yum repository"
        return
    fi

    curl -s https://packagecloud.io/install/repositories/crowdsec/crowdsec/script.rpm.sh | bash >/dev/null 2>&1
}

install_agent() {
    log_step "Installing CrowdSec agent..."

    if [ "$DRY_RUN" = true ]; then
        log_info "[DRY RUN] Would install crowdsec package"
        return
    fi

    case "$OS_FAMILY" in
        debian) apt-get install -y -qq crowdsec >/dev/null 2>&1 ;;
        rhel) dnf install -y -q crowdsec >/dev/null 2>&1 ;;
    esac

    systemctl enable --now crowdsec >/dev/null 2>&1
    log_info "CrowdSec agent installed and running"
}

install_bouncer() {
    if [ "$NO_BOUNCER" = true ]; then
        log_info "Skipping bouncer installation (--no-bouncer)"
        return
    fi

    local pkg="crowdsec-firewall-bouncer-${BOUNCER_TYPE}"
    log_step "Installing bouncer: $pkg..."

    if [ "$DRY_RUN" = true ]; then
        log_info "[DRY RUN] Would install $pkg"
        return
    fi

    case "$OS_FAMILY" in
        debian) apt-get install -y -qq "$pkg" >/dev/null 2>&1 ;;
        rhel) dnf install -y -q "$pkg" >/dev/null 2>&1 ;;
    esac

    systemctl enable --now "$pkg" >/dev/null 2>&1
    log_info "Bouncer installed: $pkg"
}

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

install_collections() {
    local collections_to_install=""

    if [ -n "$COLLECTIONS" ]; then
        # User-specified collections
        IFS=',' read -ra cols <<< "$COLLECTIONS"
        for col in "${cols[@]}"; do
            col=$(echo "$col" | xargs)
            # Add crowdsecurity/ prefix if not present
            if [[ "$col" != */* ]]; then
                col="crowdsecurity/$col"
            fi
            collections_to_install="$collections_to_install $col"
        done
    else
        # Auto-detected collections
        collections_to_install="$DETECTED_COLLECTIONS"
    fi

    log_step "Installing collections: $collections_to_install"

    if [ "$DRY_RUN" = true ]; then
        log_info "[DRY RUN] Would install: $collections_to_install"
        return
    fi

    for col in $collections_to_install; do
        if cscli collections install "$col" >/dev/null 2>&1; then
            log_info "  Installed: $col"
        else
            log_warn "  Failed to install: $col"
        fi
    done
}

configure_allowlist() {
    if [ -z "$ALLOWLIST" ]; then
        return
    fi

    log_step "Configuring allowlist..."

    if [ "$DRY_RUN" = true ]; then
        log_info "[DRY RUN] Would whitelist: $ALLOWLIST"
        return
    fi

    IFS=',' read -ra ips <<< "$ALLOWLIST"
    for ip in "${ips[@]}"; do
        ip=$(echo "$ip" | xargs)
        if cscli decisions add --ip "$ip" --type whitelist --duration 87600h >/dev/null 2>&1; then
            log_info "  Whitelisted: $ip"
        else
            log_warn "  Failed to whitelist: $ip"
        fi
    done
}

configure_prometheus() {
    if [ "$PROMETHEUS" != true ]; then
        return
    fi

    log_step "Enabling Prometheus metrics on :6060..."

    if [ "$DRY_RUN" = true ]; then
        log_info "[DRY RUN] Would enable Prometheus metrics"
        return
    fi

    local config="/etc/crowdsec/config.yaml"
    if [ -f "$config" ]; then
        # Prometheus is enabled by default in CrowdSec, verify
        if grep -q "prometheus:" "$config"; then
            log_info "Prometheus metrics already configured"
        fi
    fi

    log_info "Prometheus metrics available at http://localhost:6060/metrics"
}

enroll_console() {
    if [ -z "$ENROLL_KEY" ]; then
        return
    fi

    log_step "Enrolling with CrowdSec console..."

    if [ "$DRY_RUN" = true ]; then
        log_info "[DRY RUN] Would enroll with key: $ENROLL_KEY"
        return
    fi

    if cscli console enroll "$ENROLL_KEY" >/dev/null 2>&1; then
        log_info "Enrolled with CrowdSec console"
    else
        log_warn "Console enrollment failed — verify enrollment key"
    fi
}

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

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

    echo ""

    # CrowdSec agent
    if systemctl is-active --quiet crowdsec 2>/dev/null; then
        log_info "✓ CrowdSec agent: running"
    else
        log_error "✗ CrowdSec agent: not running"
    fi

    # Bouncer
    if [ "$NO_BOUNCER" != true ]; then
        local bouncer_svc="crowdsec-firewall-bouncer-${BOUNCER_TYPE}"
        if systemctl is-active --quiet "$bouncer_svc" 2>/dev/null; then
            log_info "✓ Firewall bouncer: running"
        else
            log_error "✗ Firewall bouncer: not running"
        fi
    fi

    # Collections
    log_info "Installed collections:"
    cscli collections list 2>/dev/null | grep -E "enabled|installed" || true

    # Bouncers
    log_info "Registered bouncers:"
    cscli bouncers list 2>/dev/null || true

    echo ""
    log_info "Installation complete"
    echo ""
    log_info "Useful commands:"
    echo "  cscli decisions list      — view active decisions"
    echo "  cscli alerts list         — view recent alerts"
    echo "  cscli metrics             — view metrics summary"
    echo "  cscli hub list            — view installed hub items"
}

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

main() {
    parse_args "$@"
    check_root
    detect_os

    echo ""
    log_info "=== CrowdSec Installation Script v1.0 ==="
    echo ""

    if [ "$DRY_RUN" = true ]; then
        log_warn "DRY RUN MODE — no changes will be made"
        echo ""
    fi

    detect_services

    case "$OS_FAMILY" in
        debian) add_repo_debian ;;
        rhel) add_repo_rhel ;;
    esac

    install_agent
    install_collections
    configure_allowlist
    install_bouncer
    configure_prometheus
    enroll_console

    if [ "$DRY_RUN" != true ]; then
        verify_installation
    fi
}

main "$@"