organizing and refactoring

agent/system/structs.go

@@ -0,0 +1,13 @@
+package system
+
+import "time"
+
+type CmdOptions struct {
+	Shell        string
+	Command      string
+	Args         []string
+	Timeout      time.Duration
+	IsScript     bool
+	IsExecutable bool
+	Detached     bool
+}

agent/system/system.go

@@ -0,0 +1,101 @@
+package system
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"os/exec"
+	"time"
+
+	"github.com/amidaware/rmmagent/agent/utils"
+	gocmd "github.com/go-cmd/cmd"
+)
+
+type CmdStatus struct {
+	Status gocmd.Status
+	Stdout string
+	Stderr string
+}
+
+func CmdV2(c *CmdOptions) CmdStatus {
+
+	ctx, cancel := context.WithTimeout(context.Background(), c.Timeout * time.Second)
+	defer cancel()
+
+	// Disable output buffering, enable streaming
+	cmdOptions := gocmd.Options{
+		Buffered:  false,
+		Streaming: true,
+	}
+
+	// have a child process that is in a different process group so that
+	// parent terminating doesn't kill child
+	if c.Detached {
+		cmdOptions.BeforeExec = []func(cmd *exec.Cmd){
+			func(cmd *exec.Cmd) {
+				cmd.SysProcAttr = SetDetached()
+			},
+		}
+	}
+
+	var envCmd *gocmd.Cmd
+	if c.IsScript {
+		envCmd = gocmd.NewCmdOptions(cmdOptions, c.Shell, c.Args...) // call script directly
+	} else if c.IsExecutable {
+		envCmd = gocmd.NewCmdOptions(cmdOptions, c.Shell, c.Command) // c.Shell: bin + c.Command: args as one string
+	} else {
+		envCmd = gocmd.NewCmdOptions(cmdOptions, c.Shell, "-c", c.Command) // /bin/bash -c 'ls -l /var/log/...'
+	}
+
+	var stdoutBuf bytes.Buffer
+	var stderrBuf bytes.Buffer
+	// Print STDOUT and STDERR lines streaming from Cmd
+	doneChan := make(chan struct{})
+	go func() {
+		defer close(doneChan)
+		// Done when both channels have been closed
+		// https://dave.cheney.net/2013/04/30/curious-channels
+		for envCmd.Stdout != nil || envCmd.Stderr != nil {
+			select {
+			case line, open := <-envCmd.Stdout:
+				if !open {
+					envCmd.Stdout = nil
+					continue
+				}
+
+				fmt.Fprintln(&stdoutBuf, line)
+
+			case line, open := <-envCmd.Stderr:
+				if !open {
+					envCmd.Stderr = nil
+					continue
+				}
+
+				fmt.Fprintln(&stderrBuf, line)
+			}
+		}
+	}()
+
+	// Run and wait for Cmd to return, discard Status
+	envCmd.Start()
+
+	go func() {
+		select {
+		case <-doneChan:
+			return
+		case <-ctx.Done():
+			pid := envCmd.Status().PID
+			KillProc(int32(pid))
+		}
+	}()
+
+	// Wait for goroutine to print everything
+	<-doneChan
+	ret := CmdStatus{
+		Status: envCmd.Status(),
+		Stdout: utils.CleanString(stdoutBuf.String()),
+		Stderr: utils.CleanString(stderrBuf.String()),
+	}
+
+	return ret
+}

agent/system/system_linux.go

