#!/usr/bin/env bash

#########################################################################################
#### contabo-firewall-auditor.sh — Audit Contabo firewall rules across all instances ####
#### Finds open ports, missing firewalls, overly permissive rules, and misconfigs    ####
#### Requires: bash 4+, curl, jq                                                    ####
####                                                                                 ####
#### Author: Phil Connor                                                             ####
#### Contact: contact@mylinux.work                                                   ####
#### License: MIT                                                                    ####
#### Version 1.01                                                                    ####
####                                                                                 ####
#### Usage:                                                                          ####
####   ./contabo-firewall-auditor.sh --full                                          ####
####                                                                                 ####
#### See --help for all options.                                                     ####
#########################################################################################

set -euo pipefail

# ── Colors (pre-initialized) ─────────────────────────────────────────
RED="" GREEN="" YELLOW="" BLUE="" CYAN="" BOLD="" DIM="" RESET=""

setup_colors() {
    if [[ "${COLOR:-auto}" == "never" ]]; then
        return
    fi
    if [[ "${COLOR:-auto}" == "always" ]] || [[ -t 1 ]]; then
        RED='\033[0;31m'
        GREEN='\033[0;32m'
        YELLOW='\033[0;33m'
        BLUE='\033[0;34m'
        CYAN='\033[0;36m'
        BOLD='\033[1m'
        DIM='\033[2m'
        RESET='\033[0m'
    fi
}

# ── Logging ───────────────────────────────────────────────────────────
log()     { echo -e "${BLUE}[INFO]${RESET} $*"; }
warn()    { echo -e "${YELLOW}[WARN]${RESET} $*" >&2; }
err()     { echo -e "${RED}[ERROR]${RESET} $*" >&2; }
verbose() { if [[ "$VERBOSE" == "true" ]]; then echo -e "${DIM}[DEBUG]${RESET} $*"; fi; }
die()     { err "$*"; exit 1; }

# ── Severity counters ────────────────────────────────────────────────
TOTAL_CRIT=0
TOTAL_WARN=0
TOTAL_INFO=0
TOTAL_OK=0

flag_crit() { ((TOTAL_CRIT++)) || true; }
flag_warn() { ((TOTAL_WARN++)) || true; }
flag_info() { ((TOTAL_INFO++)) || true; }
flag_ok()   { ((TOTAL_OK++)) || true; }

# ── Defaults ──────────────────────────────────────────────────────────
RUN_MODE=""
DANGEROUS_PORTS="${DANGEROUS_PORTS:-22,3389,3306,5432,1433,6379,27017,9200,8080,8443}"
VERBOSE="${VERBOSE:-false}"
COLOR="${COLOR:-auto}"

# ── Credentials ───────────────────────────────────────────────────────
CONTABO_CLIENT_ID="${CONTABO_CLIENT_ID:-}"
CONTABO_CLIENT_SECRET="${CONTABO_CLIENT_SECRET:-}"
CONTABO_API_USER="${CONTABO_API_USER:-}"
CONTABO_API_PASS="${CONTABO_API_PASS:-}"

# ── State ─────────────────────────────────────────────────────────────
SCRIPT_NAME="$(basename "$0")"
readonly SCRIPT_NAME
START_TIME=""

# ── API helpers ──────────────────────────────────────────────────────
contabo_token() {
    local resp
    resp=$(curl -s -d "client_id=${CONTABO_CLIENT_ID}" \
        -d "client_secret=${CONTABO_CLIENT_SECRET}" \
        --data-urlencode "username=${CONTABO_API_USER}" \
        --data-urlencode "password=${CONTABO_API_PASS}" \
        -d "grant_type=password" \
        "https://auth.contabo.com/auth/realms/contabo/protocol/openid-connect/token")
    local token
    token=$(echo "$resp" | jq -r '.access_token // empty' 2>/dev/null)
    if [[ -z "$token" ]]; then
        die "Failed to obtain access token — check credentials"
    fi
    echo "$token"
}

