#!/usr/bin/env python3
"""
check_python_venv_security_updates

Icinga-compatible check plugin that:
 - Ensures pip-audit is installed and up to date in the specified virtualenv
 - Runs pip-audit and pip list to check for security and regular updates
 - Produces Nagios/Icinga style output with detailed results and perfdata
"""
import argparse
import json
import logging
import os
import subprocess  # nosec
import sys
from getpass import getuser
from pathlib import Path
from typing import Any, Dict, List, Tuple

from packaging.version import InvalidVersion, Version

# Nagios/Icinga exit codes
OK = 0
WARNING = 1
CRITICAL = 2
UNKNOWN = 3
STATES_LABEL = {OK: "OK", WARNING: "WARNING", CRITICAL: "CRITICAL", UNKNOWN: "UNKNOWN"}

logger = logging.getLogger(__name__)


def get_python_bin(venv_path: Path) -> Path:
    """Get python binary path inside the virtualenv."""
    return Path.joinpath(venv_path, Path("bin/python"))


def check_venv(venv_path: Path) -> Tuple[bool, Path]:
    """Check virtualenv path by checking if Python binary is present inside the virtualenv."""
    python_bin = get_python_bin(venv_path)
    return python_bin.exists(), python_bin


def run_python_cmd(
    venv_path: Path, module: str, *cmd_args: List[str]
) -> Tuple[int, str, str]:
    """Run a shell command and return (exit_code, stdout, stderr)."""
    cmd = [
        str(get_python_bin(venv_path)),
        "-m",
        module,
        *[arg for arg in cmd_args if arg is not None],
    ]
    logger.debug("Executing command: %s", " ".join(cmd))
    try:
        result = subprocess.run(  # nosec
            cmd,
            text=True,
            capture_output=True,
            check=False,
        )
        logger.debug("Return code: %d", result.returncode)
        if result.stdout:
            logger.debug("STDOUT:\n%s", result.stdout.strip())
        if result.stderr:
            logger.debug("STDERR:\n%s", result.stderr.strip())
        return result.returncode, result.stdout.strip(), result.stderr.strip()
    except OSError as err:
        logger.exception("Failed to execute command: %s", cmd)
        return 1, "", str(err)


def check_venv_owner(venv_path: Path) -> Tuple[bool, str]:
    """Check virtualenv is owned by the current user"""
    current_user = getuser()
    logger.info("Check virtualenv is owned by current user (%s)", current_user)
    venv_owner = venv_path.owner()
    if venv_owner == current_user:
        logger.info("Virtualenv is owned by the current user (%s)", current_user)
        return True, venv_owner, current_user
    return False, venv_owner, current_user


def ensure_pip_audit(venv_path: Path) -> Tuple[bool, bool, str | None]:
    """Ensure pip-audit is installed in the virtualenv."""
    logger.info("Check if pip-audit is installed...")
    code, out, _ = run_python_cmd(venv_path, "pip", "show", "pip-audit")
    if code == 0 and out:
        logger.info("pip-audit is already installed")
        return True, False, None
    logger.info("pip-audit not found, installing...")
    success, venv_owner, current_user = check_venv_owner(venv_path)
    if not success:
        return (
            False,
            False,
            f"Virtualenv is owned by {venv_owner} but check is run as {current_user}: can't "
            "install pip-audit",
        )
    code, out, err = run_python_cmd(venv_path, "pip", "install", "pip-audit")
    if code == 0:
        logger.info("pip-audit successfully installed.")
        return True, True, "pip-audit installed"
    return False, False, f"failed to install pip-audit: {err}"


def update_pip_audit(
    venv_path: Path, outdated: List[Dict[str, Any]]
) -> Tuple[bool, str]:
    """Check and update pip-audit if a new version is available."""
    logger.info("Check if pip-audit is up to date...")
    for pkg in outdated:
        if pkg["name"].lower() == "pip-audit":
            logger.info(
                "pip-audit update available: %s -> %s, install it...",
                pkg["version"],
                pkg["latest_version"],
            )
            success, venv_owner, current_user = check_venv_owner(venv_path)
            if not success:
                return (
                    False,
                    f"Virtualenv is owned by {venv_owner} but check is run as {current_user}: "
                    "can't update pip-audit",
                )
            code, out, err = run_python_cmd(
                venv_path, "pip", "install", "--upgrade", "pip-audit"
            )
            if code == 0:
                logger.info("pip-audit upgraded")
                return (
                    True,
                    f"pip-audit upgraded from {pkg['version']} to {pkg['latest_version']}",
                )
            logger.info("Failed to upgrade pip-audit, continue")
            return (
                False,
                f"failed to upgrade pip-audit ({pkg['version']} -> {pkg['latest_version']}): "
                f"{err or out}",
            )

    return True, "pip-audit already installed and up to date"