@@ -0,0 +1,154 @@
+package system
+
+import (
+	"fmt"
+	"os"
+	"strings"
+	"syscall"
+	"time"
+
+	"github.com/amidaware/rmmagent/agent/utils"
+	"github.com/shirou/gopsutil/process"
+	psHost "github.com/shirou/gopsutil/v3/host"
+	"github.com/wh1te909/trmm-shared"
+)
+
+func NewCMDOpts() *CmdOptions {
+	return &CmdOptions{
+		Shell:   "/bin/bash",
+		Timeout: 30,
+	}
+}
+
+func SetDetached() *syscall.SysProcAttr {
+	return &syscall.SysProcAttr{Setpgid: true}
+}
+
+func ShowStatus(version string) {
+	fmt.Println(version)
+}
+
+func SystemRebootRequired() (bool, error) {
+	// deb
+	paths := [2]string{"/var/run/reboot-required", "/run/reboot-required"}
+	for _, p := range paths {
+		if trmm.FileExists(p) {
+			return true, nil
+		}
+	}
+	// rhel
+	bins := [2]string{"/usr/bin/needs-restarting", "/bin/needs-restarting"}
+	for _, bin := range bins {
+		if trmm.FileExists(bin) {
+			opts := NewCMDOpts()
+			// https://man7.org/linux/man-pages/man1/needs-restarting.1.html
+			// -r Only report whether a full reboot is required (exit code 1) or not (exit code 0).
+			opts.Command = fmt.Sprintf("%s -r", bin)
+			out := CmdV2(opts)
+
+			if out.Status.Error != nil {
+				continue
+			}
+
+			if out.Status.Exit == 1 {
+				return true, nil
+			}
+
+			return false, nil
+		}
+	}
+
+	return false, nil
+}
+
+func LoggedOnUser() string {
+	var ret string
+	users, err := psHost.Users()
+	if err != nil {
+		return ret
+	}
+
+	// return the first logged in user
+	for _, user := range users {
+		if user.User != "" {
+			ret = user.User
+			break
+		}
+	}
+
+	return ret
+}
+
+func OsString() string {
+	h, err := psHost.Info()
+	if err != nil {
+		return "error getting host info"
+	}
+
+	return fmt.Sprintf("%s %s %s %s", strings.Title(h.Platform), h.PlatformVersion, h.KernelArch, h.KernelVersion)
+}
+
+// KillProc kills a process and its children
+func KillProc(pid int32) error {
+	p, err := process.NewProcess(pid)
+	if err != nil {
+		return err
+	}
+
+	children, err := p.Children()
+	if err == nil {
+		for _, child := range children {
+			if err := child.Kill(); err != nil {
+				continue
+			}
+		}
+	}
+
+	if err := p.Kill(); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func RunScript(code string, shell string, args []string, timeout int) (stdout, stderr string, exitcode int, e error) {
+	code = utils.RemoveWinNewLines(code)
+	content := []byte(code)
+
+	f, err := utils.CreateTmpFile()
+	if err != nil {
+		return "", err.Error(), 85, err
+	}
+
+	defer os.Remove(f.Name())
+
+	if _, err := f.Write(content); err != nil {
+		return "", err.Error(), 85, err
+	}
+
+	if err := f.Close(); err != nil {
+		return "", err.Error(), 85, err
+	}
+
+	if err := os.Chmod(f.Name(), 0770); err != nil {
+		return "", err.Error(), 85, err
+	}
+
+	opts := NewCMDOpts()
+	opts.IsScript = true
+	opts.Shell = f.Name()
+	opts.Args = args
+	opts.Timeout = time.Duration(timeout)
+	out := CmdV2(opts)
+	retError := ""
+	if out.Status.Error != nil {
+		retError += utils.CleanString(out.Status.Error.Error())
+		retError += "\n"
+	}
+
+	if len(out.Stderr) > 0 {
+		retError += out.Stderr
+	}
+
+	return out.Stdout, retError, out.Status.Exit, nil
+}

agent/system/system_linux_test.go

@@ -0,0 +1,67 @@
+package system
+
+import (
+	"testing"
+	"github.com/amidaware/rmmagent/agent/utils"
+)
+
+func TestNewCMDOpts(t *testing.T) {
+	opts := NewCMDOpts()
+	if opts.Shell != "/bin/bash" {
+		t.Fatalf("Expected /bin/bash, got %s", opts.Shell)
+	}
+}
+
+func TestSystemRebootRequired(t *testing.T) {
+	required, err := SystemRebootRequired()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	t.Logf("System Reboot Required %t", required)
+}
+
+func TestShowStatus(t *testing.T) {
+	output := utils.CaptureOutput(func() {
+		ShowStatus("1.0.0")
+	});
+
+	if output != "1.0.0\n" {
+		t.Fatalf("Expected 1.0.0, got %s", output)
+	}
+}
+
+func TestLoggedOnUser(t *testing.T) {
+	user := LoggedOnUser()
+	if user == "" {
+		t.Fatalf("Expected a user, got empty")
+	}
+
+	t.Logf("Logged on user: %s", user)
+}
+
+func TestOsString(t *testing.T) {
+	osString := OsString()
+	if osString == "error getting host info" {
+		t.Fatalf("Unable to get OS string")
+	}
+
+	t.Logf("OS String: %s", osString)
+}
+
+func TestRunScript(t *testing.T) {
+	stdout, stderr, exitcode, err := RunScript("#!/bin/sh\ncat /etc/os-release", "/bin/sh", nil, 30)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if stderr != "" {
+		t.Fatal(stderr)
+	}
+
+	if exitcode != 0 {
+		t.Fatalf("Error: Exit Code %d", exitcode)
+	}
+
+	t.Logf("Result: %s", stdout)
+}