Minor refactor & Implement Rocky support

Author: Ricky Cousins

distros/distros.go

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

distros/rhel/rocky.go

@@ -0,0 +1,36 @@
+package rhel
+
+import (
+	"log"
+	"os/exec"
+
+	utils "github.com/Puppet-Finland/updates-exporter/distros"
+)
+
+type Rocky struct{}
+
+func (Rocky) GetSecurityUpdates() int {
+	cmd := exec.Command("sh", "-c", "dnf updateinfo list --sec-severity=Critical --sec-severity=Important --all | wc -l")
+	out, err := cmd.Output()
+	if err != nil {
+		log.Printf("Error running dnf: %v", err)
+		return -1
+	}
+	return utils.ParseUpdateCount(string(out))
+}
+
+func (Rocky) GetTotalUpdates() int {
+	cmd := exec.Command("sh", "-c", "dnf updateinfo list --all | wc -l")
+	out, err := cmd.Output()
+	if err != nil {
+		log.Printf("Error running dnf: %v", err)
+		return -1
+	}
+	return utils.ParseUpdateCount(string(out))
+}
+
+func (Rocky) GetRebootRequired() bool {
+	cmd := exec.Command("needs-restarting", "-r")
+	err := cmd.Run()
+	return err != nil
+}

distros/ubuntu/ubuntu.go

@@ -1,46 +1,40 @@
 package ubuntu
 
 import (
-    "log"
-    "os"
-    "os/exec"
-    "strconv"
-    "strings"
+	"log"
+	"os"
+	"os/exec"
+
+	utils "github.com/Puppet-Finland/updates-exporter/distros"
 )
 
 type Ubuntu struct{}
 
 var rebootFile = "/var/run/reboot-required"
 
-func parseUpdateCount(out string) int {
-    count, _ := strconv.Atoi(strings.TrimSpace(out))
-    return count
-}
-
 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 -1
-    }
-    return parseUpdateCount(string(output))
+	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
+	}
+	return utils.ParseUpdateCount(string(output))
 }
 
 func (Ubuntu) GetTotalUpdates() int {
-    cmd := exec.Command("sh", "-c", `apt-get -s dist-upgrade | grep "^Inst" | wc -l`)
-    output, err := cmd.Output()
-    if err != nil {
-        log.Printf("Error running apt-get: %v", err)
-        return -1
-    }
-    return parseUpdateCount(string(output))
+	cmd := exec.Command("sh", "-c", `apt-get -s dist-upgrade | grep "^Inst" | wc -l`)
+	output, err := cmd.Output()
+	if err != nil {
+		log.Printf("Error running apt-get: %v", err)
+		return -1
+	}
+	return utils.ParseUpdateCount(string(output))
 }
 
 func (Ubuntu) GetRebootRequired() bool {
-    if _, err := os.Stat(rebootFile); err == nil {
-        return true
-    }
-    return false
+	if _, err := os.Stat(rebootFile); err == nil {
+		return true
+	}
+	return false
 }
-

distros/ubuntu/ubuntu_test.go

@@ -1,36 +1,27 @@
 package ubuntu
 
 import (
-    "os"
-    "testing"
+	"os"
+	"testing"
 )
 
 var ubuntu = Ubuntu{}
 
-func TestParseUpdateCount(t *testing.T) {
-    output := "3\n"
-    expected := 3
-    got := parseUpdateCount(output)
-    if got != expected {
-        t.Errorf("Expected %d, got %d", expected, got)
-    }
-}
-
-func TestRebootRequired(t *testing.T) {
-    // Create a fake reboot-required file
-    tmpfile := "/tmp/reboot-required"
-    err := os.WriteFile(tmpfile, []byte("reboot"), 0644)
-    if err != nil {
-        t.Fatalf("failed to create temp file: %v", err)
-    }
+func TestUbuntuRebootRequired(t *testing.T) {
+	// Create a fake reboot-required file
+	tmpfile := "/tmp/reboot-required"
+	err := os.WriteFile(tmpfile, []byte("reboot"), 0644)
+	if err != nil {
+		t.Fatalf("failed to create temp file: %v", err)
+	}
 
-    oldFile := rebootFile
-    rebootFile = tmpfile
-    defer func() {
-        rebootFile = oldFile
-    }()
+	oldFile := rebootFile
+	rebootFile = tmpfile
+	defer func() {
+		rebootFile = oldFile
+	}()
 
-    if !ubuntu.GetRebootRequired() {
-        t.Errorf("Expected reboot required = true")
-    }
+	if !ubuntu.GetRebootRequired() {
+		t.Errorf("Expected reboot required = true")
+	}
 }

distros/utils.go

@@ -1,29 +1,41 @@
 package distros
 
 import (
-    "runtime"
-    "strings"
-    "os/exec"
+	"os"
+	"os/exec"
+	"runtime"
+	"strconv"
+	"strings"
 )
 
-// DetectLinuxDistro returns "ubuntu", "rhel", or "unknown"
+var rockyReleaseFile = "/etc/rocky-release"
+
+func ParseUpdateCount(out string) int {
+	count, _ := strconv.Atoi(strings.TrimSpace(out))
+	return count
+}
+
 func GetLinuxDistro() string {
-    if runtime.GOOS != "linux" {
-        return "unknown"
-    }
+	if runtime.GOOS != "linux" {
+		return "unknown"
+	}
+
+	if _, err := os.Stat(rockyReleaseFile); err == nil {
+		return "rocky"
+	}
 
-    out, err := exec.Command("sh", "-c", "cat /etc/os-release").Output()
-    if err != nil {
-        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"
-    }
+	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"
+	}
 }

