#!/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}"
GITLEAKS_IMAGE="${GITLEAKS_IMAGE:-ghcr.io/gitleaks/gitleaks:latest}"

truthy() {
    case "${1,,}" in
        1 | true | yes | y | on | enabled)
            return 0
            ;;
        *)
            return 1
            ;;
    esac
}

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

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.
  SECURITY_WEB_TARGET    Single URL override for web header checks.
  LAB_BACKSTAGE_BRAIN_ENABLED=true enables LLM-assisted Lynis triage.
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
}

prepare_security_scans() {
    local image
    local -a images=(
        "${TRIVY_IMAGE}"
        "${ZAP_IMAGE}"
        "${KUBE_BENCH_IMAGE}"
        "${NUCLEI_IMAGE}"
        "${GITLEAKS_IMAGE}"
    )

    ensure_report_dir
    require_docker

    if ! docker info >/dev/null 2>&1; then
        echo "docker is installed but not reachable by this user." >&2
        exit 1
    fi
    if [[ ! -s "${SECURITY_TARGETS_FILE}" ]]; then
        echo "Missing target file: ${SECURITY_TARGETS_FILE}" >&2
        exit 1
    fi
    if [[ ! -w "${SECURITY_REPORT_DIR}" ]]; then
        echo "Report directory is not writable: ${SECURITY_REPORT_DIR}" >&2
        exit 1
    fi

    echo "Preparing security scan images..."
    for image in "${images[@]}"; do
        echo "--> docker pull ${image}"
        docker pull "${image}"
    done

    echo "Security scan prerequisites are ready."
    echo "Report directory: ${SECURITY_REPORT_DIR}"
}

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}"
}

run_gitleaks() {
    local stamp
    local report_file
    stamp="$(timestamp)"
    ensure_report_dir
    report_file="${SECURITY_REPORT_DIR}/gitleaks-${stamp}.json"

    echo "Running gitleaks secret scan against this repo..."
    if command -v gitleaks >/dev/null 2>&1; then
        gitleaks detect \
            --source "${REPO_ROOT}" \
            --report-format json \
            --report-path "${report_file}" \
            --redact
    else
        require_docker
        docker run --rm \
            -v "${REPO_ROOT}:/repo:ro" \
            -v "${SECURITY_REPORT_DIR}:/reports:rw" \
            "${GITLEAKS_IMAGE}" detect \
            --source /repo \
            --report-format json \
            --report-path "/reports/gitleaks-${stamp}.json" \
            --redact
    fi

    echo "gitleaks report: ${report_file}"
}

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}"
}

header_value() {
    local headers_file="$1"
    local header_name="$2"

    awk -v name="${header_name}" '
        BEGIN { IGNORECASE = 1 }
        index($0, name ":") == 1 {
            sub(/^[^:]+:[[:space:]]*/, "")
            sub(/\r$/, "")
            print
            exit
        }
    ' "${headers_file}"
}

run_security_web() {
    local stamp
    local target
    local safe_name
    local headers_file
    local body_file
    local report_file
    local status
    local failures=0
    local missing=0
    local -a required_headers=(
        "strict-transport-security"
        "x-content-type-options"
        "x-frame-options"
        "referrer-policy"
    )

    stamp="$(timestamp)"
    ensure_report_dir
    report_file="${SECURITY_REPORT_DIR}/security-web-${stamp}.txt"

    {
        echo "Security web quick check"
        echo "Generated: ${stamp}"
        echo
    } >"${report_file}"

    while IFS= read -r target; do
        [[ -n "${target}" ]] || continue
        safe_name="$(safe_target_name "${target}")"
        headers_file="${SECURITY_REPORT_DIR}/headers-${safe_name}-${stamp}.txt"
        body_file="${SECURITY_REPORT_DIR}/headers-${safe_name}-${stamp}.body"

        echo "== ${target}" | tee -a "${report_file}"
        if [[ "${target}" != https://* ]]; then
            echo "fail: target is not HTTPS" | tee -a "${report_file}"
            failures=$((failures + 1))
        fi

        status="$(
            curl -k -sS -L \
                --connect-timeout "${SECURITY_WEB_CONNECT_TIMEOUT:-8}" \
                --max-time "${SECURITY_WEB_MAX_TIME:-20}" \
                -D "${headers_file}" \
                -o "${body_file}" \
                -w '%{http_code}' \
                "${target}" 2>>"${report_file}" || true
        )"

        if [[ ! "${status}" =~ ^2|3[0-9][0-9]$ ]]; then
            echo "fail: HTTP status ${status:-curl failed}" | tee -a "${report_file}"
            failures=$((failures + 1))
        else
            echo "ok: HTTP status ${status}" | tee -a "${report_file}"
        fi

        missing=0
        for header_name in "${required_headers[@]}"; do
            if [[ -n "$(header_value "${headers_file}" "${header_name}")" ]]; then
                echo "ok: ${header_name}" | tee -a "${report_file}"
            else
                echo "warn: missing ${header_name}" | tee -a "${report_file}"
                missing=$((missing + 1))
            fi
        done

        if [[ -n "$(header_value "${headers_file}" "content-security-policy")" ]]; then
            echo "ok: content-security-policy" | tee -a "${report_file}"
        else
            echo "info: content-security-policy missing or intentionally deferred" | tee -a "${report_file}"
        fi

        if ((missing > 0)); then
            failures=$((failures + 1))
        fi
        echo | tee -a "${report_file}" >/dev/null
        rm -f "${body_file}"
    done < <(scan_targets "${SECURITY_WEB_TARGET:-}")

    echo "Security web report: ${report_file}"
    if ((failures > 0)); then
        echo "security-web found ${failures} target(s) needing review." >&2
        return 1
    fi
}

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"
}

