From fb0c0f85aa9f12733d7e065bb8f38a25b7a9f67a Mon Sep 17 00:00:00 2001
From: egvimo <egvimo@gmail.com>
Date: Mon, 9 Oct 2023 15:33:27 +0200
Subject: [PATCH] feat: add root dir

---
 app.py      |  5 ++++-
 apt_info.py | 25 +++++++++++++++++++------
 2 files changed, 23 insertions(+), 7 deletions(-)

diff --git a/app.py b/app.py
index 9876f27..2223044 100644
--- a/app.py
+++ b/app.py
@@ -1,11 +1,14 @@
+import os
 from fastapi import FastAPI
 from fastapi.responses import PlainTextResponse
 import apt_info
 
+root_dir = os.environ.get("APT_ROOT_DIR", "/")
+
 app = FastAPI()
 
 
 @app.get("/metrics", response_class=PlainTextResponse)
 def metrics():
-    data = apt_info._main()  # pylint: disable=protected-access
+    data = apt_info.generate_metrics(root_dir=root_dir)
     return data
diff --git a/apt_info.py b/apt_info.py
index 00055dc..be2a43f 100644
--- a/apt_info.py
+++ b/apt_info.py
@@ -10,6 +10,7 @@
 
 # pylint: disable=invalid-name
 
+import argparse
 import collections
 import contextlib
 import os
@@ -76,14 +77,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,
@@ -98,9 +99,21 @@ 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, 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__":
-    print(_main().decode(), end='')
+    _main()