distros/utils_test.go

@@ -0,0 +1,14 @@
+package distros
+
+import (
+	"testing"
+)
+
+func TestParseUpdateCount(t *testing.T) {
+	output := "3\n"
+	expected := 3
+	got := ParseUpdateCount(output)
+	if got != expected {
+		t.Errorf("Expected %d, got %d", expected, got)
+	}
+}

main.go

@@ -1,86 +1,89 @@
 package main
 
 import (
-    "flag"
-    "fmt"
-    "log"
-    "net/http"
-    "time"
+	"flag"
+	"fmt"
+	"log"
+	"net/http"
+	"time"
 
-    "github.com/prometheus/client_golang/prometheus"
-    "github.com/prometheus/client_golang/prometheus/promhttp"
+	"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"
+	"github.com/Puppet-Finland/updates-exporter/distros"
+	"github.com/Puppet-Finland/updates-exporter/distros/rhel"
+	"github.com/Puppet-Finland/updates-exporter/distros/ubuntu"
 )
 
 const (
-    DEFAULT_INTERVAL = 3600
-    DEFAULT_PORT     = 9101
+	DEFAULT_INTERVAL = 3600
+	DEFAULT_PORT     = 9101
 )
+
 var (
-    securityUpdates = prometheus.NewGauge(prometheus.GaugeOpts{
-        Name: "pending_security_updates",
-        Help: "Number of pending security updates",
-    })
-    totalUpdates = prometheus.NewGauge(prometheus.GaugeOpts{
-        Name: "pending_updates",
-        Help: "Total number of pending updates",
-    })
-    rebootRequired = prometheus.NewGauge(prometheus.GaugeOpts{
-        Name: "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",
+	})
+	totalUpdates = prometheus.NewGauge(prometheus.GaugeOpts{
+		Name: "pending_updates",
+		Help: "Total number of pending updates",
+	})
+	rebootRequired = prometheus.NewGauge(prometheus.GaugeOpts{
+		Name: "reboot_required",
+		Help: "1 if a reboot is required, 0 otherwise",
+	})
 )
 
 func getDistro() distros.Distro {
-    switch distros.GetLinuxDistro() {
-        case "ubuntu":
-            return ubuntu.Ubuntu{}
-        default:
-            return nil
-    }
+	switch distros.GetLinuxDistro() {
+	case "ubuntu":
+		return ubuntu.Ubuntu{}
+	case "rocky":
+		return rhel.Rocky{}
+	default:
+		return nil
+	}
 }
 
 func updateMetrics(d distros.Distro) {
-    if d == nil {
-        return
-    }
-    securityUpdates.Set(float64(d.GetSecurityUpdates()))
-    totalUpdates.Set(float64(d.GetTotalUpdates()))
+	if d == nil {
+		return
+	}
+	securityUpdates.Set(float64(d.GetSecurityUpdates()))
+	totalUpdates.Set(float64(d.GetTotalUpdates()))
 
-    if d.GetRebootRequired() {
-        rebootRequired.Set(1)
-    } else {
-        rebootRequired.Set(0)
-    }
+	if d.GetRebootRequired() {
+		rebootRequired.Set(1)
+	} else {
+		rebootRequired.Set(0)
+	}
 
 }
 
 func main() {
-    port     := flag.Int("port", DEFAULT_PORT, "HTTP port")
-    interval := flag.Int("interval", DEFAULT_INTERVAL, "Metrics refresh interval (seconds)")
+	port := flag.Int("port", DEFAULT_PORT, "HTTP port")
+	interval := flag.Int("interval", DEFAULT_INTERVAL, "Metrics refresh interval (seconds)")
 
-    flag.Parse()
+	flag.Parse()
 
-    prometheus.MustRegister(securityUpdates)
-    prometheus.MustRegister(totalUpdates)
-    prometheus.MustRegister(rebootRequired)
+	prometheus.MustRegister(securityUpdates)
+	prometheus.MustRegister(totalUpdates)
+	prometheus.MustRegister(rebootRequired)
 
-    distro := getDistro()
-    if distro == nil {
-        log.Fatal("Error: Distro not detected")
-    }
+	distro := getDistro()
+	if distro == nil {
+		log.Fatal("Error: Distro not detected")
+	}
 
-    go func() {
-        for {
-            updateMetrics(distro)
-            time.Sleep(time.Duration(*interval) * time.Second)
-        }
-    }()
+	go func() {
+		for {
+			updateMetrics(distro)
+			time.Sleep(time.Duration(*interval) * time.Second)
+		}
+	}()
 
-    http.Handle("/metrics", promhttp.Handler())
-    fmt.Printf("Starting updates_exporter on :%d, updating every %d seconds\n", *port, *interval)
-    log.Fatal(http.ListenAndServe(fmt.Sprintf(":%d", *port), nil))
+	http.Handle("/metrics", promhttp.Handler())
+	fmt.Printf("Starting updates_exporter on :%d, updating every %d seconds\n", *port, *interval)
+	log.Fatal(http.ListenAndServe(fmt.Sprintf(":%d", *port), nil))
 }
-