From 12954d24988dcc747940f6f0f2f6402f0ea031d6 Mon Sep 17 00:00:00 2001 From: egvimo Date: Mon, 18 Sep 2023 12:51:56 +0000 Subject: [PATCH] feat: allow setting different root dir for apt Signed-off-by: egvimo --- apt_info.py | 31 ++++++++++++++++++++++--------- 1 file changed, 22 insertions(+), 9 deletions(-) diff --git a/apt_info.py b/apt_info.py index eb1a642..ab73749 100755 --- a/apt_info.py +++ b/apt_info.py @@ -8,10 +8,11 @@ # Authors: Kyle Fazzari # Daniel Swarbrick -import apt +import argparse import collections import contextlib import os +import apt from prometheus_client import CollectorRegistry, Gauge, generate_latest _UpgradeInfo = collections.namedtuple("_UpgradeInfo", ["labels", "count"]) @@ -26,12 +27,12 @@ def _convert_candidates_to_upgrade_infos(candidates): ) changes_dict[",".join(origins)][candidate.architecture] += 1 - changes_list = list() + changes_list = [] for origin in sorted(changes_dict.keys()): for arch in sorted(changes_dict[origin].keys()): changes_list.append( _UpgradeInfo( - labels=dict(origin=origin, arch=arch), + labels={"origin": origin, "arch": arch}, count=changes_dict[origin][arch], ) ) @@ -74,14 +75,14 @@ def _write_autoremove_pending(registry, cache): g.set(len(autoremovable_packages)) -def _write_reboot_required(registry): +def _write_reboot_required(registry, root_dir): g = Gauge('node_reboot_required', "Node reboot is required for software updates.", registry=registry) - g.set(int(os.path.isfile('/run/reboot-required'))) + g.set(int(os.path.isfile(os.path.join(root_dir, 'run/reboot-required')))) -def _main(): - cache = apt.cache.Cache() +def generate_metrics(root_dir: str = '/') -> bytes: + cache = apt.cache.Cache(rootdir=root_dir) # First of all, attempt to update the index. If we don't have permission # to do so (or it fails for some reason), it's not the end of the world, @@ -96,8 +97,20 @@ def _main(): _write_pending_upgrades(registry, cache) _write_held_upgrades(registry, cache) _write_autoremove_pending(registry, cache) - _write_reboot_required(registry) - print(generate_latest(registry).decode(), end='') + _write_reboot_required(registry, root_dir) + + return generate_latest(registry) + + +def _main(): + parser = argparse.ArgumentParser() + parser.add_argument('-r', '--root-dir', dest='root_dir', type=str, default='/', + help="Set root directory to a different path than /") + args = parser.parse_args() + + metrics = generate_metrics(args.root_dir) + + print(metrics.decode(), end='') if __name__ == "__main__":