contabo_api() {
    local method="$1" endpoint="$2"
    shift 2
    local attempt=0 max_attempts=3

    while (( attempt < max_attempts )); do
        local http_code
        http_code=$(curl -s -o /tmp/cfa_resp.json -w "%{http_code}" \
            -X "$method" \
            -H "Authorization: Bearer $(contabo_token)" \
            -H "Content-Type: application/json" \
            -H "x-request-id: $(cat /proc/sys/kernel/random/uuid 2>/dev/null || date +%s%N)" \
            "https://api.contabo.com/v1${endpoint}" "$@")

        verbose "API ${method} ${endpoint} → HTTP ${http_code}"

        if [[ "$http_code" == "429" ]]; then
            ((attempt++)) || true
            local wait=$(( attempt * 5 ))
            warn "Rate limited — retrying in ${wait}s (attempt ${attempt}/${max_attempts})"
            sleep "$wait"
            continue
        fi

        cat /tmp/cfa_resp.json
        return 0
    done

    err "API request failed after ${max_attempts} attempts: ${method} ${endpoint}"
    return 1
}

check_credentials() {
    [[ -z "$CONTABO_CLIENT_ID" ]] && die "CONTABO_CLIENT_ID not set"
    [[ -z "$CONTABO_CLIENT_SECRET" ]] && die "CONTABO_CLIENT_SECRET not set"
    [[ -z "$CONTABO_API_USER" ]] && die "CONTABO_API_USER not set"
    [[ -z "$CONTABO_API_PASS" ]] && die "CONTABO_API_PASS not set"
}

check_deps() {
    command -v curl &>/dev/null || die "curl is required"
    command -v jq &>/dev/null || die "jq is required"
}

# ── Instance helpers ─────────────────────────────────────────────────
get_all_instances() {
    local page=1 size=100 result="[]"
    while true; do
        local resp
        resp=$(contabo_api GET "/compute/instances?page=${page}&size=${size}")
        local page_data
        page_data=$(echo "$resp" | jq '.data // []' 2>/dev/null)
        local count
        count=$(echo "$page_data" | jq 'length' 2>/dev/null || echo 0)
        [[ "$count" -eq 0 ]] && break
        result=$(echo "$result" "$page_data" | jq -s '.[0] + .[1]')
        (( count < size )) && break
        ((page++)) || true
    done
    echo "$result"
}

# ── Firewall helpers ─────────────────────────────────────────────────
get_all_firewalls() {
    local page=1 size=100 result="[]"
    while true; do
        local resp
        resp=$(contabo_api GET "/firewalls?page=${page}&size=${size}")
        local page_data
        page_data=$(echo "$resp" | jq '.data // []' 2>/dev/null)
        local count
        count=$(echo "$page_data" | jq 'length' 2>/dev/null || echo 0)
        [[ "$count" -eq 0 ]] && break
        result=$(echo "$result" "$page_data" | jq -s '.[0] + .[1]')
        (( count < size )) && break
        ((page++)) || true
    done
    echo "$result"
}

# ── Port-to-service mapping ─────────────────────────────────────────
port_to_service() {
    local port="$1"
    case "$port" in
        22)    echo "SSH" ;;
        80)    echo "HTTP" ;;
        443)   echo "HTTPS" ;;
        3306)  echo "MySQL" ;;
        5432)  echo "PostgreSQL" ;;
        1433)  echo "MSSQL" ;;
        3389)  echo "RDP" ;;
        6379)  echo "Redis" ;;
        27017) echo "MongoDB" ;;
        9200)  echo "Elasticsearch" ;;
        8080)  echo "HTTP-Alt" ;;
        8443)  echo "HTTPS-Alt" ;;
        53)    echo "DNS" ;;
        25)    echo "SMTP" ;;
        5900)  echo "VNC" ;;
        11211) echo "Memcached" ;;
        2379)  echo "etcd" ;;
        9090)  echo "Prometheus" ;;
        *)     echo "" ;;
    esac
}

