#!/usr/bin/env bash

######################################################################################
#### port-exposure-scanner.sh — Audit listening ports against an allow-list       ####
#### Detects unauthorized listeners, missing expected services, and reports       ####
#### per-port detail including process names and PIDs.                            ####
#### Requires: bash 4+, ss (iproute2)                                             ####
####                                                                              ####
#### Author: Phil Connor                                                          ####
#### Contact: contact@mylinux.work                                                ####
#### License: MIT                                                                 ####
#### Version 1.00                                                                 ####
####                                                                              ####
#### Usage:                                                                       ####
####   sudo ./port-exposure-scanner.sh --role webserver                           ####
####   sudo ./port-exposure-scanner.sh --allow-list /etc/port-exposure.conf       ####
####                                                                              ####
#### See --help for all options.                                                  ####
######################################################################################

set -euo pipefail

# ── Defaults ──────────────────────────────────────────────────────────
ALLOW_LIST_FILE="/etc/port-exposure-scanner/allow-list.conf"
ALLOW_LIST_PROVIDED=false
ROLE=""
INCLUDE_UDP=false
TEXTFILE_MODE=false
PROM_FILE="/var/lib/node_exporter/port_exposure.prom"
OUTPUT_FORMAT="${OUTPUT_FORMAT:-text}"  # text, tap, junit
JUNIT_FILE="${JUNIT_FILE:-port-exposure-results.xml}"
VERBOSE="${VERBOSE:-false}"
COLOR="${COLOR:-auto}"

# ── State ─────────────────────────────────────────────────────────────
PASS=0
FAIL=0
WARN=0
TOTAL=0
RESULTS=()
START_TIME=""
CAN_SHOW_PROCESS=true

# Allow-list entries: associative array  port/proto -> description
declare -A ALLOWED_PORTS

# Discovered listeners: associative array  port/proto -> process_info
declare -A LISTENING_PORTS

# Unauthorized port details for Prometheus labels
declare -a UNAUTHORIZED_DETAILS=()

# ── Colors ────────────────────────────────────────────────────────────
setup_colors() {
    if [[ "$COLOR" == "never" ]]; then
        RED="" GREEN="" YELLOW="" BLUE="" BOLD="" RESET=""
        return
    fi
    if [[ "$COLOR" == "always" ]] || [[ -t 1 ]]; then
        RED='\033[0;31m'
        GREEN='\033[0;32m'
        YELLOW='\033[0;33m'
        BLUE='\033[0;34m'
        BOLD='\033[1m'
        RESET='\033[0m'
    else
        RED="" GREEN="" YELLOW="" BLUE="" BOLD="" RESET=""
    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 "${BLUE}[DEBUG]${RESET} $*"; fi; }

# ── Test Result Recording ─────────────────────────────────────────────
record_pass() {
    local name="$1"
    local detail="${2:-}"
    ((PASS++)) || true
    ((TOTAL++)) || true
    RESULTS+=("PASS|${name}|${detail}")
    if [[ "$OUTPUT_FORMAT" == "tap" ]]; then
        echo "ok ${TOTAL} - ${name}"
    elif [[ "$OUTPUT_FORMAT" == "text" ]]; then
        echo -e "  ${GREEN}✓${RESET} ${name}${detail:+ — ${detail}}"
    fi
}

record_fail() {
    local name="$1"
    local detail="${2:-}"
    ((FAIL++)) || true
    ((TOTAL++)) || true
    RESULTS+=("FAIL|${name}|${detail}")
    if [[ "$OUTPUT_FORMAT" == "tap" ]]; then
        echo "not ok ${TOTAL} - ${name}"
        [[ -n "$detail" ]] && echo "  # ${detail}"
    elif [[ "$OUTPUT_FORMAT" == "text" ]]; then
        echo -e "  ${RED}✗${RESET} ${name}${detail:+ — ${detail}}"
    fi
}

