POC: Refactor Distros

Author: Ricky Cousins

.gitignore

@@ -1 +1,2 @@
 bin/*
+updates-exporter

distros/distros.go

@@ -0,0 +1,6 @@
+package distros
+
+type Distro interface {
+    GetSecurityUpdates() int
+    GetRebootRequired() bool
+}

distros/ubuntu/ubuntu.go

@@ -0,0 +1,30 @@
+package ubuntu
+
+import (
+    "log"
+    "os"
+    "os/exec"
+    "strconv"
+    "strings"
+)
+
+type Ubuntu struct{}
+
+func (Ubuntu) GetSecurityUpdates() int {
+    cmd := exec.Command("sh", "-c", `apt-get -s dist-upgrade | grep "^Inst" | grep security | wc -l`)
+    output, err := cmd.Output()
+    if err != nil {
+        log.Printf("Error running apt-get: %v", err)
+        return 0
+    }
+    count, _ := strconv.Atoi(strings.TrimSpace(string(output)))
+    return count
+}
+
+func (Ubuntu) GetRebootRequired() bool {
+    if _, err := os.Stat("/var/run/reboot-required"); err == nil {
+        return true
+    }
+    return false
+}
+

distros/utils.go

@@ -0,0 +1,29 @@
+package distros
+
+import (
+    "runtime"
+    "strings"
+    "os/exec"
+)
+
+// DetectLinuxDistro returns "ubuntu", "rhel", or "unknown"
+func GetLinuxDistro() string {
+    if runtime.GOOS != "linux" {
+        return "unknown"
+    }
+
+    out, err := exec.Command("sh", "-c", "cat /etc/os-release").Output()
+    if err != nil {
+        return "unknown"
+    }
+
+    s := strings.ToLower(string(out))
+    switch {
+    case strings.Contains(s, "ubuntu"):
+        return "ubuntu"
+    case strings.Contains(s, "rhel"), strings.Contains(s, "centos"), strings.Contains(s, "fedora"):
+        return "rhel"
+    default:
+        return "unknown"
+    }
+}

main.go

@@ -5,81 +5,66 @@ import (
     "fmt"
     "log"
     "net/http"
-    "os"
-    "os/exec"
-    "strconv"
-    "strings"
     "time"
 
     "github.com/prometheus/client_golang/prometheus"
     "github.com/prometheus/client_golang/prometheus/promhttp"
+
+    "github.com/Puppet-Finland/updates-exporter/distros"
+    "github.com/Puppet-Finland/updates-exporter/distros/ubuntu"
 )
 
 const (
-    DEFAULT_INTERVAL = 86400
+    DEFAULT_INTERVAL = 3600
     DEFAULT_PORT     = 9101
 )
 var (
-    securityUpdates = prometheus.NewGauge(
-        prometheus.GaugeOpts{
-            Name: "ubuntu_pending_security_updates",
-            Help: "Number of pending security updates on Ubuntu",
-        },
-    )
-    rebootRequired = prometheus.NewGauge(
-        prometheus.GaugeOpts{
-            Name: "ubuntu_reboot_required",
-            Help: "1 if a reboot is required, 0 otherwise",
-        },
-    )
+    securityUpdates = prometheus.NewGauge(prometheus.GaugeOpts{
+        Name: "pending_security_updates",
+        Help: "Number of pending security updates",
+    })
+    rebootRequired = prometheus.NewGauge(prometheus.GaugeOpts{
+        Name: "reboot_required",
+        Help: "1 if a reboot is required, 0 otherwise",
+    })
 )
 
-func init() {
-    prometheus.MustRegister(securityUpdates)
-    prometheus.MustRegister(rebootRequired)
-}
-
-func getPendingSecurityUpdates() int {
-    cmd := exec.Command("sh", "-c", `apt-get -s dist-upgrade | grep "^Inst" | grep security | wc -l`)
-    output, err := cmd.Output()
-    if err != nil {
-        log.Printf("Error running apt-get: %v", err)
-        return -1
-    }
-    count, err := strconv.Atoi(strings.TrimSpace(string(output)))
-    if err != nil {
-        log.Printf("Error parsing security update count: %v", err)
-        return -1
+func getDistro() distros.Distro {
+    switch distros.GetLinuxDistro() {
+        case "ubuntu":
+            return ubuntu.Ubuntu{}
+        default:
+            return nil
     }
-    return count
 }
 
-func checkRebootRequired() bool {
-    if _, err := os.Stat("/var/run/reboot-required"); err == nil {
-        return true
+func updateMetrics(d distros.Distro) {
+    if d == nil {
+        return
     }
-    return false
-}
+    securityUpdates.Set(float64(d.GetSecurityUpdates()))
 
-func updateMetrics() {
-    securityUpdates.Set(float64(getPendingSecurityUpdates()))
-    if checkRebootRequired() {
+    if d.GetRebootRequired() {
         rebootRequired.Set(1)
     } else {
         rebootRequired.Set(0)
     }
 }
 
 func main() {
-    // Command-line flags
-    port := flag.Int("port", DEFAULT_PORT, "HTTP port to serve metrics")
-    interval := flag.Int("interval", DEFAULT_INTERVAL, "Metrics update interval in seconds")
+    port     := flag.Int("port", DEFAULT_PORT, "HTTP port")
+    interval := flag.Int("interval", DEFAULT_INTERVAL, "Metrics refresh interval (seconds)")
+
     flag.Parse()
 
-    // Update metrics periodically
+    prometheus.MustRegister(securityUpdates)
+    prometheus.MustRegister(rebootRequired)
+
+    distro := getDistro()
+
     go func() {
         for {
-            updateMetrics()
+            updateMetrics(distro)
             time.Sleep(time.Duration(*interval) * time.Second)
         }
     }()
@@ -88,3 +73,4 @@ func main() {
     fmt.Printf("Starting updates_exporter on :%d, updating every %d seconds\n", *port, *interval)
     log.Fatal(http.ListenAndServe(fmt.Sprintf(":%d", *port), nil))
 }
+