# ── Check if port is in dangerous list ───────────────────────────────
is_dangerous_port() {
    local port="$1"
    local IFS=','
    for dp in $DANGEROUS_PORTS; do
        if [[ "$port" == "$dp" ]]; then
            return 0
        fi
    done
    return 1
}

# ══════════════════════════════════════════════════════════════════════
#                       OPEN PORTS AUDIT
# ══════════════════════════════════════════════════════════════════════
audit_open_ports() {
    log "Auditing firewall rules for dangerous open ports..."
    log "Dangerous ports: ${DANGEROUS_PORTS}"
    echo ""

    printf "  %-10s %-22s %-8s %-8s %-18s %-12s %s\n" \
        "FW_ID" "FW_NAME" "PORT" "PROTO" "SOURCE" "SERVICE" "SEVERITY"
    printf "  %s\n" "$(printf '%.0s─' {1..95})"

    local fw_json
    fw_json=$(get_all_firewalls)

    echo "$fw_json" | jq -c '.[]' 2>/dev/null | while IFS= read -r fw; do
        local fw_id fw_name
        fw_id=$(echo "$fw" | jq -r '.firewallId' 2>/dev/null)
        fw_name=$(echo "$fw" | jq -r '.displayName // .name // "unnamed"' 2>/dev/null)

        echo "$fw" | jq -c '.rules[]? // empty' 2>/dev/null | while IFS= read -r rule; do
            local action protocol port_str source_cidr
            action=$(echo "$rule" | jq -r '.action // "accept"' 2>/dev/null)
            protocol=$(echo "$rule" | jq -r '.protocol // "tcp"' 2>/dev/null)
            port_str=$(echo "$rule" | jq -r '.port // .destPorts // .destinationPorts // ""' 2>/dev/null)
            source_cidr=$(echo "$rule" | jq -r '.srcCidr // .source // .ipRange // "0.0.0.0/0"' 2>/dev/null)

            [[ "$action" != "accept" && "$action" != "allow" ]] && continue
            [[ "$source_cidr" != "0.0.0.0/0" && "$source_cidr" != "::/0" ]] && continue

            if [[ -z "$port_str" || "$port_str" == "null" ]]; then
                local IFS=','
                for dp in $DANGEROUS_PORTS; do
                    local svc
                    svc=$(port_to_service "$dp")
                    printf "  %-10s %-22s %-8s %-8s %-18s %-12s %b%s%b\n" \
                        "$fw_id" "${fw_name:0:20}" "$dp" "$protocol" \
                        "$source_cidr" "${svc:-unknown}" "$RED" "CRITICAL" "$RESET"
                    flag_crit
                done
                continue
            fi

            local IFS=','
            for port_entry in $port_str; do
                local single_port="$port_entry"
                if [[ "$port_entry" == *-* ]]; then
                    local range_start range_end
                    range_start="${port_entry%-*}"
                    range_end="${port_entry#*-}"
                    local IFS=','
                    for dp in $DANGEROUS_PORTS; do
                        if [[ "$dp" -ge "$range_start" && "$dp" -le "$range_end" ]]; then
                            local svc
                            svc=$(port_to_service "$dp")
                            printf "  %-10s %-22s %-8s %-8s %-18s %-12s %b%s%b\n" \
                                "$fw_id" "${fw_name:0:20}" "$dp" "$protocol" \
                                "$source_cidr" "${svc:-unknown}" "$RED" "CRITICAL" "$RESET"
                            flag_crit
                        fi
                    done
                    continue
                fi

                if is_dangerous_port "$single_port"; then
                    local svc
                    svc=$(port_to_service "$single_port")
                    printf "  %-10s %-22s %-8s %-8s %-18s %-12s %b%s%b\n" \
                        "$fw_id" "${fw_name:0:20}" "$single_port" "$protocol" \
                        "$source_cidr" "${svc:-unknown}" "$RED" "CRITICAL" "$RESET"
                    flag_crit
                elif [[ "$single_port" == "80" || "$single_port" == "443" ]]; then
                    local svc
                    svc=$(port_to_service "$single_port")
                    printf "  %-10s %-22s %-8s %-8s %-18s %-12s %b%s%b\n" \
                        "$fw_id" "${fw_name:0:20}" "$single_port" "$protocol" \
                        "$source_cidr" "${svc:-$single_port}" "$CYAN" "INFO" "$RESET"
                    flag_info
                fi
            done
        done
    done

    echo ""
}