record_warn() {
    local name="$1"
    local detail="${2:-}"
    ((WARN++)) || true
    ((TOTAL++)) || true
    RESULTS+=("WARN|${name}|${detail}")
    if [[ "$OUTPUT_FORMAT" == "tap" ]]; then
        echo "ok ${TOTAL} - ${name} # SKIP ${detail}"
    elif [[ "$OUTPUT_FORMAT" == "text" ]]; then
        echo -e "  ${YELLOW}⊘${RESET} ${name}${detail:+ — ${detail}}"
    fi
}

# ── JUnit XML Writer ──────────────────────────────────────────────────
write_junit() {
    local end_time
    end_time=$(date +%s)
    local duration=$(( end_time - START_TIME ))

    cat > "$JUNIT_FILE" <<JUNIT_EOF
<?xml version="1.0" encoding="UTF-8"?>
<testsuites tests="${TOTAL}" failures="${FAIL}" skipped="${WARN}" time="${duration}">
  <testsuite name="port-exposure-scanner" tests="${TOTAL}" failures="${FAIL}" skipped="${WARN}" time="${duration}">
JUNIT_EOF

    for result in "${RESULTS[@]}"; do
        local status name detail
        status=$(echo "$result" | cut -d'|' -f1)
        name=$(echo "$result" | cut -d'|' -f2)
        detail=$(echo "$result" | cut -d'|' -f3)

        # XML-escape
        name="${name//&/&amp;}"
        name="${name//</&lt;}"
        name="${name//>/&gt;}"
        name="${name//\"/&quot;}"
        detail="${detail//&/&amp;}"
        detail="${detail//</&lt;}"
        detail="${detail//>/&gt;}"
        detail="${detail//\"/&quot;}"

        case "$status" in
            PASS)
                echo "    <testcase name=\"${name}\" classname=\"port-exposure\">" >> "$JUNIT_FILE"
                [[ -n "$detail" ]] && echo "      <system-out>${detail}</system-out>" >> "$JUNIT_FILE"
                echo "    </testcase>" >> "$JUNIT_FILE"
                ;;
            FAIL)
                echo "    <testcase name=\"${name}\" classname=\"port-exposure\">" >> "$JUNIT_FILE"
                echo "      <failure message=\"${detail}\">FAILED: ${name} — ${detail}</failure>" >> "$JUNIT_FILE"
                echo "    </testcase>" >> "$JUNIT_FILE"
                ;;
            WARN)
                echo "    <testcase name=\"${name}\" classname=\"port-exposure\">" >> "$JUNIT_FILE"
                echo "      <skipped message=\"${detail}\"/>" >> "$JUNIT_FILE"
                echo "    </testcase>" >> "$JUNIT_FILE"
                ;;
        esac
    done

    echo "  </testsuite>" >> "$JUNIT_FILE"
    echo "</testsuites>" >> "$JUNIT_FILE"

    log "JUnit report written to ${JUNIT_FILE}"
}

# ══════════════════════════════════════════════════════════════════════
#                        ALLOW-LIST LOADING
# ══════════════════════════════════════════════════════════════════════

# ── Built-in roles ────────────────────────────────────────────────────
load_role() {
    local role="$1"
    case "$role" in
        minimal)
            ALLOWED_PORTS["22/tcp"]="SSH"
            ;;
        webserver)
            ALLOWED_PORTS["22/tcp"]="SSH"
            ALLOWED_PORTS["80/tcp"]="HTTP"
            ALLOWED_PORTS["443/tcp"]="HTTPS"
            ;;
        database)
            ALLOWED_PORTS["22/tcp"]="SSH"
            ALLOWED_PORTS["3306/tcp"]="MySQL"
            ALLOWED_PORTS["5432/tcp"]="PostgreSQL"
            ;;
        monitoring)
            ALLOWED_PORTS["22/tcp"]="SSH"
            ALLOWED_PORTS["9090/tcp"]="Prometheus"
            ALLOWED_PORTS["9093/tcp"]="Alertmanager"
            ALLOWED_PORTS["9100/tcp"]="node_exporter"
            ALLOWED_PORTS["3000/tcp"]="Grafana"
            ;;
        *)
            err "Unknown role: ${role}"
            err "Available roles: minimal, webserver, database, monitoring"
            exit 1
            ;;
    esac
    verbose "Loaded role '${role}' with ${#ALLOWED_PORTS[@]} allowed ports"
}

