#!/usr/bin/python3
# Monitoring plugin to check if some packages are upgradeable (even if they are marked for keep).
#
# Author: Benjamin Renard <brenard@easter-eggs.com>
# Source: https://gitlab.easter-eggs.com/brenard/check_apt_package
#
# Copyright (C) 2024 Easter-eggs
#
# check_apt_package is free software; you can redistribute
# it and/or modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# This software is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
# cSpell:words depcache
"""
Monitoring plugin to check if some packages are upgradeable (even if they are marked for keep).
"""

import argparse
import json
import logging
import os.path
import re
import sys
from json.decoder import JSONDecodeError

import apt_pkg

default_cache_path = os.path.join("/var/cache", f"{os.path.basename(sys.argv[0])}.json")

parser = argparse.ArgumentParser(description=__doc__)

parser.add_argument("-d", "--debug", action="store_true", help="Show debug messages")
parser.add_argument("-v", "--verbose", action="store_true", help="Show verbose messages")
parser.add_argument("-w", "--warning", action="store_true", help="Show warning messages")
parser.add_argument("-l", "--log-file", action="store", type=str, help="Log file path")
parser.add_argument(
    "-C",
    "--console",
    action="store_true",
    help="Also log on console (even if log file is provided)",
)

parser.add_argument("packages", nargs="+", help="Packages to check")
parser.add_argument(
    "-k",
    "--check-mark-for-keep",
    action="store_true",
    help="Check that specified packages are marked for keep",
)
parser.add_argument(
    "-I",
    "--ignore-upgrade",
    action="store_true",
    help="Do not trigger alert if packages are upgradable",
)

clean_version_args = parser.add_mutually_exclusive_group()
clean_version_args.add_argument(
    "-i",
    "--ignore-characters-after",
    help=(
        "Ignore characters in package versions after a specified character. For example, "
        "with '-i \"-\"', the version '6.2-foo' will be cut as '6.2' before comparison."
    ),
)
clean_version_args.add_argument(
    "-c",
    "--clean-version",
    type=re.compile,
    help=(
        "Regex to use to clean package versions. The specified regex must match one part that "
        r"will be used as cleaned package version. Eg: with '^([0-9\.]+).*', the version '6.2-foo'"
        "will be cleaned as '6.2' before comparison. Note: if have multiple capturing groups in "
        "your regex, the fist one will be used. To force to use another capturing group, named "
        r"it as 'version', eg: '^(?P<version>[0-9\.]+).*'."
    ),
)

parser.add_argument(
    "-a",
    "--ack-version",
    action="store_true",
    help="Acknowledge the available version of the packages to prevent further alerts about them.",
)
parser.add_argument(
    "-p",
    "--cache-path",
    help=(
        "Cache file path to store acknowledged packages's version "
        f"(default: {default_cache_path})."
    ),
    default=default_cache_path,
)

args = parser.parse_args()

# Initialize log
log = logging.getLogger()
log_format = logging.Formatter(
    "%(asctime)s - " + os.path.basename(sys.argv[0]) + " - %(levelname)s : %(message)s"
)

if args.debug:
    log.setLevel(logging.DEBUG)
elif args.verbose:
    log.setLevel(logging.INFO)
elif args.warning:
    log.setLevel(logging.WARNING)
else:
    log.setLevel(logging.FATAL)

if args.log_file:
    log_file = logging.FileHandler(args.log_file)
    log_file.setFormatter(log_format)
    log.addHandler(log_file)

if not args.log_file or args.console:
    log_console = logging.StreamHandler()
    log_console.setFormatter(log_format)
    log.addHandler(log_console)


