feat: allow setting different root dir for apt

Signed-off-by: egvimo <egvimo@gmail.com>

Author: egvimo

apt_info.py

@@ -8,6 +8,7 @@
 # Authors: Kyle Fazzari <kyrofa@ubuntu.com>
 #          Daniel Swarbrick <dswarbrick@debian.org>
 
+import argparse
 import apt
 import collections
 import contextlib
@@ -74,14 +75,19 @@ 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()
+    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()
+
+    cache = apt.cache.Cache(rootdir=args.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,7 +102,7 @@ def _main():
     _write_pending_upgrades(registry, cache)
     _write_held_upgrades(registry, cache)
     _write_autoremove_pending(registry, cache)
-    _write_reboot_required(registry)
+    _write_reboot_required(registry, args.root_dir)
     print(generate_latest(registry).decode(), end='')