Refactor runAsUser into tablehelper package

kolide · Mar 14, 2024 · 1b9bcda · 1b9bcda
1 parent 0013ad0
commit 1b9bcda
Show file tree

Hide file tree

Showing 3 changed files with 77 additions and 157 deletions.
diff --git a/ee/tables/homebrew/upgradeable.go b/ee/tables/homebrew/upgradeable.go
@@ -6,13 +6,8 @@ package brew_upgradeable
 import (
 	"context"
 	"fmt"
-	"io"
 	"log/slog"
-	"os/exec"
-	"os/user"
-	"strconv"
 	"strings"
-	"syscall"
 
 	"github.com/kolide/launcher/ee/allowedcmd"
 	"github.com/kolide/launcher/ee/dataflatten"
@@ -25,7 +20,6 @@ const allowedCharacters = "0123456789"
 
 type Table struct {
 	slogger *slog.Logger
-	execCC  allowedcmd.AllowedCommand
 }
 
 func TablePlugin(slogger *slog.Logger) *table.Plugin {
@@ -35,7 +29,6 @@ func TablePlugin(slogger *slog.Logger) *table.Plugin {
 
 	t := &Table{
 		slogger: slogger.With("table", "kolide_brew_upgradeable"),
-		execCC:  allowedcmd.Brew,
 	}
 
 	return table.NewPlugin("kolide_brew_upgradeable", columns, t.generate)
@@ -49,9 +42,14 @@ func (t *Table) generate(ctx context.Context, queryContext table.QueryContext) (
 		return results, fmt.Errorf("kolide_brew_upgradeable requires at least one user id to be specified")
 	}
 
+	cmd, err := allowedcmd.Brew(ctx, "outdated", "--json")
+	if err != nil {
+		return results, fmt.Errorf("creating brew outdated command: %w", err)
+	}
+
 	for _, uid := range uids {
 		for _, dataQuery := range tablehelpers.GetConstraints(queryContext, "query", tablehelpers.WithDefaults("*")) {
-			output, err := t.getBrewOutdated(ctx, uid)
+			output, err := tablehelpers.runCmdAsUser(cmd, uid)
 			if err != nil {
 				t.slogger.Log(ctx, slog.LevelInfo, "failure querying user brew installed packages", "err", err, "target_uid", uid)
 				continue
@@ -78,69 +76,3 @@ func (t *Table) generate(ctx context.Context, queryContext table.QueryContext) (
 
 	return results, nil
 }
-
-func (t *Table) getBrewOutdated(ctx context.Context, uid string) ([]byte, error) {
-	cmd, err := t.execCC(ctx, "outdated", "--json")
-	if err != nil {
-		return nil, fmt.Errorf("creating brew outdated command: %w", err)
-	}
-
-	stdout, err := cmd.StdoutPipe()
-	if err != nil {
-		return nil, fmt.Errorf("assigning StdoutPipe for brew outdated command: %w", err)
-	}
-
-	if err := runAsUser(ctx, uid, cmd); err != nil {
-		return nil, fmt.Errorf("runAsUser brew outdated command as user %s: %w", uid, err)
-	}
-
-	data, err := io.ReadAll(stdout)
-	if err != nil {
-		return nil, fmt.Errorf("ReadAll brew outdated stdout: %w", err)
-	}
-
-	if err := cmd.Wait(); err != nil {
-		return nil, fmt.Errorf("deallocation of brew outdated command as user %s: %w", uid, err)
-	}
-
-	return data, nil
-}
-
-func runAsUser(ctx context.Context, uid string, cmd *exec.Cmd) error {
-	currentUser, err := user.Current()
-	if err != nil {
-		return fmt.Errorf("getting current user: %w", err)
-	}
-
-	runningUser, err := user.LookupId(uid)
-	if err != nil {
-		return fmt.Errorf("looking up user with uid %s: %w", uid, err)
-	}
-
-	if currentUser.Uid != "0" {
-		if currentUser.Uid != runningUser.Uid {
-			return fmt.Errorf("current user %s is not root and can't start process for other user %s", currentUser.Uid, uid)
-		}
-
-		return cmd.Start()
-	}
-
-	runningUserUid, err := strconv.ParseUint(runningUser.Uid, 10, 32)
-	if err != nil {
-		return fmt.Errorf("converting uid %s to int: %w", runningUser.Uid, err)
-	}
-
-	runningUserGid, err := strconv.ParseUint(runningUser.Gid, 10, 32)
-	if err != nil {
-		return fmt.Errorf("converting gid %s to int: %w", runningUser.Gid, err)
-	}
-
-	cmd.SysProcAttr = &syscall.SysProcAttr{
-		Credential: &syscall.Credential{
-			Uid: uint32(runningUserUid),
-			Gid: uint32(runningUserGid),
-		},
-	}
-
-	return cmd.Start()
-}
diff --git a/ee/tables/nix_env/upgradeable/upgradeable.go b/ee/tables/nix_env/upgradeable/upgradeable.go
@@ -6,13 +6,8 @@ package nix_env_upgradeable
 import (
 	"context"
 	"fmt"
-	"io"
 	"log/slog"
-	"os/exec"
-	"os/user"
-	"strconv"
 	"strings"
-	"syscall"
 
 	"github.com/kolide/launcher/ee/allowedcmd"
 	"github.com/kolide/launcher/ee/dataflatten"
@@ -25,7 +20,6 @@ const allowedCharacters = "0123456789"
 
 type Table struct {
 	slogger *slog.Logger
-	execCC  allowedcmd.AllowedCommand
 }
 
 func TablePlugin(slogger *slog.Logger) *table.Plugin {
@@ -35,7 +29,6 @@ func TablePlugin(slogger *slog.Logger) *table.Plugin {
 
 	t := &Table{
 		slogger: slogger.With("table", "kolide_nix_upgradeable"),
-		execCC:  allowedcmd.NixEnv,
 	}
 
 	return table.NewPlugin("kolide_nix_upgradeable", columns, t.generate)
@@ -49,15 +42,16 @@ func (t *Table) generate(ctx context.Context, queryContext table.QueryContext) (
 		return results, fmt.Errorf("kolide_nix_upgradeable requires at least one user id to be specified")
 	}
 
+	cmd, err := allowedcmd.NixEnv(ctx, "--query", "--installed", "-c", "--xml")
+	if err != nil {
+		return results, fmt.Errorf("creating nix-env package query command: %w", err)
+	}
+
 	for _, uid := range uids {
 		for _, dataQuery := range tablehelpers.GetConstraints(queryContext, "query", tablehelpers.WithDefaults("*")) {
-			output, err := t.getUserPackages(ctx, uid)
+			output, err := tablehelpers.runCmdAsUser(cmd, uid)
 			if err != nil {
-				t.slogger.Log(ctx, slog.LevelInfo,
-					"failure querying user installed packages",
-					"err", err,
-					"target_uid", uid,
-				)
+				t.slogger.Log(ctx, slog.LevelInfo, "failure querying user installed packages", "err", err, "target_uid", uid)
 				continue
 			}
 
@@ -68,10 +62,7 @@ func (t *Table) generate(ctx context.Context, queryContext table.QueryContext) (
 
 			flattened, err := dataflatten.Xml(output, flattenOpts...)
 			if err != nil {
-				t.slogger.Log(ctx, slog.LevelInfo,
-					"failure flattening output",
-					"err", err,
-				)
+				t.slogger.Log(ctx, slog.LevelInfo, "failure flattening output", "err", err)
 				continue
 			}
 
@@ -85,69 +76,3 @@ func (t *Table) generate(ctx context.Context, queryContext table.QueryContext) (
 
 	return results, nil
 }
-
-func (t *Table) getUserPackages(ctx context.Context, uid string) ([]byte, error) {
-	cmd, err := t.execCC(ctx, "--query", "--installed", "-c", "--xml")
-	if err != nil {
-		return nil, fmt.Errorf("creating nix-env command: %w", err)
-	}
-
-	stdout, err := cmd.StdoutPipe()
-	if err != nil {
-		return nil, fmt.Errorf("assigning StdoutPipe for nix-env command: %w", err)
-	}
-
-	if err := runAsUser(ctx, uid, cmd); err != nil {
-		return nil, fmt.Errorf("runAsUser nix-env command as user %s: %w", uid, err)
-	}
-
-	data, err := io.ReadAll(stdout)
-	if err != nil {
-		return nil, fmt.Errorf("ReadAll nix-env stdout: %w", err)
-	}
-
-	if err := cmd.Wait(); err != nil {
-		return nil, fmt.Errorf("deallocation of nix-env command as user %s: %w", uid, err)
-	}
-
-	return data, nil
-}
-
-func runAsUser(ctx context.Context, uid string, cmd *exec.Cmd) error {
-	currentUser, err := user.Current()
-	if err != nil {
-		return fmt.Errorf("getting current user: %w", err)
-	}
-
-	runningUser, err := user.LookupId(uid)
-	if err != nil {
-		return fmt.Errorf("looking up user with uid %s: %w", uid, err)
-	}
-
-	if currentUser.Uid != "0" {
-		if currentUser.Uid != runningUser.Uid {
-			return fmt.Errorf("current user %s is not root and can't start process for other user %s", currentUser.Uid, uid)
-		}
-
-		return cmd.Start()
-	}
-
-	runningUserUid, err := strconv.ParseUint(runningUser.Uid, 10, 32)
-	if err != nil {
-		return fmt.Errorf("converting uid %s to int: %w", runningUser.Uid, err)
-	}
-
-	runningUserGid, err := strconv.ParseUint(runningUser.Gid, 10, 32)
-	if err != nil {
-		return fmt.Errorf("converting gid %s to int: %w", runningUser.Gid, err)
-	}
-
-	cmd.SysProcAttr = &syscall.SysProcAttr{
-		Credential: &syscall.Credential{
-			Uid: uint32(runningUserUid),
-			Gid: uint32(runningUserGid),
-		},
-	}
-
-	return cmd.Start()
-}
diff --git a/ee/tables/tablehelpers/run_as_user.go b/ee/tables/tablehelpers/run_as_user.go
@@ -0,0 +1,63 @@
+package tablehelpers
+
+import (
+	"fmt"
+	"io"
+	"os/exec"
+	"os/user"
+	"strconv"
+	"syscall"
+)
+
+func runCmdAsUser(cmd *exec.Cmd, uid string) ([]byte, error) {
+	currentUser, err := user.Current()
+	if err != nil {
+		return nil, fmt.Errorf("getting current user: %w", err)
+	}
+
+	runningUser, err := user.LookupId(uid)
+	if err != nil {
+		return nil, fmt.Errorf("looking up user with uid %s: %w", uid, err)
+	}
+
+	if currentUser.Uid != "0" && currentUser.Uid != runningUser.Uid {
+		return nil, fmt.Errorf("current user %s is not root and can't start process for other user %s", currentUser.Uid, uid)
+	}
+
+	runningUserUid, err := strconv.ParseUint(runningUser.Uid, 10, 32)
+	if err != nil {
+		return nil, fmt.Errorf("converting uid %s to int: %w", runningUser.Uid, err)
+	}
+
+	runningUserGid, err := strconv.ParseUint(runningUser.Gid, 10, 32)
+	if err != nil {
+		return nil, fmt.Errorf("converting gid %s to int: %w", runningUser.Gid, err)
+	}
+
+	cmd.SysProcAttr = &syscall.SysProcAttr{
+		Credential: &syscall.Credential{
+			Uid: uint32(runningUserUid),
+			Gid: uint32(runningUserGid),
+		},
+	}
+
+	stdout, err := cmd.StdoutPipe()
+	if err != nil {
+		return nil, fmt.Errorf("assigning command StdoutPipe: %w", err)
+	}
+
+	if err := cmd.Start(); err != nil {
+		return nil, fmt.Errorf("starting command: %w", err)
+	}
+
+	data, err := io.ReadAll(stdout)
+	if err != nil {
+		return nil, fmt.Errorf("ReadAll command stdout: %w", err)
+	}
+
+	if err := cmd.Wait(); err != nil {
+		return nil, fmt.Errorf("deallocate command: %w", err)
+	}
+
+	return data, nil
+}