def clean_version(input_version):
    """Clean specified package version based on script arguments"""
    version = str(input_version)
    if args.ignore_characters_after:
        version = version.split(args.ignore_characters_after, maxsplit=1)[0]
        logging.debug(
            "Version '%s' cleaned as '%s' using -i/--ignore-characters-after '%s'",
            input_version,
            version,
            args.ignore_characters_after,
        )
    elif args.clean_version:
        matches = args.clean_version.match(version)
        if matches:
            # cSpell:words groupdict
            version = matches.groupdict().get("version", matches.group(1))
            logging.debug(
                "Version '%s' cleaned as '%s' using -c/--clean-version '%s'",
                input_version,
                version,
                args.clean_version.pattern,
            )
    return version


def get_current_and_candidate_package_version(pkg):
    """Get current and candidate versions of the specified package"""
    current_version_str = clean_version(pkg.current_ver.ver_str)

    # Retrieve candidate version based on APT policy (avoids selecting unprioritized backports)
    candidate_pkg = depcache.get_candidate_ver(pkg)
    if not candidate_pkg:
        return current_version_str, current_version_str

    candidate_version_str = clean_version(candidate_pkg.ver_str)

    # Check if a candidate version is different (assume newer) and can be downloaded/resolved
    if candidate_version_str != current_version_str and candidate_pkg.downloadable:
        return (current_version_str, candidate_version_str)

    return current_version_str, current_version_str


cache_data = None
cache_data_updated = False


def load_cache_data():
    """Load cache data from JSON file"""
    global cache_data  # pylint: disable=global-statement
    if cache_data is not None:
        return True
    if os.path.exists(args.cache_path):
        try:
            with open(args.cache_path, encoding="utf8") as file_desc:
                cache_data = json.load(file_desc)
            return True
        except OSError as err:
            errors.append(f"Failed to load cache file ({args.cache_path}): {err}")
        except JSONDecodeError as err:
            errors.append(f"Failed to parse JSON data from cache file ({args.cache_path}): {err}")
    cache_data = {}
    return False


def save_cache_data():
    """Save cache data to JSON file"""
    if not cache_data_updated:
        logging.debug("Cache data not updated, don't update cache data file (%s)", args.cache_path)
        return True
    print()
    try:
        with open(args.cache_path, "w", encoding="utf8") as file_desc:
            json.dump(cache_data, file_desc)
        logging.info("Cache data file updated (%s)", args.cache_path)
        return True
    except OSError as err:
        errors.append(f"Failed to load cache file ({args.cache_path}): {err}")
    return False


def is_acknowledged_package_version(pkg, version):
    """Check if specified package candidate version is acknowledged"""
    load_cache_data()
    return cache_data.get("ack_packages_versions", {}).get(pkg) == version


def acknowledge_package_version(pkg, version):
    """Set specified package candidate version as acknowledged"""
    global cache_data_updated  # pylint: disable=global-statement
    load_cache_data()
    if not isinstance(cache_data.get("ack_packages_versions"), dict):
        cache_data["ack_packages_versions"] = {}
    cache_data["ack_packages_versions"][pkg] = version
    print(f"Version {version} of package {pkg} acknowledged")
    cache_data_updated = True


apt_pkg.init()
cache = apt_pkg.Cache(None)
# Use DepCache to resolve dependencies and respect APT policy/pinning rules
depcache = apt_pkg.DepCache(cache)

errors = []
unknown_packages = []
upgradable_packages = {}
up_to_date_packages = {}
ack_packages_versions = {}

logging.info("Specified packages: %s", ", ".join(args.packages))
for package_name in args.packages:
    try:
        package = cache[package_name]
    except KeyError:
        unknown_packages.append(package_name)

    if package.current_state != apt_pkg.CURSTATE_INSTALLED:  # cSpell:words curstate
        errors.append(f"Package {package_name} is not installed")
        continue

    current_version, candidate_version = get_current_and_candidate_package_version(package)
    if current_version != candidate_version:
        if is_acknowledged_package_version(package_name, candidate_version):
            ack_packages_versions[package_name] = (
                f"{current_version} installed, {candidate_version} available but acknowledged"
            )
        else:
            upgradable_packages[package_name] = f"{current_version} => {candidate_version}"
            if args.ack_version:
                acknowledge_package_version(package_name, candidate_version)
    else:
        up_to_date_packages[package_name] = current_version

    if args.check_mark_for_keep:
        if package.selected_state != apt_pkg.INSTSTATE_HOLD:  # cSpell:words inststate
            errors.append(f"Package {package_name} is not marked for keep")
        else:
            logging.debug("Package %s is marked for keep", package_name)