# ══════════════════════════════════════════════════════════════════════
#                       UNPROTECTED INSTANCES
# ══════════════════════════════════════════════════════════════════════
audit_unprotected() {
    log "Checking for instances without firewalls..."
    echo ""

    printf "  %-10s %-22s %-16s %-10s %s\n" \
        "INST_ID" "NAME" "IP" "STATUS" "FIREWALL"
    printf "  %s\n" "$(printf '%.0s─' {1..75})"

    local instances
    instances=$(get_all_instances)

    local fw_json
    fw_json=$(get_all_firewalls)

    local assigned_instances
    assigned_instances=$(echo "$fw_json" | jq -r \
        '[.[].assignedInstances // [] | .[]] | unique | .[]' 2>/dev/null || true)

    echo "$instances" | jq -c '.[]' 2>/dev/null | while IFS= read -r inst; do
        local iid iname ip status
        iid=$(echo "$inst" | jq -r '.instanceId' 2>/dev/null)
        iname=$(echo "$inst" | jq -r '.name // .displayName // "unknown"' 2>/dev/null)
        ip=$(echo "$inst" | jq -r '.ipConfig.v4.ip // "N/A"' 2>/dev/null)
        status=$(echo "$inst" | jq -r '.status // "unknown"' 2>/dev/null)

        local has_fw="false"
        if echo "$assigned_instances" | grep -q "^${iid}$" 2>/dev/null; then
            has_fw="true"
        fi

        if [[ "$has_fw" == "false" ]]; then
            printf "  %-10s %-22s %-16s %-10s %b%s%b\n" \
                "$iid" "${iname:0:20}" "$ip" "$status" \
                "$RED" "NONE — UNPROTECTED" "$RESET"
            flag_crit
        else
            printf "  %-10s %-22s %-16s %-10s %b%s%b\n" \
                "$iid" "${iname:0:20}" "$ip" "$status" \
                "$GREEN" "✓ Protected" "$RESET"
            flag_ok
        fi
    done

    echo ""
}

