Some fixes and improvments

Signed-off-by: RomainMou <[email protected]>

needrestart_info.py

@@ -1,19 +1,21 @@
 #!/usr/bin/env python3
-#
-#
-# Description: Expose metrics from needrestart.
-#
-# This script runs needrestart in batch mode. It will never ask for input
-# and will never restart or upgrade anything.
-#
-# Dependencies: python >= 3.5, python3-prometheus-client, needrestart
-#
-# Authors: RomainMou
+
+"""
+Description: Expose metrics from needrestart.
+
+This script runs needrestart in batch mode. It will never ask for input
+and will never restart or upgrade anything.
+
+Dependencies: python >= 3.5, python3-prometheus-client, needrestart
+
+Authors: RomainMou
+"""
 
 import time
 import subprocess
 from collections import Counter
 from enum import Enum
+
 from prometheus_client import (
     CollectorRegistry,
     Gauge,
@@ -34,7 +36,7 @@ class MicroCodeStatus(Enum):
     OBSELETE = 2
 
 
-class NeedrestartParser:
+class NeedRestartData:
     def __init__(self, needrestart_output):
         # Some default value
         self.timestamp = int(time.time())
@@ -48,8 +50,8 @@ def __init__(self, needrestart_output):
         needrestart_counter = Counter()
 
         # Parse the cmd output
-        for line in needrestart_output.stdout.decode().splitlines():
-            key, value = line.split(": ")
+        for line in needrestart_output.splitlines():
+            key, value = line.split(": ", maxsplit=1)
             if key == "NEEDRESTART-VER":
                 self.version = value
             # Kernel informations
@@ -75,7 +77,7 @@ def __init__(self, needrestart_output):
         self.sessions_count = needrestart_counter["NEEDRESTART-SESS"]
 
 
-def _write_timestamp(registry, needrestart_data):
+def write_timestamp(registry, needrestart_data):
     g = Gauge(
         "needrestart_timestamp",
         "information about the version and when it was last run",
@@ -85,7 +87,7 @@ def _write_timestamp(registry, needrestart_data):
     g.labels(needrestart_data.version).set(needrestart_data.timestamp)
 
 
-def _write_kernel(registry, needrestart_data):
+def write_kernel(registry, needrestart_data):
     if needrestart_data.kernel_status:
         e = Gauge(
             "needrestart_kernel_status",
@@ -99,7 +101,7 @@ def _write_kernel(registry, needrestart_data):
         ).set(needrestart_data.kernel_status.value)
 
 
-def _write_microcode(registry, needrestart_data):
+def write_microcode(registry, needrestart_data):
     if needrestart_data.microcode_status:
         e = Gauge(
             "needrestart_microcode_status",
@@ -113,7 +115,7 @@ def _write_microcode(registry, needrestart_data):
         ).set(needrestart_data.microcode_status.value)
 
 
-def _write_services(registry, needrestart_data):
+def write_services(registry, needrestart_data):
     g = Gauge(
         "needrestart_services_count",
         "number of services requiring a restart",
@@ -122,7 +124,7 @@ def _write_services(registry, needrestart_data):
     g.set(needrestart_data.services_count)
 
 
-def _write_containers(registry, needrestart_data):
+def write_containers(registry, needrestart_data):
     g = Gauge(
         "needrestart_containers_count",
         "number of containers requiring a restart",
@@ -131,7 +133,7 @@ def _write_containers(registry, needrestart_data):
     g.set(needrestart_data.containers_count)
 
 
-def _write_sessions(registry, needrestart_data):
+def write_sessions(registry, needrestart_data):
     g = Gauge(
         "needrestart_sessions_count",
         "number of sessions requiring a restart",
@@ -140,21 +142,23 @@ def _write_sessions(registry, needrestart_data):
     g.set(needrestart_data.sessions_count)
 
 
-def _main():
+def main():
     registry = CollectorRegistry()
-    needrestart_data = NeedrestartParser(
-        subprocess.run(["needrestart", "-b"], stdout=subprocess.PIPE)
-    )
 
-    _write_timestamp(registry, needrestart_data)
-    _write_kernel(registry, needrestart_data)
-    _write_microcode(registry, needrestart_data)
-    _write_services(registry, needrestart_data)
-    _write_containers(registry, needrestart_data)
-    _write_sessions(registry, needrestart_data)
+    needrestart_output = subprocess.run(
+        ["needrestart", "-b"], capture_output=True, text=True
+    ).stdout
+    needrestart_data = NeedRestartData(needrestart_output)
+
+    write_timestamp(registry, needrestart_data)
+    write_kernel(registry, needrestart_data)
+    write_microcode(registry, needrestart_data)
+    write_services(registry, needrestart_data)
+    write_containers(registry, needrestart_data)
+    write_sessions(registry, needrestart_data)
 
     print(generate_latest(registry).decode(), end="")
 
 
 if __name__ == "__main__":
-    _main()
+    main()