#!/usr/bin/env bash

#########################################################################################
#### port-scan-reporter.sh — Scan hosts for open ports and compare against baseline  ####
#### Reports new open ports, closed ports, and service version changes               ####
#### Requires: bash 4+, nmap                                                         ####
####                                                                                 ####
#### Author: Phil Connor                                                             ####
#### Contact: contact@mylinux.work                                                   ####
#### License: MIT                                                                    ####
#### Version 1.00                                                                    ####
####                                                                                 ####
#### Usage:                                                                          ####
####   ./port-scan-reporter.sh --target 192.168.1.1                                  ####
####   ./port-scan-reporter.sh --targets-file hosts.txt --save-baseline              ####
####                                                                                 ####
#### See --help for all options.                                                     ####
#########################################################################################

set -euo pipefail

# ── Defaults & State ──────────────────────────────────────────────────
SCRIPT_NAME="$(basename "$0")"; readonly SCRIPT_NAME
TARGET=""; TARGETS_FILE=""; BASELINE_DIR="${HOME:-/tmp}/.port-scan-baselines"
SAVE_BASELINE=false; OUTPUT_JSON=false; TEXTFILE_PATH=""; PORTS=""; COLOR="auto"
TOTAL_HOSTS=0; HOSTS_WITH_CHANGES=0; TOTAL_NEW=0; TOTAL_CLOSED=0; TOTAL_CHANGED=0
JSON_RESULTS=(); PROM_METRICS=()

# ── Colors ────────────────────────────────────────────────────────────
RED="" GREEN="" YELLOW="" BLUE="" CYAN="" BOLD="" RESET=""
setup_colors() {
    if [[ "$COLOR" == "never" ]]; then
        RED="" GREEN="" YELLOW="" BLUE="" CYAN="" 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'
        CYAN='\033[0;36m'
        BOLD='\033[1m'
        RESET='\033[0m'
    else
        RED="" GREEN="" YELLOW="" BLUE="" CYAN="" BOLD="" RESET=""
    fi
}

# ── Logging ───────────────────────────────────────────────────────────
log()  { echo -e "${BLUE}[INFO]${RESET} $*"; }
warn() { echo -e "${YELLOW}[WARN]${RESET} $*" >&2; }
err()  { echo -e "${RED}[ERROR]${RESET} $*" >&2; }

# ── Dependency check ─────────────────────────────────────────────────
check_deps() {
    if ! command -v nmap &>/dev/null; then
        err "nmap is required but not installed"
        echo "  Install: apt install nmap  /  yum install nmap  /  brew install nmap" >&2
        exit 2
    fi
}

# ── Usage ─────────────────────────────────────────────────────────────
usage() {
    cat <<EOF
${SCRIPT_NAME} — Scan hosts for open ports and compare against baseline

USAGE:
  ${SCRIPT_NAME} --target <host>
  ${SCRIPT_NAME} --targets-file <file> [OPTIONS]

OPTIONS:
  --target <host>          Single host to scan (IP or hostname)
  --targets-file <file>    File with one host per line
  --baseline-dir <path>    Baseline directory (default: ~/.port-scan-baselines/)
  --save-baseline          Save current scan as the new baseline
  --json                   Output results as JSON
  --textfile <path>        Write Prometheus metrics to file
  --ports <range>          Port range (e.g., "22,80,443" or "1-1024", default: nmap top 1000)
  --no-color               Disable colored output
  --help                   Show this help

EXAMPLES:
  ${SCRIPT_NAME} --target 192.168.1.1
  ${SCRIPT_NAME} --targets-file hosts.txt --save-baseline
  ${SCRIPT_NAME} --target 10.0.0.5 --ports 22,80,443,8080 --json
  ${SCRIPT_NAME} --targets-file /etc/hosts.txt --textfile /var/lib/node_exporter/port_scan.prom

EXIT CODES:  0 = no changes, 1 = changes detected, 2 = error
EOF
}