# ══════════════════════════════════════════════════════════════════════
#                       PERMISSIVE RULES AUDIT
# ══════════════════════════════════════════════════════════════════════
audit_permissive() {
    log "Auditing overly permissive firewall rules..."
    echo ""

    printf "  %-10s %-22s %-10s %-8s %-18s %-14s %s\n" \
        "FW_ID" "FW_NAME" "PORTS" "PROTO" "SOURCE" "ISSUE" "SEVERITY"
    printf "  %s\n" "$(printf '%.0s─' {1..100})"

    local fw_json
    fw_json=$(get_all_firewalls)

    echo "$fw_json" | jq -c '.[]' 2>/dev/null | while IFS= read -r fw; do
        local fw_id fw_name
        fw_id=$(echo "$fw" | jq -r '.firewallId' 2>/dev/null)
        fw_name=$(echo "$fw" | jq -r '.displayName // .name // "unnamed"' 2>/dev/null)

        echo "$fw" | jq -c '.rules[]? // empty' 2>/dev/null | while IFS= read -r rule; do
            local action protocol port_str source_cidr
            action=$(echo "$rule" | jq -r '.action // "accept"' 2>/dev/null)
            protocol=$(echo "$rule" | jq -r '.protocol // "tcp"' 2>/dev/null)
            port_str=$(echo "$rule" | jq -r '.port // .destPorts // .destinationPorts // ""' 2>/dev/null)
            source_cidr=$(echo "$rule" | jq -r '.srcCidr // .source // .ipRange // ""' 2>/dev/null)

            [[ "$action" != "accept" && "$action" != "allow" ]] && continue

            if [[ -z "$port_str" || "$port_str" == "null" ]] && [[ "$source_cidr" == "0.0.0.0/0" || "$source_cidr" == "::/0" ]]; then
                printf "  %-10s %-22s %-10s %-8s %-18s %-14s %b%s%b\n" \
                    "$fw_id" "${fw_name:0:20}" "ALL" "$protocol" \
                    "$source_cidr" "all-ports" "$RED" "CRITICAL" "$RESET"
                flag_crit
                continue
            fi

            if [[ "$protocol" == "all" || "$protocol" == "-1" ]] && [[ "$source_cidr" == "0.0.0.0/0" || "$source_cidr" == "::/0" ]]; then
                printf "  %-10s %-22s %-10s %-8s %-18s %-14s %b%s%b\n" \
                    "$fw_id" "${fw_name:0:20}" "${port_str:-ALL}" "all" \
                    "$source_cidr" "all-protocols" "$RED" "CRITICAL" "$RESET"
                flag_crit
                continue
            fi

            if [[ -n "$source_cidr" && "$source_cidr" != "null" ]]; then
                if [[ "$source_cidr" == *"/8" || "$source_cidr" == *"/16" ]]; then
                    printf "  %-10s %-22s %-10s %-8s %-18s %-14s %b%s%b\n" \
                        "$fw_id" "${fw_name:0:20}" "${port_str:-ALL}" "$protocol" \
                        "${source_cidr:0:16}" "wide-cidr" "$YELLOW" "WARN" "$RESET"
                    flag_warn
                fi
            fi
        done
    done

    echo ""
}

# ══════════════════════════════════════════════════════════════════════
#                       UNUSED FIREWALLS
# ══════════════════════════════════════════════════════════════════════
audit_unused() {
    log "Checking for unused firewalls..."
    echo ""

    printf "  %-10s %-28s %-8s %s\n" \
        "FW_ID" "FW_NAME" "RULES" "STATUS"
    printf "  %s\n" "$(printf '%.0s─' {1..60})"

    local fw_json
    fw_json=$(get_all_firewalls)

    echo "$fw_json" | jq -c '.[]' 2>/dev/null | while IFS= read -r fw; do
        local fw_id fw_name rule_count assigned_count
        fw_id=$(echo "$fw" | jq -r '.firewallId' 2>/dev/null)
        fw_name=$(echo "$fw" | jq -r '.displayName // .name // "unnamed"' 2>/dev/null)
        rule_count=$(echo "$fw" | jq '[.rules[]?] | length' 2>/dev/null || echo 0)
        assigned_count=$(echo "$fw" | jq '[.assignedInstances // [] | .[]] | length' 2>/dev/null || echo 0)

        if [[ "$assigned_count" -eq 0 ]]; then
            printf "  %-10s %-28s %-8s %b%s%b\n" \
                "$fw_id" "${fw_name:0:26}" "$rule_count" \
                "$YELLOW" "UNUSED" "$RESET"
            flag_warn
        else
            verbose "Firewall ${fw_id} (${fw_name}): assigned to ${assigned_count} instance(s)"
            flag_ok
        fi
    done

    echo ""
}