lynis_findings() {
    local report_file="$1"

    sed -n 's/^warning\[\]=/warning|/p; s/^suggestion\[\]=/suggestion|/p' "${report_file}" 2>/dev/null |
        awk -F'|' 'NF >= 2 { source = $1; finding = $2; test = $NF; print source "|" finding (NF > 2 ? " [" test "]" : "") }'
}

print_lynis_fallback_focus() {
    local findings_file="$1"
    local selected

    selected="$(awk -F'|' '$1 == "warning" { print; exit }' "${findings_file}")"
    if [[ -z "${selected}" ]]; then
        selected="$(awk -F'|' '$1 == "suggestion" { print; exit }' "${findings_file}")"
    fi

    if [[ -z "${selected}" ]]; then
        printf '\n== Highest-Risk Lynis Finding ==\n'
        printf 'No Lynis warnings or suggestions were found.\n'
        return 0
    fi

    printf '\n== Highest-Risk Lynis Finding ==\n'
    printf 'Source: %s\n' "$(cut -d'|' -f1 <<<"${selected}")"
    printf 'Finding: %s\n' "$(cut -d'|' -f2- <<<"${selected}")"
    printf 'Why this one: LLM triage was unavailable, so Jeannie selected the first Lynis warning, falling back to the first suggestion.\n'
    printf 'Mode: report-only. Add an explicit narrow Jeannie hardening command before changing host state.\n'
}

