From 783eaea96a9768971400eb63dfe659be958f3c18 Mon Sep 17 00:00:00 2001
From: Melora Hugues <dev@faercol.me>
Date: Fri, 13 Dec 2024 10:57:10 +0100
Subject: [PATCH] Add reboot check to apt collector

---
 collector/apt/apt.go | 30 ++++++++++++++++++++++++++++--
 1 file changed, 28 insertions(+), 2 deletions(-)

diff --git a/collector/apt/apt.go b/collector/apt/apt.go
index b0cf396..93d1b66 100644
--- a/collector/apt/apt.go
+++ b/collector/apt/apt.go
@@ -2,7 +2,10 @@ package apt
 
 import (
 	"context"
+	"errors"
 	"fmt"
+	"io/fs"
+	"os"
 	"os/exec"
 	"strings"
 	"time"
@@ -12,9 +15,14 @@ import (
 	"go.uber.org/zap"
 )
 
+const rebootRequiredPath = "/var/run/reboot-required"
+
 type AptCollector struct {
 	installedPackages int
 	pendingUpdates    int
+	rebootRequired    int
+
+	rebootRequiredPath string
 
 	promExporter *prometheus.GaugeVec
 	updateFreq   time.Duration
@@ -56,6 +64,19 @@ func (c *AptCollector) updatePendingUpdates() error {
 	return nil
 }
 
+func (c *AptCollector) updateRebootRequired() error {
+	_, err := os.Stat(c.rebootRequiredPath)
+	if err != nil {
+		if errors.Is(err, fs.ErrNotExist) {
+			c.rebootRequired = 0
+			return nil
+		}
+		return fmt.Errorf("failed to check if reboot-required exists: %w", err)
+	}
+	c.rebootRequired = 1
+	return nil
+}
+
 func (c *AptCollector) update() error {
 	if err := c.updateInstalledPackages(); err != nil {
 		return fmt.Errorf("failed to update count of installed packages: %w", err)
@@ -63,9 +84,13 @@ func (c *AptCollector) update() error {
 	if err := c.updatePendingUpdates(); err != nil {
 		return fmt.Errorf("failed to update count of pending updates: %w", err)
 	}
+	if err := c.updateRebootRequired(); err != nil {
+		return fmt.Errorf("failed to update reboot status: %w", err)
+	}
 
 	c.promExporter.WithLabelValues("total").Set(float64(c.installedPackages))
 	c.promExporter.WithLabelValues("updates").Set(float64(c.pendingUpdates))
+	c.promExporter.WithLabelValues("reboot_required").Set(float64(c.rebootRequired))
 	return nil
 }
 
@@ -97,8 +122,9 @@ func (c *AptCollector) PromCollector() prometheus.Collector {
 
 func New() *AptCollector {
 	return &AptCollector{
-		updateFreq:   30 * time.Minute,
-		promExporter: prometheus.NewGaugeVec(prometheus.GaugeOpts{Namespace: "packages", Subsystem: "apt", Name: "packages_count", Help: "Count of apt packages"}, []string{"status"}),
+		rebootRequiredPath: rebootRequiredPath,
+		updateFreq:         30 * time.Minute,
+		promExporter:       prometheus.NewGaugeVec(prometheus.GaugeOpts{Namespace: "packages", Subsystem: "apt", Name: "packages_count", Help: "Count of apt packages"}, []string{"status"}),
 	}
 }