logging.info(
    "Unknown packages: %s",
    ", ".join(unknown_packages) if unknown_packages else "no unknown package",
)
logging.info(
    "Up-to-date packages: %s",
    ", ".join(up_to_date_packages) if up_to_date_packages else "no up-to-date package",
)
logging.info(
    "Upgradable packages: %s",
    ", ".join(upgradable_packages) if upgradable_packages else "no upgradable package",
)
logging.info(
    "Upgradable packages to acknowledged version : %s",
    (
        ", ".join(ack_packages_versions)
        if ack_packages_versions
        else "no upgradable package to acknowledged version "
    ),
)

if args.ack_version:
    save_cache_data()
    if errors:
        print("\n".join(errors))
    sys.exit(1 if errors else 0)

STATUS = "OK"
EXIT_CODE = 0
MESSAGES = []
if unknown_packages:
    MESSAGES.append(
        f"Package {unknown_packages[0]} is upgradable"
        if len(unknown_packages) == 1
        else f"{len(unknown_packages)} upgradable packages"
    )
    EXIT_CODE = 3
if upgradable_packages:
    message = (
        f"Package {next(iter(upgradable_packages))} is upgradable"
        if len(upgradable_packages) == 1
        else f"{len(upgradable_packages)} upgradable packages"
    )
    if ack_packages_versions:
        message += " and " + (
            f"package {next(iter(ack_packages_versions))} is upgradable to an acknowledged version"
            if len(ack_packages_versions) == 1
            else f"{len(ack_packages_versions)} upgradable packages to an acknowledged version"
        )
    if not args.ignore_upgrade:
        errors.insert(0, message)
    MESSAGES.append(message)
    EXIT_CODE = max(EXIT_CODE, 1)
    STATUS = STATUS if STATUS != "OK" else "WARNING"
elif ack_packages_versions:
    if up_to_date_packages:
        MESSAGES.append(
            f"Package {next(iter(up_to_date_packages))} is up-to-date"
            if len(up_to_date_packages) == 1
            else f"{len(up_to_date_packages)} up-to-date packages"
        )
    MESSAGES.append(
        f"Package {next(iter(ack_packages_versions))} is upgradable to an acknowledged version"
        if len(ack_packages_versions) == 1
        else f"{len(ack_packages_versions)} upgradable packages to an acknowledged version"
    )
else:
    MESSAGES.append(
        "All packages are up-to-date"
        if len(args.packages) > 1
        else f"Package {args.packages[0]} is up-to-date"
    )

print(f"{STATUS} - {' / '.join(MESSAGES)}")


if unknown_packages:
    print()
    print(f"The {len(unknown_packages)} following packages are unknown:")
    print("\n".join([f"- {msg}" for msg in unknown_packages]))

if upgradable_packages:
    print()
    print(f"The {len(upgradable_packages)} following packages are upgradable:")
    print("\n".join([f"- {pkg} ({msg})" for pkg, msg in upgradable_packages.items()]))

if ack_packages_versions:
    print()
    print(
        f"The {len(ack_packages_versions)} following packages are upgradable to acknowledged "
        "version:"
    )
    print("\n".join([f"- {pkg} ({msg})" for pkg, msg in ack_packages_versions.items()]))

if up_to_date_packages:
    print()
    print(f"The {len(up_to_date_packages)} following packages are already up-to-date:")
    print("\n".join([f"- {pkg} ({msg})" for pkg, msg in up_to_date_packages.items()]))

sys.exit(EXIT_CODE)
