From fa1543212116cc721c932027a8dc6def25e32033 Mon Sep 17 00:00:00 2001
From: Tai Groot <tai@taigrr.com>
Date: Thu, 8 Aug 2024 15:37:04 -0700
Subject: [PATCH] add ability to list all units

---
 helpers.go      | 28 ++++++++++++++++++++++++++++
 helpers_test.go | 48 ++++++++++++++++++++++++++++++++++++++++++++++++
 structs.go      |  8 ++++++++
 3 files changed, 84 insertions(+)

diff --git a/helpers.go b/helpers.go
index bb7aec9..01cd05c 100644
--- a/helpers.go
+++ b/helpers.go
@@ -55,6 +55,34 @@ func GetPID(ctx context.Context, unit string, opts Options) (int, error) {
 	return strconv.Atoi(value)
 }
 
+func GetUnits(ctx context.Context, opts Options) ([]Unit, error) {
+	args := []string{"list-units", "--all", "--no-legend", "--full", "--no-pager"}
+	if opts.UserMode {
+		args = append(args, "--user")
+	}
+	stdout, stderr, _, err := execute(ctx, args)
+	if err != nil {
+		return []Unit{}, errors.Join(err, filterErr(stderr))
+	}
+	lines := strings.Split(stdout, "\n")
+	units := []Unit{}
+	for _, line := range lines {
+		entry := strings.Fields(line)
+		if len(entry) < 4 {
+			continue
+		}
+		unit := Unit{
+			Name:        entry[0],
+			Load:        entry[1],
+			Active:      entry[2],
+			Sub:         entry[3],
+			Description: strings.Join(entry[4:], " "),
+		}
+		units = append(units, unit)
+	}
+	return units, nil
+}
+
 func GetMaskedUnits(ctx context.Context, opts Options) ([]string, error) {
 	args := []string{"list-unit-files", "--state=masked"}
 	if opts.UserMode {
diff --git a/helpers_test.go b/helpers_test.go
index d4b91e1..9b91c9d 100644
--- a/helpers_test.go
+++ b/helpers_test.go
@@ -238,6 +238,54 @@ func TestGetMemoryUsage(t *testing.T) {
 	})
 }
 
+func TestGetUnits(t *testing.T) {
+	type testCase struct {
+		err       error
+		runAsUser bool
+		opts      Options
+	}
+	testCases := []testCase{{
+		// Run these tests only as a user
+		runAsUser: true,
+		opts:      Options{UserMode: true},
+		err:       nil,
+	}}
+	for _, tc := range testCases {
+		t.Run(fmt.Sprintf("as %s", userString), func(t *testing.T) {
+			if (userString == "root" || userString == "system") && tc.runAsUser {
+				t.Skip("skipping user test while running as superuser")
+			} else if (userString != "root" && userString != "system") && !tc.runAsUser {
+				t.Skip("skipping superuser test while running as user")
+			}
+			ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
+			defer cancel()
+			units, err := GetUnits(ctx, tc.opts)
+			if !errors.Is(err, tc.err) {
+				t.Errorf("error is %v, but should have been %v", err, tc.err)
+			}
+			if len(units) == 0 {
+				t.Errorf("Expected at least one unit, but got none")
+			}
+			unit := units[0]
+			if unit.Name == "" {
+				t.Errorf("Expected unit name to be non-empty, but got empty")
+			}
+			if unit.Load == "" {
+				t.Errorf("Expected unit load state to be non-empty, but got empty")
+			}
+			if unit.Active == "" {
+				t.Errorf("Expected unit active state to be non-empty, but got empty")
+			}
+			if unit.Sub == "" {
+				t.Errorf("Expected unit sub state to be non-empty, but got empty")
+			}
+			if unit.Description == "" {
+				t.Errorf("Expected unit description to be non-empty, but got empty")
+			}
+		})
+	}
+}
+
 func TestGetPID(t *testing.T) {
 	type testCase struct {
 		unit      string
diff --git a/structs.go b/structs.go
index 2a48a73..0a0458a 100644
--- a/structs.go
+++ b/structs.go
@@ -3,3 +3,11 @@ package systemctl
 type Options struct {
 	UserMode bool
 }
+
+type Unit struct {
+	Name        string
+	Load        string
+	Active      string
+	Sub         string
+	Description string
+}