# ── Parse allow-list config file ──────────────────────────────────────
load_allow_list_file() {
    local file="$1"
    if [[ ! -f "$file" ]]; then
        err "Allow-list file not found: ${file}"
        exit 1
    fi
    if [[ ! -r "$file" ]]; then
        err "Allow-list file not readable: ${file}"
        exit 1
    fi

    local line_num=0
    while IFS= read -r line || [[ -n "$line" ]]; do
        ((line_num++)) || true
        # Strip leading/trailing whitespace
        line="${line#"${line%%[![:space:]]*}"}"
        line="${line%"${line##*[![:space:]]}"}"

        # Skip empty lines and comments
        [[ -z "$line" || "$line" =~ ^# ]] && continue

        # Extract description from inline comment
        local description=""
        if [[ "$line" =~ \#(.*) ]]; then
            description="${BASH_REMATCH[1]}"
            description="${description#"${description%%[![:space:]]*}"}"
        fi

        # Extract port/protocol token (first field)
        local token
        token=$(echo "$line" | awk '{print $1}')

        local port proto
        if [[ "$token" =~ ^([0-9]+)/(.+)$ ]]; then
            port="${BASH_REMATCH[1]}"
            proto="${BASH_REMATCH[2]}"
        elif [[ "$token" =~ ^[0-9]+$ ]]; then
            port="$token"
            proto="tcp"
        else
            warn "Ignoring invalid line ${line_num} in ${file}: ${line}"
            continue
        fi

        ALLOWED_PORTS["${port}/${proto}"]="${description:-port ${port}}"
        verbose "Allow-listed: ${port}/${proto} (${description:-port ${port}})"
    done < "$file"

    verbose "Loaded ${#ALLOWED_PORTS[@]} allowed ports from ${file}"
}

# ══════════════════════════════════════════════════════════════════════
#                        PORT SCANNING
# ══════════════════════════════════════════════════════════════════════

# ── Parse ss output into LISTENING_PORTS ──────────────────────────────
scan_ports() {
    local proto="$1"
    local ss_flag="$2"

    verbose "Scanning ${proto} listeners with ss ${ss_flag}"

    local ss_output
    if [[ "$CAN_SHOW_PROCESS" == "true" ]]; then
        ss_output=$(ss "${ss_flag}" -n -p 2>/dev/null) || ss_output=$(ss "${ss_flag}" -n 2>/dev/null)
    else
        ss_output=$(ss "${ss_flag}" -n 2>/dev/null) || true
    fi

    # Skip header line; parse each listening entry
    while IFS= read -r line; do
        [[ "$line" =~ ^State ]] && continue
        [[ -z "$line" ]] && continue

        # Extract local address:port (4th column)
        local local_addr
        local_addr=$(echo "$line" | awk '{print $4}')

        local port
        # Handle IPv6 [::]:port and IPv4 0.0.0.0:port and *:port
        if [[ "$local_addr" =~ ]:([0-9]+)$ ]]; then
            port="${BASH_REMATCH[1]}"
        elif [[ "$local_addr" =~ :([0-9]+)$ ]]; then
            port="${BASH_REMATCH[1]}"
        else
            continue
        fi

        # Extract process info if available
        local process_info=""
        if [[ "$line" =~ users:\(\(\"([^\"]+)\",pid=([0-9]+) ]]; then
            process_info="${BASH_REMATCH[1]}:${BASH_REMATCH[2]}"
        fi

        local key="${port}/${proto}"
        # Only record once per port/proto (first match wins)
        if [[ -z "${LISTENING_PORTS[$key]+x}" ]]; then
            LISTENING_PORTS["$key"]="${process_info:-unknown}"
            verbose "Found listener: ${key} (${process_info:-unknown})"
        fi
    done <<< "$ss_output"
}

# ══════════════════════════════════════════════════════════════════════
#                        ANALYSIS
# ══════════════════════════════════════════════════════════════════════

analyze_ports() {
    local auth_listening=0
    local auth_missing=0
    local unauth_listening=0

    # ── Check authorized ports ────────────────────────────────────────
    if [[ "$OUTPUT_FORMAT" == "text" ]]; then
        echo ""
        echo -e "${BOLD}Authorized Ports${RESET}"
    fi

    for key in $(echo "${!ALLOWED_PORTS[@]}" | tr ' ' '\n' | sort -t/ -k1 -n); do
        local desc="${ALLOWED_PORTS[$key]}"
        local port="${key%%/*}"
        local proto="${key##*/}"

        if [[ -n "${LISTENING_PORTS[$key]+x}" ]]; then
            local proc="${LISTENING_PORTS[$key]}"
            record_pass "Port ${port}/${proto} listening" "${desc} (${proc})"
            ((auth_listening++)) || true
        else
            record_warn "Port ${port}/${proto} not listening" "expected: ${desc}"
            ((auth_missing++)) || true
        fi
    done

    # ── Check unauthorized ports ──────────────────────────────────────
    local has_unauthorized=false
    for key in $(echo "${!LISTENING_PORTS[@]}" | tr ' ' '\n' | sort -t/ -k1 -n); do
        if [[ -z "${ALLOWED_PORTS[$key]+x}" ]]; then
            if [[ "$has_unauthorized" == "false" && "$OUTPUT_FORMAT" == "text" ]]; then
                echo ""
                echo -e "${BOLD}Unauthorized Ports${RESET}"
                has_unauthorized=true
            fi

            local port="${key%%/*}"
            local proto="${key##*/}"
            local proc="${LISTENING_PORTS[$key]}"
            local proc_name="" proc_pid=""

            if [[ "$proc" =~ ^(.+):([0-9]+)$ ]]; then
                proc_name="${BASH_REMATCH[1]}"
                proc_pid="${BASH_REMATCH[2]}"
            else
                proc_name="${proc}"
                proc_pid="n/a"
            fi

            record_fail "Port ${port}/${proto} unauthorized" "process=${proc_name} pid=${proc_pid}"
            UNAUTHORIZED_DETAILS+=("${port}|${proto}|${proc_name}")
            ((unauth_listening++)) || true
        fi
    done

    if [[ "$has_unauthorized" == "false" && "$OUTPUT_FORMAT" == "text" ]]; then
        echo ""
        echo -e "${BOLD}Unauthorized Ports${RESET}"
        echo -e "  ${GREEN}✓${RESET} No unauthorized listeners detected"
    fi

    # Store counts for Prometheus / summary
    AUTH_LISTENING=$auth_listening
    AUTH_MISSING=$auth_missing
    UNAUTH_LISTENING=$unauth_listening
}

# ══════════════════════════════════════════════════════════════════════
#                        PROMETHEUS TEXTFILE
# ══════════════════════════════════════════════════════════════════════

write_prometheus() {
    local prom_dir
    prom_dir=$(dirname "$PROM_FILE")
    if [[ ! -d "$prom_dir" ]]; then
        warn "Prometheus textfile directory does not exist: ${prom_dir}"
        warn "Skipping textfile write. Create the directory or use --prom-file."
        return
    fi

    local tmpfile="${PROM_FILE}.$$"

    {
        echo "# HELP port_exposure_authorized_listening Number of authorized ports that are listening"
        echo "# TYPE port_exposure_authorized_listening gauge"
        echo "port_exposure_authorized_listening ${AUTH_LISTENING}"
        echo ""
        echo "# HELP port_exposure_authorized_missing Number of authorized ports not listening"
        echo "# TYPE port_exposure_authorized_missing gauge"
        echo "port_exposure_authorized_missing ${AUTH_MISSING}"
        echo ""
        echo "# HELP port_exposure_unauthorized_listening Number of unauthorized ports listening"
        echo "# TYPE port_exposure_unauthorized_listening gauge"
        echo "port_exposure_unauthorized_listening ${UNAUTH_LISTENING}"
        echo ""
        echo "# HELP port_exposure_scan_timestamp Unix timestamp of last scan"
        echo "# TYPE port_exposure_scan_timestamp gauge"
        echo "port_exposure_scan_timestamp $(date +%s)"
        echo ""
        echo "# HELP port_exposure_unauthorized_port Per-port indicator for unauthorized listeners"
        echo "# TYPE port_exposure_unauthorized_port gauge"
        for entry in "${UNAUTHORIZED_DETAILS[@]}"; do
            local port proto process
            port=$(echo "$entry" | cut -d'|' -f1)
            proto=$(echo "$entry" | cut -d'|' -f2)
            process=$(echo "$entry" | cut -d'|' -f3)
            echo "port_exposure_unauthorized_port{port=\"${port}\",protocol=\"${proto}\",process=\"${process}\"} 1"
        done
    } > "$tmpfile"

    mv "$tmpfile" "$PROM_FILE"
    log "Prometheus metrics written to ${PROM_FILE}"
}

# ══════════════════════════════════════════════════════════════════════
#                          OUTPUT
# ══════════════════════════════════════════════════════════════════════

print_summary() {
    local end_time
    end_time=$(date +%s)
    local duration=$(( end_time - START_TIME ))

    echo ""
    echo -e "${BOLD}────────────────────────────────────────${RESET}"
    echo -e "${BOLD}Summary${RESET}"
    echo -e "  Authorized listening: ${AUTH_LISTENING}"
    echo -e "  Authorized missing:   ${AUTH_MISSING}"
    echo -e "  Unauthorized:         ${UNAUTH_LISTENING}"
    echo -e "  ${GREEN}${PASS} passed${RESET}  ${RED}${FAIL} failed${RESET}  ${YELLOW}${WARN} warnings${RESET}  (${duration}s)"
    echo -e "${BOLD}────────────────────────────────────────${RESET}"

    if [[ $FAIL -eq 0 ]]; then
        echo -e "${GREEN}${BOLD}All checks passed.${RESET}"
    else
        echo -e "${RED}${BOLD}${FAIL} unauthorized port(s) detected.${RESET}"
    fi
}

print_tap_header() {
    echo "TAP version 13"
}

print_tap_footer() {
    echo "1..${TOTAL}"
    echo "# pass ${PASS}"
    echo "# fail ${FAIL}"
    echo "# warn ${WARN}"
}

# ══════════════════════════════════════════════════════════════════════
#                          MAIN
# ══════════════════════════════════════════════════════════════════════

usage() {
    cat <<EOF
Usage: $(basename "$0") [OPTIONS]

Audit listening TCP/UDP ports against an allow-list. Flags unauthorized
listeners, detects missing expected services, and reports per-port process
detail. Read-only — never modifies services or kills processes.

Requires: bash 4+, ss (iproute2)

Options:
  --allow-list PATH   Path to allow-list config file
                      (default: /etc/port-exposure-scanner/allow-list.conf)
  --role ROLE         Use a built-in role profile instead of a file
                      Roles: minimal, webserver, database, monitoring
  --include-udp       Also scan UDP listeners
  --textfile          Write Prometheus metrics to textfile
  --prom-file PATH    Custom Prometheus textfile path
                      (default: /var/lib/node_exporter/port_exposure.prom)
  --format FORMAT     Output: text (default), tap, junit
  --junit-file FILE   JUnit output path (default: port-exposure-results.xml)
  --verbose           Show debug output
  --no-color          Disable colored output
  -h, --help          Show this help

Allow-list format (one port per line):
  # Comments start with #
  # Format: port[/protocol] [# description]
  22/tcp      # SSH
  80/tcp      # HTTP
  443/tcp     # HTTPS
  9100/tcp    # node_exporter

Built-in roles:
  minimal      22/tcp
  webserver    22, 80, 443/tcp
  database     22, 3306, 5432/tcp
  monitoring   22, 9090, 9093, 9100, 3000/tcp

Examples:
  # Scan with a role profile
  sudo ./$(basename "$0") --role webserver

  # Scan with a custom allow-list
  sudo ./$(basename "$0") --allow-list /etc/myports.conf

  # Scan TCP+UDP, write Prometheus metrics
  sudo ./$(basename "$0") --role minimal --include-udp --textfile

  # TAP output for CI pipeline
  sudo ./$(basename "$0") --role webserver --format tap

  # JUnit report
  sudo ./$(basename "$0") --allow-list ./ports.conf --format junit
EOF
}

main() {
    setup_colors

    # Parse arguments
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --allow-list)     ALLOW_LIST_FILE="$2"; ALLOW_LIST_PROVIDED=true; shift ;;
            --role)           ROLE="$2"; shift ;;
            --include-udp)    INCLUDE_UDP=true ;;
            --textfile)       TEXTFILE_MODE=true ;;
            --prom-file)      PROM_FILE="$2"; shift ;;
            --format)         OUTPUT_FORMAT="$2"; shift ;;
            --junit-file)     JUNIT_FILE="$2"; shift ;;
            --verbose)        VERBOSE=true ;;
            --no-color)       COLOR=never ;;
            --help|-h)        usage; exit 0 ;;
            *)                err "Unknown option: $1"; usage; exit 1 ;;
        esac
        shift
    done

    # re-apply colors in case --no-color was passed
    setup_colors

    # ── Check for root (needed for process names via ss -p) ───────────
    if [[ $EUID -ne 0 ]]; then
        warn "Not running as root — process names will not be available (ss -p requires root)"
        CAN_SHOW_PROCESS=false
    fi

    # ── Load allow-list ───────────────────────────────────────────────
    if [[ -n "$ROLE" && "$ALLOW_LIST_PROVIDED" == "true" ]]; then
        err "Cannot specify both --role and --allow-list"
        exit 1
    fi

    if [[ -n "$ROLE" ]]; then
        load_role "$ROLE"
    elif [[ "$ALLOW_LIST_PROVIDED" == "true" ]]; then
        load_allow_list_file "$ALLOW_LIST_FILE"
    elif [[ -f "$ALLOW_LIST_FILE" ]]; then
        load_allow_list_file "$ALLOW_LIST_FILE"
    else
        err "No allow-list found. Provide --allow-list PATH, --role ROLE,"
        err "or create ${ALLOW_LIST_FILE}"
        exit 1
    fi

    if [[ ${#ALLOWED_PORTS[@]} -eq 0 ]]; then
        err "Allow-list is empty — nothing to check"
        exit 1
    fi

    # ── Verify ss is available ────────────────────────────────────────
    if ! command -v ss &>/dev/null; then
        err "ss command not found. Install iproute2."
        exit 1
    fi

    START_TIME=$(date +%s)

    if [[ "$OUTPUT_FORMAT" == "tap" ]]; then
        print_tap_header
    elif [[ "$OUTPUT_FORMAT" == "text" ]]; then
        echo ""
        echo -e "${BOLD}Port Exposure Scanner${RESET}"
        echo -e "Time:       $(date -u +%Y-%m-%dT%H:%M:%SZ)"
        if [[ -n "$ROLE" ]]; then
            echo -e "Role:       ${ROLE}"
        else
            echo -e "Allow-list: ${ALLOW_LIST_FILE}"
        fi
        local proto_str="TCP"
        [[ "$INCLUDE_UDP" == "true" ]] && proto_str="TCP, UDP"
        echo -e "Protocols:  ${proto_str}"
    fi

    # ── Scan listening ports ──────────────────────────────────────────
    scan_ports "tcp" "-tln"

    if [[ "$INCLUDE_UDP" == "true" ]]; then
        scan_ports "udp" "-uln"
    fi

    verbose "Found ${#LISTENING_PORTS[@]} unique listening ports"

    # ── Analyze ───────────────────────────────────────────────────────
    analyze_ports

    # ── Output ────────────────────────────────────────────────────────
    if [[ "$OUTPUT_FORMAT" == "tap" ]]; then
        print_tap_footer
    elif [[ "$OUTPUT_FORMAT" == "junit" ]]; then
        print_summary
        write_junit
    else
        print_summary
    fi

    # ── Prometheus textfile ───────────────────────────────────────────
    if [[ "$TEXTFILE_MODE" == "true" ]]; then
        write_prometheus
    fi

    # Exit code
    [[ $FAIL -eq 0 ]] && exit 0 || exit 1
}

main "$@"