def run_pip_audit(
    venv_path: Path, ignore_system_packages: bool
) -> Tuple[bool, List[Dict[str, Any]] | str]:
    """Run pip-audit in the given virtualenv and return parsed JSON result."""
    logger.info(
        "Run pip-audit on virtualenv packages..."
        if ignore_system_packages
        else "Run pip-audit..."
    )
    _, out, _ = run_python_cmd(
        venv_path,
        "pip_audit",
        "--format",
        "json",
        "--local" if ignore_system_packages else None,
    )
    logger.info("Parse JSON output of pip-audit...")
    try:
        return True, json.loads(out)
    except json.JSONDecodeError:
        return False, "invalid JSON output from pip-audit"


def get_outdated_packages(
    venv_path: Path, ignore_system_packages: bool
) -> Tuple[bool, List[Dict[str, Any]], str | None]:
    """List outdated packages via pip list inside the virtualenv."""
    logger.info(
        "List virtualenv outdated packages..."
        if ignore_system_packages
        else "List outdated packages..."
    )
    code, out, err = run_python_cmd(
        venv_path,
        "pip",
        "list",
        "--outdated",
        "--format=json",
        "--local" if ignore_system_packages else None,
    )
    if code != 0:
        return (
            False,
            [],
            f"Failed to list outdated packages (using `pip list --outdated`): {err}",
        )
    try:
        data = json.loads(out)
        logger.info("%d outdated packages found", len(data))
        return True, data, None
    except json.JSONDecodeError:
        return (
            False,
            [],
            "Failed to parse pip list JSON output, can't get list of outdated packages",
        )


def get_target_version(
    name: str, version: str, fix_versions: List[str]
) -> str | None | bool:
    """Get target version of a package to fix its issues"""
    to_update_version = None
    try:
        cur = Version(version)
        for v in fix_versions:
            v = Version(v)
            if v < cur:
                continue
            if to_update_version is None or v < to_update_version:
                to_update_version = v
    except InvalidVersion as err:
        logger.debug(
            "Package %s: %s. Can't determine the right targeted version to fix its issue",
            name,
            err,
        )
        return False
    return to_update_version


