#!/usr/bin/env bash
#
# fix-openvas-scan-0.sh
#
# Fixes Kali/OpenVAS scans failing at 0% with ospd-openvas crashing on:
#   TypeError: cannot pickle '_thread.RLock' object
#
# Root cause:
#   Python 3.14 changed the default multiprocessing start method away from
#   'fork'. ospd-openvas keeps a Redis connection pool (containing an
#   unpicklable _thread.RLock) on the daemon object, so starting a queued
#   scan process now fails while pickling and the scanner dies at 0%.
#
# The official fix (ospd-openvas v22.10.5, commit a99fe321,
# https://github.com/greenbone/ospd-openvas/releases) adds one line to
# ospd/ospd.py:
#
#     multiprocessing.set_start_method('fork', force=True)
#
# This script applies exactly that fix to the ospd package installed on the
# system, backs up the original file, verifies the result and restarts the
# scanner service.
#
# Usage:  sudo ./fix-openvas-scan-0.sh
#
set -euo pipefail

RED=$'\e[31m'; GREEN=$'\e[32m'; YELLOW=$'\e[33m'; BOLD=$'\e[1m'; RESET=$'\e[0m'
info()  { printf '%s[*]%s %s\n' "$BOLD" "$RESET" "$*"; }
ok()    { printf '%s[+]%s %s\n' "$GREEN" "$RESET" "$*"; }
warn()  { printf '%s[!]%s %s\n' "$YELLOW" "$RESET" "$*"; }
die()   { printf '%s[x]%s %s\n' "$RED" "$RESET" "$*" >&2; exit 1; }

# --- 0. Parse arguments -----------------------------------------------------
# "check" -> only report whether the fix is already applied, then exit.
CHECK_ONLY=0
if [[ "${1:-}" == "check" ]]; then
    CHECK_ONLY=1
fi

# --- 1. Must run as root ----------------------------------------------------
# Check-only mode is read-only, so it does not require root.
if [[ "${CHECK_ONLY}" -eq 0 && ${EUID:-$(id -u)} -ne 0 ]]; then
    die "Please run as root:  sudo $0"
fi

# --- 2. Locate the ospd package -------------------------------------------
info "Locating the installed 'ospd' Python package..."
OSPD_PY=""
if command -v python3 >/dev/null 2>&1; then
    OSPD_PY=$(python3 - <<'PY' 2>/dev/null || true
import importlib.util
spec = importlib.util.find_spec("ospd.ospd")
print(spec.origin if spec and spec.origin else "")
PY
)
fi

# Fall back to common Debian/Kali locations if import lookup failed.
if [[ -z "${OSPD_PY}" || ! -f "${OSPD_PY}" ]]; then
    for cand in \
        /usr/lib/python3/dist-packages/ospd/ospd.py \
        /usr/local/lib/python3*/dist-packages/ospd/ospd.py \
        /usr/lib/python3*/site-packages/ospd/ospd.py; do
        if [[ -f "$cand" ]]; then OSPD_PY="$cand"; break; fi
    done
fi

[[ -n "${OSPD_PY}" && -f "${OSPD_PY}" ]] \
    || die "Could not find ospd/ospd.py. Is ospd-openvas installed?"
ok "Found: ${OSPD_PY}"

# --- 3. Check whether the fix is already present ----------------------------
if grep -Eq "set_start_method\(\s*['\"]fork['\"]" "${OSPD_PY}"; then
    ok "Fix already present in ${OSPD_PY}."
    [[ "${CHECK_ONLY}" -eq 1 ]] && exit 0
    PATCHED_ALREADY=1
else
    if [[ "${CHECK_ONLY}" -eq 1 ]]; then
        warn "Fix NOT applied in ${OSPD_PY}. Run without 'check' as root to apply it."
        exit 1
    fi
    PATCHED_ALREADY=0

    # --- 4. Back up the original --------------------------------------------
    BACKUP="${OSPD_PY}.bak.$(date +%Y%m%d%H%M%S)"
    cp -a "${OSPD_PY}" "${BACKUP}"
    ok "Backup created: ${BACKUP}"

    # --- 5. Apply the patch with Python (safe, idempotent) ------------------
    info "Applying multiprocessing start-method fix..."
    python3 - "${OSPD_PY}" <<'PY'
import io, re, sys

path = sys.argv[1]
with io.open(path, "r", encoding="utf-8") as f:
    src = f.read()

fix_line = "multiprocessing.set_start_method('fork', force=True)\n"

# Ensure the multiprocessing module is imported.
if not re.search(r"^\s*import\s+multiprocessing\b", src, re.MULTILINE):
    # Insert the import after the last top-level import line.
    lines = src.splitlines(keepends=True)
    last_import = 0
    for i, line in enumerate(lines):
        if re.match(r"^\s*(import|from)\s+\S", line):
            last_import = i + 1
    lines.insert(last_import, "import multiprocessing\n")
    src = "".join(lines)

# Insert the fix right after the module logger definition (matches upstream),
# or after the imports as a fallback.
if "set_start_method('fork'" not in src and 'set_start_method("fork"' not in src:
    anchor = re.search(r"^logger\s*=\s*logging\.getLogger\(__name__\).*$",
                       src, re.MULTILINE)
    if anchor:
        idx = anchor.end()
        src = src[:idx] + "\n\n" + fix_line + src[idx:]
    else:
        lines = src.splitlines(keepends=True)
        last_import = 0
        for i, line in enumerate(lines):
            if re.match(r"^\s*(import|from)\s+\S", line):
                last_import = i + 1
        lines.insert(last_import, "\n" + fix_line)
        src = "".join(lines)

with io.open(path, "w", encoding="utf-8") as f:
    f.write(src)

print("patched")
PY

    # --- 6. Verify the patched file still compiles --------------------------
    info "Verifying Python syntax of the patched file..."
    if ! python3 -m py_compile "${OSPD_PY}"; then
        warn "Syntax check failed. Restoring original from backup."
        cp -a "${BACKUP}" "${OSPD_PY}"
        die "Patch reverted. No changes made to a working file."
    fi
    grep -Eq "set_start_method\(\s*['\"]fork['\"]" "${OSPD_PY}" \
        || { cp -a "${BACKUP}" "${OSPD_PY}"; die "Fix not detected after patch; reverted."; }
    ok "Patch applied and verified."
fi

# --- 7. Restart the scanner service ----------------------------------------
SERVICE=""
for svc in ospd-openvas gvmd; do
    if systemctl list-unit-files 2>/dev/null | grep -q "^${svc}\.service"; then
        SERVICE="ospd-openvas"; break
    fi
done

if [[ -n "${SERVICE}" ]] && command -v systemctl >/dev/null 2>&1; then
    info "Restarting ospd-openvas service..."
    systemctl restart ospd-openvas.service || warn "Could not restart ospd-openvas via systemctl."
    sleep 2
    if systemctl is-active --quiet ospd-openvas.service; then
        ok "ospd-openvas is running."
    else
        warn "ospd-openvas is not active. Check: journalctl -u ospd-openvas -e"
    fi
else
    warn "systemd ospd-openvas.service not found. Restart the scanner manually."
fi

echo
if [[ "${PATCHED_ALREADY}" -eq 1 ]]; then
    ok "System already patched. If scans still fail, ensure the feed sync is complete."
else
    ok "Done. Start a scan again — it should now progress past 0%."
fi
info "Verify with:   gvm-check-setup    and    journalctl -u ospd-openvas -e"