#!/usr/bin/env bash
set -euo pipefail

REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
HOMELAB_STATE_DIR="${HOMELAB_STATE_DIR:-${XDG_DATA_HOME:-${HOME}/.local/share}/homelab}"
SECURITY_REPORT_DIR="${SECURITY_REPORT_DIR:-${HOMELAB_STATE_DIR}/security-reports}"
SECURITY_TARGETS_FILE="${SECURITY_TARGETS_FILE:-${REPO_ROOT}/security/targets.txt}"
TRIVY_IMAGE="${TRIVY_IMAGE:-aquasec/trivy:latest}"
ZAP_IMAGE="${ZAP_IMAGE:-ghcr.io/zaproxy/zaproxy:stable}"
KUBE_BENCH_IMAGE="${KUBE_BENCH_IMAGE:-aquasec/kube-bench:latest}"
NUCLEI_IMAGE="${NUCLEI_IMAGE:-projectdiscovery/nuclei:latest}"

usage() {
    cat <<'EOF'
Usage: scripts/security-scan {all|trivy|zap|kube-bench|nuclei}

Environment:
  SECURITY_REPORT_DIR    Report output directory.
  SECURITY_TARGETS_FILE  Target allowlist for web scans.
  SECURITY_ZAP_TARGET    Single URL override for ZAP.
  SECURITY_NUCLEI_TARGET Single URL override for nuclei.
EOF
}

timestamp() {
    date -u +%Y%m%dT%H%M%SZ
}

ensure_report_dir() {
    mkdir -p "${SECURITY_REPORT_DIR}"
}

require_docker() {
    if ! command -v docker >/dev/null 2>&1; then
        echo "docker is required for containerized security scans." >&2
        exit 1
    fi
}

run_trivy() {
    local stamp
    stamp="$(timestamp)"
    ensure_report_dir

    echo "Running Trivy filesystem and IaC/config scans..."
    if command -v trivy >/dev/null 2>&1; then
        trivy fs --scanners vuln,secret,misconfig --format table --output "${SECURITY_REPORT_DIR}/trivy-fs-${stamp}.txt" "${REPO_ROOT}"
        trivy config --format table --output "${SECURITY_REPORT_DIR}/trivy-config-${stamp}.txt" "${REPO_ROOT}"
    else
        require_docker
        docker run --rm \
            -v "${REPO_ROOT}:/repo:ro" \
            -v "${SECURITY_REPORT_DIR}:/reports" \
            "${TRIVY_IMAGE}" fs --scanners vuln,secret,misconfig --format table --output "/reports/trivy-fs-${stamp}.txt" /repo
        docker run --rm \
            -v "${REPO_ROOT}:/repo:ro" \
            -v "${SECURITY_REPORT_DIR}:/reports" \
            "${TRIVY_IMAGE}" config --format table --output "/reports/trivy-config-${stamp}.txt" /repo
    fi

    echo "Trivy reports written to ${SECURITY_REPORT_DIR}"
}

scan_targets() {
    local target_override="$1"

    if [[ -n "${target_override}" ]]; then
        printf '%s\n' "${target_override}"
        return 0
    fi

    if [[ ! -s "${SECURITY_TARGETS_FILE}" ]]; then
        echo "Missing target file: ${SECURITY_TARGETS_FILE}" >&2
        exit 1
    fi

    sed -e 's/#.*//' -e '/^[[:space:]]*$/d' "${SECURITY_TARGETS_FILE}"
}

safe_target_name() {
    printf '%s' "$1" | sed -E 's#^https?://##; s#[^A-Za-z0-9._-]+#_#g; s#_+$##'
}

run_zap() {
    local stamp
    local target
    local safe_name
    stamp="$(timestamp)"
    ensure_report_dir
    require_docker

    echo "Running OWASP ZAP baseline passive scans..."
    while IFS= read -r target; do
        [[ -n "${target}" ]] || continue
        safe_name="$(safe_target_name "${target}")"
        echo "--> ZAP baseline ${target}"
        docker run --rm \
            -v "${SECURITY_REPORT_DIR}:/zap/wrk:rw" \
            "${ZAP_IMAGE}" zap-baseline.py \
            -t "${target}" \
            -I \
            -r "zap-${safe_name}-${stamp}.html" \
            -J "zap-${safe_name}-${stamp}.json"
    done < <(scan_targets "${SECURITY_ZAP_TARGET:-}")

    echo "ZAP reports written to ${SECURITY_REPORT_DIR}"
}

run_nuclei() {
    local stamp
    local target
    local safe_name
    stamp="$(timestamp)"
    ensure_report_dir
    require_docker

    echo "Running nuclei low-noise HTTP exposure scans..."
    while IFS= read -r target; do
        [[ -n "${target}" ]] || continue
        safe_name="$(safe_target_name "${target}")"
        echo "--> nuclei ${target}"
        docker run --rm \
            -v "${SECURITY_REPORT_DIR}:/reports:rw" \
            "${NUCLEI_IMAGE}" \
            -u "${target}" \
            -severity low,medium,high,critical \
            -rate-limit "${SECURITY_NUCLEI_RATE_LIMIT:-5}" \
            -retries 1 \
            -timeout 5 \
            -o "/reports/nuclei-${safe_name}-${stamp}.txt"
    done < <(scan_targets "${SECURITY_NUCLEI_TARGET:-}")

    echo "nuclei reports written to ${SECURITY_REPORT_DIR}"
}

run_kube_bench() {
    local stamp
    stamp="$(timestamp)"
    ensure_report_dir
    require_docker

    echo "Running kube-bench CIS checks against this kubeadm host..."
    docker run --rm --pid=host \
        -v /etc:/etc:ro \
        -v /var:/var:ro \
        -v /usr/bin:/usr/local/mount-from-host/bin:ro \
        "${KUBE_BENCH_IMAGE}" run --targets master,node | tee "${SECURITY_REPORT_DIR}/kube-bench-${stamp}.txt"
}

case "${1:-all}" in
    all)
        run_trivy
        run_zap
        ;;
    trivy)
        run_trivy
        ;;
    zap)
        run_zap
        ;;
    nuclei)
        run_nuclei
        ;;
    kube-bench)
        run_kube_bench
        ;;
    -h | --help | help)
        usage
        ;;
    *)
        usage >&2
        exit 1
        ;;
esac