# ══════════════════════════════════════════════════════════════════════
#                       LIST ALL RULES
# ══════════════════════════════════════════════════════════════════════
list_rules() {
    log "Listing all firewall rules..."
    echo ""

    printf "  %-10s %-20s %-8s %-8s %-12s %-18s %s\n" \
        "FW_ID" "FW_NAME" "ACTION" "PROTO" "PORTS" "SOURCE" "SERVICE"
    printf "  %s\n" "$(printf '%.0s─' {1..90})"

    local fw_json
    fw_json=$(get_all_firewalls)

    echo "$fw_json" | jq -c '.[]' 2>/dev/null | while IFS= read -r fw; do
        local fw_id fw_name
        fw_id=$(echo "$fw" | jq -r '.firewallId' 2>/dev/null)
        fw_name=$(echo "$fw" | jq -r '.displayName // .name // "unnamed"' 2>/dev/null)

        echo "$fw" | jq -c '.rules[]? // empty' 2>/dev/null | while IFS= read -r rule; do
            local action protocol port_str source_cidr
            action=$(echo "$rule" | jq -r '.action // "accept"' 2>/dev/null)
            protocol=$(echo "$rule" | jq -r '.protocol // "tcp"' 2>/dev/null)
            port_str=$(echo "$rule" | jq -r '.port // .destPorts // .destinationPorts // "all"' 2>/dev/null)
            source_cidr=$(echo "$rule" | jq -r '.srcCidr // .source // .ipRange // "any"' 2>/dev/null)

            [[ "$port_str" == "null" ]] && port_str="all"
            [[ "$source_cidr" == "null" ]] && source_cidr="any"

            local svc=""
            if [[ "$port_str" =~ ^[0-9]+$ ]]; then
                svc=$(port_to_service "$port_str")
            fi

            local action_color="$GREEN"
            [[ "$action" == "drop" || "$action" == "deny" || "$action" == "reject" ]] && action_color="$RED"

            printf "  %-10s %-20s %b%-8s%b %-8s %-12s %-18s %s\n" \
                "$fw_id" "${fw_name:0:18}" "$action_color" "$action" "$RESET" \
                "$protocol" "${port_str:0:10}" "${source_cidr:0:16}" "${svc}"
        done
    done

    echo ""
}

# ══════════════════════════════════════════════════════════════════════
#                       SUMMARY
# ══════════════════════════════════════════════════════════════════════
print_summary() {
    local elapsed
    elapsed=$(( $(date +%s) - START_TIME ))

    echo ""
    echo "  ══════════════════════════════════════════"
    echo "  Firewall Audit Summary"
    echo "  ══════════════════════════════════════════"
    printf "  %-20s %b%d%b\n" "CRITICAL:" "$RED" "$TOTAL_CRIT" "$RESET"
    printf "  %-20s %b%d%b\n" "WARN:" "$YELLOW" "$TOTAL_WARN" "$RESET"
    printf "  %-20s %b%d%b\n" "INFO:" "$CYAN" "$TOTAL_INFO" "$RESET"
    printf "  %-20s %b%d%b\n" "OK:" "$GREEN" "$TOTAL_OK" "$RESET"
    echo "  ──────────────────────────────────────────"
    printf "  Completed in %ds\n" "$elapsed"
    echo ""

    if [[ "$TOTAL_CRIT" -gt 0 ]]; then
        echo -e "  ${RED}${BOLD}Action required:${RESET} ${TOTAL_CRIT} critical finding(s)"
        echo ""
        echo "  Top recommendations:"
        echo "    • Assign firewalls to all unprotected instances"
        echo "    • Close 0.0.0.0/0 rules on SSH (22), RDP (3389), and database ports"
        echo "    • Replace all-port allow rules with specific port lists"
        echo "    • Remove unused firewalls to reduce configuration sprawl"
        echo ""
    elif [[ "$TOTAL_WARN" -gt 0 ]]; then
        echo -e "  ${YELLOW}Review recommended:${RESET} ${TOTAL_WARN} warning(s)"
        echo ""
        echo "  Suggestions:"
        echo "    • Review wide CIDR rules and narrow where possible"
        echo "    • Delete unused firewalls"
        echo "    • Restrict outbound where applicable"
        echo ""
    else
        echo -e "  ${GREEN}All checks passed${RESET}"
        echo ""
    fi
}

