#!/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/>.
"""
Monitoring plugin to check if some packages are upgradeable (even if they are marked for keep).
"""

import argparse
import logging
import os.path
import re
import sys

import apt_pkg

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", "--logfile", 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\.]+).*'."
    ),
)


args = parser.parse_args()

# Initialize log
log = logging.getLogger()
logformat = 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.logfile:
    logfile = logging.FileHandler(args.logfile)
    logfile.setFormatter(logformat)
    log.addHandler(logfile)

if not args.logfile or args.console:
    logconsole = logging.StreamHandler()
    logconsole.setFormatter(logformat)
    log.addHandler(logconsole)


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:
            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


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 = []
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:
        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:
        upgradable_packages.append(f"{package_name} ({current_version} => {candidate_version})")
    else:
        up_to_date_packages.append(f"{package_name} ({current_version})")

    if args.check_mark_for_keep:
        if package.selected_state != apt_pkg.INSTSTATE_HOLD:
            errors.append(f"Package {package_name} is not marked for keep")
        else:
            logging.debug("Package %s is marked for keep", package_name)
logging.info(
    "Upgradable packages: %s",
    ", ".join(upgradable_packages) if upgradable_packages else "no upgradable packages",
)

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 {upgradable_packages[0]} is upgradable"
        if len(upgradable_packages) == 1
        else f"{len(upgradable_packages)} upgradable packages"
    )
    if not args.ignore_upgrade:
        errors.insert(0, message)
    MESSAGES.append(message)
    EXIT_CODE = max(EXIT_CODE, 1)
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"- {msg}" for msg in upgradable_packages]))

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"- {msg}" for msg in up_to_date_packages]))

sys.exit(EXIT_CODE)
