#!/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 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",
)
parser.add_argument(
    "-i",
    "--ignore-after",
    action="store",
    type=str,
    help="Ignore pkg version after a character, useful to compare ex. 6.2 instead of 6.2-foo",
)

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)

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 = []
upgradable_packages = []
up_to_date_packages = []
logging.info("Specified packages: %s", ", ".join(args.packages))
for package in args.packages:
    try:
        pkg = cache[package]
    except KeyError:
        print(f"UNKNOWN - Package {package} is not available")
        sys.exit(3)

    if pkg.current_state != apt_pkg.CURSTATE_INSTALLED:
        errors.append(f"Package {package} is not installed")
        continue

    # Retrieve candidate version based on APT policy (avoids selecting unprioritized backports)
    candidate_version = depcache.get_candidate_ver(pkg)

    # Check if a candidate version exists, is newer, and can be downloaded/resolved
    if (
        candidate_version
        and candidate_version != pkg.current_ver
        and candidate_version.downloadable
    ):
        candidate_version_str = candidate_version.ver_str.split(args.igore_after, maxsplit=1)[0]
        pkg_version_str = pkg.current_ver.ver_str.split(args.igore_after, maxsplit=1)[0]
        if candidate_version_str != pkg_version_str:
            upgradable_packages.append(
                f"{package} ({pkg.current_ver.ver_str} => {candidate_version.ver_str})"
            )
    else:
        up_to_date_packages.append(f"{package} ({pkg.current_ver.ver_str})")

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

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)
else:
    message = (
        "All packages are up-to-date"
        if len(args.packages) > 1
        else f"Package {args.packages[0]} is up-to-date"
    )

EXIT_CODE = 0
if not errors:
    print(f"OK - {message}")
else:
    print(f"WARNING - {', '.join(errors)}")
    EXIT_CODE = 1

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)