# ── Argument parsing ─────────────────────────────────────────────────
parse_args() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --target)
                [[ $# -lt 2 ]] && { err "--target requires a host"; exit 2; }
                TARGET="$2"; shift 2 ;;
            --targets-file)
                [[ $# -lt 2 ]] && { err "--targets-file requires a path"; exit 2; }
                TARGETS_FILE="$2"; shift 2 ;;
            --baseline-dir)
                [[ $# -lt 2 ]] && { err "--baseline-dir requires a path"; exit 2; }
                BASELINE_DIR="$2"; shift 2 ;;
            --save-baseline)
                SAVE_BASELINE=true; shift ;;
            --json)
                OUTPUT_JSON=true; shift ;;
            --textfile)
                [[ $# -lt 2 ]] && { err "--textfile requires a path"; exit 2; }
                TEXTFILE_PATH="$2"; shift 2 ;;
            --ports)
                [[ $# -lt 2 ]] && { err "--ports requires a range"; exit 2; }
                PORTS="$2"; shift 2 ;;
            --no-color)
                COLOR="never"; shift ;;
            --help|-h)
                usage; exit 0 ;;
            *)
                err "Unknown option: $1"
                echo "Run ${SCRIPT_NAME} --help for usage" >&2
                exit 2 ;;
        esac
    done

    if [[ -z "$TARGET" && -z "$TARGETS_FILE" ]]; then
        err "Must specify --target or --targets-file"
        echo "Run ${SCRIPT_NAME} --help for usage" >&2
        exit 2
    fi

    if [[ -n "$TARGET" && -n "$TARGETS_FILE" ]]; then
        err "Cannot use both --target and --targets-file"
        exit 2
    fi

    if [[ -n "$TARGETS_FILE" && ! -f "$TARGETS_FILE" ]]; then
        err "Targets file not found: ${TARGETS_FILE}"
        exit 2
    fi
}

