diff --git a/agent/system/system_windows.go b/agent/system/system_windows.go index a50b371..03bc110 100644 --- a/agent/system/system_windows.go +++ b/agent/system/system_windows.go @@ -56,18 +56,15 @@ func RunScript(code string, shell string, args []string, timeout int) (stdout, s tmpfn, err := ioutil.TempFile(dir, ext) if err != nil { - //a.Logger.Errorln(err) return "", err.Error(), 85, err } defer os.Remove(tmpfn.Name()) if _, err := tmpfn.Write(content); err != nil { - //a.Logger.Errorln(err) return "", err.Error(), 85, err } if err := tmpfn.Close(); err != nil { - //a.Logger.Errorln(err) return "", err.Error(), 85, err } @@ -93,7 +90,6 @@ func RunScript(code string, shell string, args []string, timeout int) (stdout, s cmd.Stdout = &outb cmd.Stderr = &errb if cmdErr := cmd.Start(); cmdErr != nil { - //a.Logger.Debugln(cmdErr) return "", cmdErr.Error(), 65, cmdErr } @@ -474,7 +470,7 @@ func OsString() string { return osFullName } -func AddDefenderExlusions() { +func AddDefenderExlusions() error { code := ` Add-MpPreference -ExclusionPath 'C:\Program Files\TacticalAgent\*' Add-MpPreference -ExclusionPath 'C:\Windows\Temp\winagent-v*.exe' @@ -483,6 +479,8 @@ Add-MpPreference -ExclusionPath 'C:\Program Files\Mesh Agent\*' ` _, _, _, err := RunScript(code, "powershell", []string{}, 20) if err != nil { - //a.Logger.Debugln(err) + return err } + + return nil } diff --git a/agent/system/system_windows_test.go b/agent/system/system_windows_test.go new file mode 100644 index 0000000..4475b07 --- /dev/null +++ b/agent/system/system_windows_test.go @@ -0,0 +1,67 @@ +package system_test + +import ( + "errors" + "strings" + "testing" + + "github.com/amidaware/rmmagent/agent/system" +) + +func TestRunScript(t *testing.T) { + testTable := []struct { + name string + code string + shell string + args []string + timeout int + expectedStdout string + expectedStderr string + expectedExitCode int + expectedError error + }{ + { + name: "Run Script", + code: "Test-Path -Path C:\\Windows", + shell: "powershell", + args: []string{}, + timeout: 30, + expectedStdout: "True\r\n", + expectedStderr: "", + expectedExitCode: 0, + expectedError: nil, + }, + { + name: "Run Error Script", + code: "Get-ThisError", + shell: "powershell", + args: []string{}, + timeout: 30, + expectedStdout: "", + expectedStderr: "The term 'Get-ThisError' is not recognized as the name of a cmdlet", + expectedExitCode: 0, + expectedError: nil, + }, + } + + for _, tt := range testTable { + t.Run(tt.name, func(t *testing.T) { + stdout, stderr, exitcode, err := system.RunScript(tt.code, tt.shell, tt.args, tt.timeout) + if stdout != tt.expectedStdout { + t.Errorf("expected stdout %s, got %s", tt.expectedStdout, stdout) + } + + if !strings.Contains(stderr, tt.expectedStderr) { + t.Errorf("expected stderr to contain %s, got %s", tt.expectedStderr, stderr) + } + + if exitcode != tt.expectedExitCode { + t.Errorf("expected exitcode %d, got %d", tt.expectedExitCode, exitcode) + } + + if !errors.Is(err, tt.expectedError) { + t.Errorf("expected error (%v), got (%v)", tt.expectedError, err) + } + }) + } +} diff --git a/agent/utils/utils_test.go b/agent/utils/utils_test.go index 2604302..e8e33f1 100644 --- a/agent/utils/utils_test.go +++ b/agent/utils/utils_test.go @@ -1,31 +1,40 @@ -package utils +package utils_test import ( "testing" + + "github.com/amidaware/rmmagent/agent/utils" ) func TestByteCountSI(t *testing.T) { - var bytes uint64 = 1048576 - mb := ByteCountSI(bytes) - if mb != "1.0 MB" { - t.Errorf("Expected 1.0 MB, got %s", mb) - } -} - -func TestRemoveWinNewLines(t *testing.T) { - result := RemoveWinNewLines("test\r\n") - if result != "test\n" { - t.Fatalf("Expected testing\\n, got %s", result) + testTable := []struct { + name string + expected string + bytes uint64 + }{ + { + name: "Bytes to Kilobytes", + expected: "1.0 kB", + bytes: 1024, + }, + { + name: "Bytes to Megabytes", + expected: "1.0 MB", + bytes: 1048576, + }, + { + name: "Bytes to Gigabytes", + expected: "1.0 GB", + bytes: 1073741824, + }, } - t.Logf("Result: %s", result) -} - -func TestStripAll(t *testing.T) { - result := StripAll(" test\r\n ") - if result != "test" { - t.Fatalf("Expecte test, got %s", result) + for _, tt := range testTable { + t.Run(tt.name, func(t *testing.T) { + result := utils.ByteCountSI(tt.bytes) + if result != tt.expected { + t.Errorf("expected %s, got %s", tt.expected, result) + } + }) } - - t.Log("Test result expected") } diff --git a/agent/wmi/wmi_windows.go b/agent/wmi/wmi_windows.go index 38f4c50..20453f9 100644 --- a/agent/wmi/wmi_windows.go +++ b/agent/wmi/wmi_windows.go @@ -7,72 +7,73 @@ import ( "github.com/yusufpapurcu/wmi" ) -func GetWMIInfo() map[string]interface{} { +func GetWMIInfo() (map[string]interface{}, []error) { wmiInfo := make(map[string]interface{}) + errs := []error{} compSysProd, err := GetWin32_ComputerSystemProduct() if err != nil { - //a.Logger.Debugln(err) + errs = append(errs, err) } compSys, err := GetWin32_ComputerSystem() if err != nil { - //a.Logger.Debugln(err) + errs = append(errs, err) } netAdaptConfig, err := GetWin32_NetworkAdapterConfiguration() if err != nil { - //a.Logger.Debugln(err) + errs = append(errs, err) } physMem, err := GetWin32_PhysicalMemory() if err != nil { - //a.Logger.Debugln(err) + errs = append(errs, err) } winOS, err := GetWin32_OperatingSystem() if err != nil { - //a.Logger.Debugln(err) + errs = append(errs, err) } baseBoard, err := GetWin32_BaseBoard() if err != nil { - //a.Logger.Debugln(err) + errs = append(errs, err) } bios, err := GetWin32_BIOS() if err != nil { - //a.Logger.Debugln(err) + errs = append(errs, err) } disk, err := GetWin32_DiskDrive() if err != nil { - //a.Logger.Debugln(err) + errs = append(errs, err) } netAdapt, err := GetWin32_NetworkAdapter() if err != nil { - //a.Logger.Debugln(err) + errs = append(errs, err) } desktopMon, err := GetWin32_DesktopMonitor() if err != nil { - //a.Logger.Debugln(err) + errs = append(errs, err) } cpu, err := GetWin32_Processor() if err != nil { - //a.Logger.Debugln(err) + errs = append(errs, err) } usb, err := GetWin32_USBController() if err != nil { - //a.Logger.Debugln(err) + errs = append(errs, err) } graphics, err := GetWin32_VideoController() if err != nil { - //a.Logger.Debugln(err) + errs = append(errs, err) } wmiInfo["comp_sys_prod"] = compSysProd @@ -89,7 +90,7 @@ func GetWMIInfo() map[string]interface{} { wmiInfo["usb"] = usb wmiInfo["graphics"] = graphics - return wmiInfo + return wmiInfo, errs } func GetWin32_ComputerSystemProduct() ([]interface{}, error) { @@ -113,6 +114,7 @@ func GetWin32_ComputerSystemProduct() ([]interface{}, error) { if err := json.Unmarshal(b, &un); err != nil { return ret, err } + tmp = append(tmp, un) ret = append(ret, tmp) } diff --git a/agent/wmi/wmi_windows_test.go b/agent/wmi/wmi_windows_test.go new file mode 100644 index 0000000..9844e2c --- /dev/null +++ b/agent/wmi/wmi_windows_test.go @@ -0,0 +1,37 @@ +package wmi_test + +import ( + "reflect" + "testing" + + "github.com/amidaware/rmmagent/agent/wmi" +) + +func TestGetWMIInfo(t *testing.T) { + testTable := []struct { + name string + expected map[string]interface{} + atLeast int + expectedErrors []error + }{ + { + name: "Get WMI Data", + expected: make(map[string]interface{}), + atLeast: 1, + expectedErrors: []error{}, + }, + } + + for _, tt := range testTable { + t.Run(tt.name, func(t *testing.T) { + result, errs := wmi.GetWMIInfo() + if len(result) < tt.atLeast { + t.Errorf("expected at least %d, got %d", tt.atLeast, len(result)) + } + + if !reflect.DeepEqual(errs, tt.expectedErrors) { + t.Errorf("expected errors (%v), got (%v)", tt.expectedErrors, errs) + } + }) + } +}