# ══════════════════════════════════════════════════════════════════════
#                       USAGE
# ══════════════════════════════════════════════════════════════════════
show_help() {
    cat <<EOF
${BOLD}${SCRIPT_NAME}${RESET} — Contabo Firewall Auditor

Audit Contabo firewall rules across all instances for risky
configurations via the REST API.

${BOLD}MODES${RESET}
  --full              Run all audits
  --open-ports        Find dangerous ports open to 0.0.0.0/0
  --unprotected       Find instances with no firewall assigned
  --permissive        Find overly broad allow rules
  --unused            Find firewalls not assigned to any instance
  --rules             List all firewall rules

${BOLD}OPTIONS${RESET}
  --ports PORTS       Override dangerous ports (comma-separated)
  --verbose           Debug output
  --no-color          Disable colored output
  --help              Show this help message

${BOLD}ENVIRONMENT VARIABLES${RESET}
  CONTABO_CLIENT_ID      OAuth2 Client ID (required)
  CONTABO_CLIENT_SECRET  OAuth2 Client Secret (required)
  CONTABO_API_USER       API username / email (required)
  CONTABO_API_PASS       API password (required)
  DANGEROUS_PORTS        Comma-separated ports to flag (default: 22,3389,...)
  VERBOSE                Enable verbose output (true/false)
  COLOR                  Color mode: auto, always, never

${BOLD}EXAMPLES${RESET}
  # Full audit
  ${SCRIPT_NAME} --full

  # Check open ports only
  ${SCRIPT_NAME} --open-ports

  # Find unprotected instances
  ${SCRIPT_NAME} --unprotected

  # Custom dangerous ports
  ${SCRIPT_NAME} --open-ports --ports "22,3389,5432,6379"

  # List all firewall rules
  ${SCRIPT_NAME} --rules

${BOLD}EXIT CODES${RESET}
  0   All checks passed
  1   Warnings found (review recommended)
  2   Critical findings (action required)
EOF
}

# ══════════════════════════════════════════════════════════════════════
#                       PARSE ARGS
# ══════════════════════════════════════════════════════════════════════
parse_args() {
    local modes=()

    while [[ $# -gt 0 ]]; do
        case "$1" in
            --full)
                modes=(open-ports unprotected permissive unused)
                shift ;;
            --open-ports)
                modes+=(open-ports); shift ;;
            --unprotected)
                modes+=(unprotected); shift ;;
            --permissive)
                modes+=(permissive); shift ;;
            --unused)
                modes+=(unused); shift ;;
            --rules)
                modes+=(rules); shift ;;
            --ports)
                DANGEROUS_PORTS="${2:?--ports requires a value}"; shift 2 ;;
            --verbose)
                VERBOSE="true"; shift ;;
            --no-color)
                COLOR="never"; shift ;;
            --help|-h)
                setup_colors; show_help; exit 0 ;;
            *)
                die "Unknown option: $1 (see --help)" ;;
        esac
    done

    if [[ ${#modes[@]} -eq 0 ]]; then
        err "No audit mode specified"
        echo "Run ${SCRIPT_NAME} --help for usage" >&2
        exit 1
    fi

    RUN_MODE="${modes[*]}"
}

# ══════════════════════════════════════════════════════════════════════
#                       MAIN
# ══════════════════════════════════════════════════════════════════════
main() {
    parse_args "$@"
    setup_colors
    check_deps
    check_credentials

    START_TIME=$(date +%s)

    echo ""
    echo -e "${BOLD}Contabo Firewall Auditor${RESET}"
    echo -e "Mode:   ${RUN_MODE}"
    echo -e "Time:   $(date -u +%Y-%m-%dT%H:%M:%SZ)"
    echo ""

    for mode in $RUN_MODE; do
        case "$mode" in
            open-ports)   audit_open_ports ;;
            unprotected)  audit_unprotected ;;
            permissive)   audit_permissive ;;
            unused)       audit_unused ;;
            rules)        list_rules ;;
        esac
    done

    print_summary

    if [[ "$TOTAL_CRIT" -gt 0 ]]; then
        exit 2
    elif [[ "$TOTAL_WARN" -gt 0 ]]; then
        exit 1
    fi
    exit 0
}

main "$@"