def main():
    """Entrypoint"""
    parser = argparse.ArgumentParser(
        description="Icinga plugin to audit Python virtualenv with pip-audit."
    )
    parser.add_argument(
        "venv_path",
        type=lambda x: Path(x).resolve(),
        help="Path to the Python virtualenv to audit.",
    )
    parser.add_argument(
        "-n",
        "--dont-warn-about-not-fixed-vuln",
        action="store_true",
        help="Dont warn about not fixed vulnerabilities",
    )
    parser.add_argument(
        "-I",
        "--ignore-system-packages",
        action="store_true",
        help="Ignore packages installed outside the virtualenv "
        "(for virtualenv created with --system-site-packages parameter)",
    )
    parser.add_argument(
        "-w",
        "--warn-on-nonsecurity-update",
        action="store_true",
        help="Treat non-security updates as WARNING instead of OK.",
    )
    parser.add_argument(
        "-v", "--verbose", action="store_true", help="Enable verbose logging."
    )
    parser.add_argument(
        "-d",
        "--debug",
        action="store_true",
        help="Enable debug logging (includes raw data).",
    )
    args = parser.parse_args()

    logging.basicConfig(
        level=(
            logging.DEBUG
            if args.debug
            else logging.INFO if args.verbose else logging.WARNING
        ),
        format="%(asctime)s [%(levelname)s] %(message)s",
    )
    python_bin = os.path.join(args.venv_path, "bin", "python")

    success, python_bin = check_venv(args.venv_path)
    if not success:
        print(f"UNKNOWN - Python binary not found in virtualenv ({python_bin})")
        sys.exit(UNKNOWN)

    state = OK
    details = []
    errors = []
    perfdata = {}

    success, installed, msg = ensure_pip_audit(args.venv_path)
    if not success:
        print(f"UNKNOWN - {msg}")
        sys.exit(UNKNOWN)
    if msg:
        details.append(msg)

    success, outdated, message = get_outdated_packages(
        args.venv_path, args.ignore_system_packages
    )
    if not success and message:
        errors.append(message)

    if not installed:
        success, msg = update_pip_audit(args.venv_path, outdated)
        if not success:
            state = WARNING
            errors.append(msg)
        elif msg:
            details.append(msg)

    success, audit_data = run_pip_audit(args.venv_path, args.ignore_system_packages)
    if not success:
        print(f"UNKNOWN - {audit_data}")
        sys.exit(UNKNOWN)

    logger.info("Analyse audit result...")
    pending_security_updates = []
    not_fixed_vulnerabilities = []
    try:
        for item in audit_data["dependencies"]:
            if not item.get("vulns"):
                continue
            fix_versions = []
            vulns_ids = []
            for vuln in item["vulns"]:
                vuln_ids = [x for x in ([vuln["id"]] + vuln["aliases"]) if x]
                vulns_ids.append("/".join(vuln_ids) or "unknown")
                if vuln["fix_versions"]:
                    fix_versions.extend(vuln["fix_versions"])
                else:
                    not_fixed_vulnerabilities.append(
                        (item["name"], item["version"], vuln_ids)
                    )

            if fix_versions:
                pending_security_updates.append(
                    (
                        item["name"],
                        item["version"],
                        list(set(fix_versions)),
                        vulns_ids,
                    )
                )
    except KeyError as err:
        print(f"UNKNOWN - Failed to analyse pip-audit result ({err})")
        sys.exit(3)

    # Format output message
    perfdata = {
        "installed_packages": len(audit_data),
        "outdated_packages": len(outdated),
        "pending_security_updates": len(pending_security_updates),
        "not_fixed_vulnerabilities": len(not_fixed_vulnerabilities),
    }

    state = OK
    messages = []
    to_update = []
    if pending_security_updates:
        state = CRITICAL
        messages.append(f"{len(pending_security_updates)} security updates available")
        errors.append("\n[PENDING SECURITY UPDATES]")
        for name, version, fix_versions, vuln_ids in pending_security_updates:
            errors.append(
                f"{name} {version} -> {', '.join(fix_versions)} ({', '.join(vuln_ids)})"
            )
            to_update_version = get_target_version(name, version, fix_versions)
            if to_update_version:
                to_update.append(f"{name}=={to_update_version}")

    if not_fixed_vulnerabilities:
        if not args.dont_warn_about_not_fixed_vuln:
            state = WARNING if state < WARNING else state
        messages.append(
            f"{len(not_fixed_vulnerabilities)} known vulnerabilities without fix"
        )
        errors.append("\n[KNOWN VULNERABILITIES]")
        for name, version, vuln_ids in not_fixed_vulnerabilities:
            errors.append(f"{name} {version} ({' / '.join(vuln_ids)})")

    if outdated:
        if args.warn_on_nonsecurity_update:
            state = WARNING if state < WARNING else state
        messages.append(f"{len(outdated)} non-security updates available")
        errors.append("\n[AVAILABLE UPDATES]")
        for pkg in outdated:
            errors.append(
                f"{pkg['name']} ({pkg['version']} -> {pkg['latest_version']})"
            )
            if args.warn_on_nonsecurity_update:
                to_update.append(f"{pkg['name']}=={pkg['latest_version']}")

    print(
        " ".join(
            [
                f"{STATES_LABEL[state]} - "
                f"{', '.join(messages) if messages else 'All packages are up to date and secure'}",
                "|",
                " ".join([f"{k}={v}" for k, v in perfdata.items()]),
            ]
        )
    )

    if errors:
        print("\n".join(errors))

    if to_update:
        print("\n[UPDATE COMMAND]")
        print(f"{get_python_bin(args.venv_path)} -m pip install {' '.join(to_update)}")

    if details:
        print("\n[DETAILS]")
        print("\n".join(details))

    sys.exit(state)


if __name__ == "__main__":
    main()