print_lynis_llm_focus() {
    local findings_file="$1"
    local endpoint="${LAB_AI_GATEWAY_URL:-${LAB_OLLAMA_URL:-http://127.0.0.1:11434}}"
    local model="${LAB_AI_GATEWAY_MODEL:-qwen2.5:0.5b}"
    local timeout="${LAB_AI_GATEWAY_TIMEOUT_SECONDS:-20}"
    local index_dir="${LAB_AI_KNOWLEDGE_INDEX_DIR:-/data/homelab-ai/index}"
    local context=""

    if ! truthy "${LAB_BACKSTAGE_BRAIN_ENABLED:-false}"; then
        return 1
    fi
    if ! command -v python3 >/dev/null 2>&1; then
        return 1
    fi

    if [[ -s "${index_dir}/index.json" && -x "${REPO_ROOT}/scripts/query-homelab-ai-index" ]]; then
        context="$("${REPO_ROOT}/scripts/query-homelab-ai-index" --index-dir "${index_dir}" --context-only --limit 4 "Lynis host hardening Debian Docker kubeadm SSH homelab" 2>/dev/null || true)"
    fi

    LYNIS_FINDINGS="$(head -120 "${findings_file}")" \
        LYNIS_CONTEXT="${context}" \
        python3 - "${endpoint}" "${model}" "${timeout}" <<'PY'
import json
import os
import sys
import urllib.error
import urllib.request

endpoint, model, timeout = sys.argv[1:4]
findings = os.environ.get("LYNIS_FINDINGS", "").strip()
context = os.environ.get("LYNIS_CONTEXT", "").strip()
if not findings:
    raise SystemExit(1)

try:
    timeout_seconds = int(timeout)
except ValueError:
    timeout_seconds = 20

try:
    with urllib.request.urlopen(f"{endpoint.rstrip('/')}/api/tags", timeout=3) as response:
        tags = json.loads(response.read().decode("utf-8"))
    available = {item.get("name", "") for item in tags.get("models", [])}
    if model not in available and f"{model}:latest" not in available:
        raise SystemExit(1)
except (OSError, TimeoutError, urllib.error.URLError, json.JSONDecodeError):
    raise SystemExit(1)

prompt = f"""You are Jeannie, a report-only homelab security triage helper.
The host is a Debian personal homelab server running Docker, kubeadm, Gitea,
Tailscale, local registry, and public edge routing. Do not recommend broad
automatic fixes. Pick exactly ONE Lynis finding to focus on first.

Rank by risk and operational value. Prefer findings that reduce real exposure
without likely breaking SSH, Docker networking, Kubernetes, or remote access.

Return exactly this format:
Finding: <one selected finding>
Why highest risk: <one sentence>
Next safe action: <one concrete manual validation or narrow remediation>
Automation stance: report-only; add an explicit Jeannie command only after review

Homelab context:
{context}

Lynis findings:
{findings}
"""

payload = {
    "model": model,
    "prompt": prompt,
    "stream": False,
    "options": {
        "temperature": 0.1,
        "num_predict": 220,
    },
}

try:
    request = urllib.request.Request(
        f"{endpoint.rstrip('/')}/api/generate",
        data=json.dumps(payload).encode("utf-8"),
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    with urllib.request.urlopen(request, timeout=timeout_seconds) as response:
        result = json.loads(response.read().decode("utf-8"))
except (OSError, TimeoutError, urllib.error.URLError, json.JSONDecodeError):
    raise SystemExit(1)

answer = (result.get("response") or "").strip()
if not answer:
    raise SystemExit(1)

print("\n== Highest-Risk Lynis Finding ==")
print(answer)
PY
}

print_lynis_focus() {
    local report_file="$1"
    local findings_file
    local hardening_index

    hardening_index="$(awk -F= '$1 == "hardening_index" { print $2; exit }' "${report_file}" 2>/dev/null || true)"

    printf '\n== Lynis Host Audit ==\n'
    if [[ -n "${hardening_index}" ]]; then
        printf 'Hardening index: %s\n' "${hardening_index}"
    fi

    findings_file="$(mktemp)"
    lynis_findings "${report_file}" >"${findings_file}"
    if ! print_lynis_llm_focus "${findings_file}"; then
        print_lynis_fallback_focus "${findings_file}"
    fi
    rm -f "${findings_file}"
}

run_lynis_host() {
    local stamp
    local report_file
    local log_file
    stamp="$(timestamp)"
    ensure_report_dir
    report_file="${SECURITY_REPORT_DIR}/lynis-host-${stamp}.dat"
    log_file="${SECURITY_REPORT_DIR}/lynis-host-${stamp}.log"

    if ! command -v lynis >/dev/null 2>&1; then
        cat >&2 <<'EOF'
lynis is not installed on this host.

Install it on Debian with:

  sudo apt install lynis

Jeannie does not install Lynis automatically because host hardening audits should
remain explicit and report-only.
EOF
        exit 1
    fi

    echo "Running Lynis host audit in report-only mode..."
    sudo lynis audit system \
        --no-colors \
        --quiet \
        --report-file "${report_file}" \
        --log-file "${log_file}" >/dev/null

    print_lynis_focus "${report_file}"
    printf '\nFull Lynis report: %s\n' "${report_file}"
    printf 'Full Lynis log:    %s\n' "${log_file}"
}

case "${1:-all}" in
    all)
        run_trivy
        run_gitleaks
        run_zap
        ;;
    prepare)
        prepare_security_scans
        ;;
    trivy)
        run_trivy
        ;;
    secrets)
        run_gitleaks
        ;;
    zap)
        run_zap
        ;;
    nuclei)
        run_nuclei
        ;;
    web)
        run_security_web
        ;;
    kube-bench)
        run_kube_bench
        ;;
    host)
        run_lynis_host
        ;;
    -h | --help | help)
        usage
        ;;
    *)
        usage >&2
        exit 1
        ;;
esac