# ── Build host list ──────────────────────────────────────────────────
build_host_list() {
    local hosts=()
    if [[ -n "$TARGET" ]]; then
        hosts+=("$TARGET")
    else
        while IFS= read -r line; do
            line="${line%%#*}"
            line="${line// /}"
            [[ -z "$line" ]] && continue
            hosts+=("$line")
        done < "$TARGETS_FILE"
    fi
    if [[ ${#hosts[@]} -eq 0 ]]; then
        err "No hosts to scan"
        exit 2
    fi
    printf '%s\n' "${hosts[@]}"
}

# ── Host to baseline filename ────────────────────────────────────────
host_to_filename() { echo "${1//[.:\/ ]/_}.baseline"; }

# ── Scan a single host ──────────────────────────────────────────────
scan_host() {
    local host="$1"
    local nmap_args=(-sV --open -T4 --host-timeout 120)
    [[ -n "$PORTS" ]] && nmap_args+=(-p "$PORTS")
    nmap_args+=("$host")
    local raw_output
    raw_output=$(nmap "${nmap_args[@]}" 2>/dev/null) || {
        warn "nmap scan failed for ${host}"
        return 1
    }

    echo "$raw_output" | awk '
        /^[0-9]+\/(tcp|udp)/ {
            port_proto = $1
            state = $2
            service = ""
            for (i = 3; i <= NF; i++) {
                if (service != "") service = service " "
                service = service $i
            }
            if (service == "") service = "unknown"
            print port_proto " " state " " service
        }
    ' | sort -t/ -k1 -n
}

# ── Compare scan to baseline ────────────────────────────────────────
compare_to_baseline() {
    local host="$1"
    local current_file="$2"
    local baseline_file="${BASELINE_DIR}/$(host_to_filename "$host")"

    local new_ports=() closed_ports=() changed_services=()
    if [[ ! -f "$baseline_file" ]]; then
        if [[ "$OUTPUT_JSON" != true ]]; then
            echo -e "  ${CYAN}No baseline found — all ports reported as new${RESET}"
        fi
        local count
        count=$(wc -l < "$current_file")
        if [[ "$count" -gt 0 ]]; then
            while IFS= read -r line; do
                new_ports+=("$line")
            done < "$current_file"
        fi
    else
        # Find NEW ports (in current, not in baseline)
        while IFS= read -r line; do
            local port_proto
            port_proto=$(echo "$line" | awk '{print $1}')
            if ! grep -q "^${port_proto} " "$baseline_file" 2>/dev/null; then
                new_ports+=("$line")
            fi
        done < "$current_file"

        # Find CLOSED ports (in baseline, not in current)
        while IFS= read -r line; do
            local port_proto
            port_proto=$(echo "$line" | awk '{print $1}')
            if ! grep -q "^${port_proto} " "$current_file" 2>/dev/null; then
                closed_ports+=("$line")
            fi
        done < "$baseline_file"

        # Find CHANGED services (same port, different service)
        while IFS= read -r line; do
            local port_proto current_rest baseline_line baseline_rest
            port_proto=$(echo "$line" | awk '{print $1}')
            current_rest=$(echo "$line" | awk '{$1=""; print}' | sed 's/^ //')
            baseline_line=$(grep "^${port_proto} " "$baseline_file" 2>/dev/null || true)
            if [[ -n "$baseline_line" ]]; then
                baseline_rest=$(echo "$baseline_line" | awk '{$1=""; print}' | sed 's/^ //')
                if [[ "$current_rest" != "$baseline_rest" ]]; then
                    changed_services+=("${port_proto} ${baseline_rest} → ${current_rest}")
                fi
            fi
        done < "$current_file"
    fi

    local host_new=${#new_ports[@]} host_closed=${#closed_ports[@]} host_changed=${#changed_services[@]}
    local host_total_changes=$(( host_new + host_closed + host_changed ))

    if [[ "$OUTPUT_JSON" != true && "$host_total_changes" -gt 0 ]]; then
        echo ""
        echo -e "  ${BOLD}Comparing against baseline...${RESET}"
        echo ""
        for entry in "${new_ports[@]}"; do
            read -r pp _ svc <<< "$entry"
            echo -e "  ${RED}[NEW]${RESET}     ${pp}  open  ${svc}"
        done
        for entry in "${closed_ports[@]}"; do
            read -r pp _ svc <<< "$entry"
            echo -e "  ${YELLOW}[CLOSED]${RESET}  ${pp}  open  ${svc}"
        done
        for entry in "${changed_services[@]}"; do
            echo -e "  ${YELLOW}[CHANGED]${RESET} ${entry}"
        done
    elif [[ "$OUTPUT_JSON" != true && "$host_total_changes" -eq 0 ]]; then
        echo ""
        echo -e "  ${GREEN}No changes from baseline${RESET}"
    fi

    TOTAL_NEW=$(( TOTAL_NEW + host_new ))
    TOTAL_CLOSED=$(( TOTAL_CLOSED + host_closed ))
    TOTAL_CHANGED=$(( TOTAL_CHANGED + host_changed ))
    if [[ "$host_total_changes" -gt 0 ]]; then
        HOSTS_WITH_CHANGES=$(( HOSTS_WITH_CHANGES + 1 ))
    fi

    local open_count
    open_count=$(wc -l < "$current_file")
    if [[ "$OUTPUT_JSON" == true ]]; then
        local json_entry
        json_entry=$(build_host_json "$host" "$current_file" "$open_count" \
            "$host_new" "$host_closed" "$host_changed")
        JSON_RESULTS+=("$json_entry")
    fi

    if [[ -n "$TEXTFILE_PATH" ]]; then
        PROM_METRICS+=("port_scan_open_ports{host=\"${host}\"} ${open_count}")
        PROM_METRICS+=("port_scan_new_ports{host=\"${host}\"} ${host_new}")
    fi
}

# ── Build JSON for a single host ─────────────────────────────────────
build_host_json() {
    local host="$1" current_file="$2" open_count="$3"
    local new_count="$4" closed_count="$5" changed_count="$6"
    local ports_json="[" first=true
    while IFS= read -r line; do
        read -r pp state svc_rest <<< "$line"
        [[ "$first" == true ]] && first=false || ports_json+=","
        ports_json+="{\"port\":\"${pp}\",\"state\":\"${state}\",\"service\":\"${svc_rest:-unknown}\"}"
    done < "$current_file"
    ports_json+="]"
    echo "{\"host\":\"${host}\",\"timestamp\":\"$(date -u +%Y-%m-%dT%H:%M:%SZ)\",\"open_ports\":${open_count},\"new_ports\":${new_count},\"closed_ports\":${closed_count},\"changed_services\":${changed_count},\"ports\":${ports_json}}"
}

# ── Print summary ────────────────────────────────────────────────────
print_summary() {
    if [[ "$OUTPUT_JSON" == true ]]; then
        return
    fi
    echo ""
    echo "  ───────────────────────────────────────────"
    echo -e "  ${BOLD}Summary${RESET}"
    printf "  %-22s %d\n" "Hosts scanned:" "$TOTAL_HOSTS"
    printf "  %-22s %d\n" "Hosts with changes:" "$HOSTS_WITH_CHANGES"
    printf "  %-22s %d\n" "New open ports:" "$TOTAL_NEW"
    printf "  %-22s %d\n" "Closed ports:" "$TOTAL_CLOSED"
    printf "  %-22s %d\n" "Changed services:" "$TOTAL_CHANGED"
    echo ""
}

# ── Write JSON output ───────────────────────────────────────────────
write_json() {
    local first=true; echo "["
    for entry in "${JSON_RESULTS[@]}"; do
        [[ "$first" == true ]] && first=false || echo ","
        echo -n "  ${entry}"
    done
    echo ""; echo "]"
}

# ── Write Prometheus textfile ────────────────────────────────────────
write_textfile() {
    local tmpfile="${TEXTFILE_PATH}.$$"
    {
        echo "# HELP port_scan_open_ports Number of open ports detected on host"
        echo "# TYPE port_scan_open_ports gauge"
        echo "# HELP port_scan_new_ports Number of new ports since last baseline"
        echo "# TYPE port_scan_new_ports gauge"
        echo "# HELP port_scan_changes_total Total port changes across all hosts"
        echo "# TYPE port_scan_changes_total gauge"
        for metric in "${PROM_METRICS[@]}"; do
            echo "$metric"
        done
        local total_changes=$(( TOTAL_NEW + TOTAL_CLOSED + TOTAL_CHANGED ))
        echo "port_scan_changes_total ${total_changes}"
    } > "$tmpfile"
    mv "$tmpfile" "$TEXTFILE_PATH"
    log "Prometheus metrics written to ${TEXTFILE_PATH}"
}

# ── Save baseline ───────────────────────────────────────────────────
save_baseline() {
    local host="$1" current_file="$2"
    local baseline_file="${BASELINE_DIR}/$(host_to_filename "$host")"
    cp "$current_file" "$baseline_file"
    log "Baseline saved for ${host} → ${baseline_file}"
}

# ── Main ─────────────────────────────────────────────────────────────
main() {
    parse_args "$@"
    setup_colors
    check_deps
    mkdir -p "$BASELINE_DIR"
    local hosts; hosts=$(build_host_list)
    if [[ "$OUTPUT_JSON" != true ]]; then
        echo ""
        echo -e "${BOLD}Port Scan Reporter${RESET}"
        echo "═══════════════════════════════════════════"
    fi

    local port_display="default top 1000"
    [[ -n "$PORTS" ]] && port_display="$PORTS"
    while IFS= read -r host; do
        TOTAL_HOSTS=$(( TOTAL_HOSTS + 1 ))
        if [[ "$OUTPUT_JSON" != true ]]; then
            echo ""
            echo -e "Scanning ${BOLD}${host}${RESET} (ports: ${port_display})..."
        fi

        local tmpfile; tmpfile=$(mktemp)
        trap "rm -f '$tmpfile'" EXIT
        if ! scan_host "$host" > "$tmpfile"; then
            warn "Skipping ${host} — scan failed"
            rm -f "$tmpfile"
            continue
        fi

        local open_count; open_count=$(wc -l < "$tmpfile")
        if [[ "$OUTPUT_JSON" != true && "$open_count" -gt 0 ]]; then
            echo ""
            printf "  %-14s %-8s %s\n" "PORT/PROTO" "STATE" "SERVICE"
            while IFS= read -r line; do
                read -r pp state svc <<< "$line"
                printf "  %-14s %-8s %s\n" "$pp" "$state" "$svc"
            done < "$tmpfile"
        elif [[ "$OUTPUT_JSON" != true ]]; then
            echo "  No open ports found"
        fi

        compare_to_baseline "$host" "$tmpfile"
        if [[ "$SAVE_BASELINE" == true ]]; then
            save_baseline "$host" "$tmpfile"
        fi
        rm -f "$tmpfile"
    done <<< "$hosts"

    if [[ "$OUTPUT_JSON" == true ]]; then
        write_json
    else
        print_summary
    fi

    [[ -n "$TEXTFILE_PATH" ]] && write_textfile
    if [[ "$SAVE_BASELINE" == true ]]; then
        exit 0
    fi
    local total_changes=$(( TOTAL_NEW + TOTAL_CLOSED + TOTAL_CHANGED ))
    if [[ "$total_changes" -gt 0 ]]; then
        exit 1
    fi
    exit 0
}

main "$@"