Merge pull request #11 from redanthrax/RefactorDebug

Refactor debug

.gitignore

@@ -10,3 +10,5 @@
 *.bmp
 build/Output
 tacticalagent-v*
+tacticalagent
+agent/testargs.json

.vscode/launch.json

@@ -0,0 +1,18 @@
+{
+    // Use IntelliSense to learn about possible attributes.
+    // Hover to view descriptions of existing attributes.
+    // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387
+    "version": "0.2.0",
+    "configurations": [
+        {
+            "name": "Launch file",
+            "type": "go",
+            "request": "launch",
+            "mode": "debug",
+            "env": {},
+            "args": ["-m", "svc", "-log", "DEBUG", "-logto", "stdout"],
+            "buildFlags": "-tags=DEBUG",
+            "program": "${workspaceRoot}"
+        }
+    ]
+}

README.md

@@ -6,4 +6,16 @@ https://github.com/amidaware/tacticalrmm
 env CGO_ENABLED=0 GOOS=<GOOS> GOARCH=<GOARCH> go build -ldflags "-s -w"
 ```
 
+### tests
+Navigate to agent directory
+```
+go test -vet=off
+```
 
+Add to settings.json
+```
+"go.testFlags": [
+    "-vet=off"
+],
+"go.testTags": "TEST"
+```

agent/agent.go

@@ -52,7 +52,7 @@ type Agent struct {
 	EXE           string
 	SystemDrive   string
 	MeshInstaller string
-	MeshSystemEXE string
+	MeshSystemBin string
 	MeshSVC       string
 	PyBin         string
 	Headers       map[string]string
@@ -114,15 +114,15 @@ func New(logger *logrus.Logger, version string) *Agent {
 		restyC.SetRootCertificate(ac.Cert)
 	}
 
-	var MeshSysExe string
+	var MeshSysBin string
 	if len(ac.CustomMeshDir) > 0 {
-		MeshSysExe = filepath.Join(ac.CustomMeshDir, "MeshAgent.exe")
+		MeshSysBin = filepath.Join(ac.CustomMeshDir, "MeshAgent.exe")
 	} else {
-		MeshSysExe = filepath.Join(os.Getenv("ProgramFiles"), "Mesh Agent", "MeshAgent.exe")
+		MeshSysBin = filepath.Join(os.Getenv("ProgramFiles"), "Mesh Agent", "MeshAgent.exe")
 	}
 
 	if runtime.GOOS == "linux" {
-		MeshSysExe = "/opt/tacticalmesh/meshagent"
+		MeshSysBin = "/opt/tacticalmesh/meshagent"
 	}
 
 	svcConf := &service.Config{
@@ -152,7 +152,7 @@ func New(logger *logrus.Logger, version string) *Agent {
 		EXE:           exe,
 		SystemDrive:   sd,
 		MeshInstaller: "meshagent.exe",
-		MeshSystemEXE: MeshSysExe,
+		MeshSystemBin: MeshSysBin,
 		MeshSVC:       meshSvcName,
 		PyBin:         pybin,
 		Headers:       headers,

agent/agent_linux.go

@@ -132,6 +132,7 @@ func NewAgentConfig() *rmm.AgentConfig {
 	viper.SetConfigType("json")
 	viper.AddConfigPath("/etc/")
 	viper.AddConfigPath(".")
+
 	err := viper.ReadInConfig()
 
 	if err != nil {
@@ -278,14 +279,14 @@ func (a *Agent) NixMeshNodeID() string {
 	meshSuccess := false
 	a.Logger.Debugln("Getting mesh node id")
 
-	if !trmm.FileExists(a.MeshSystemEXE) {
+	if !trmm.FileExists(a.MeshSystemBin) {
-		a.Logger.Debugln(a.MeshSystemEXE, "does not exist. Skipping.")
+		a.Logger.Debugln(a.MeshSystemBin, "does not exist. Skipping.")
 		return ""
 	}
 
 	opts := a.NewCMDOpts()
 	opts.IsExecutable = true
-	opts.Shell = a.MeshSystemEXE
+	opts.Shell = a.MeshSystemBin
 	opts.Command = "-nodeid"
 
 	for !meshSuccess {

agent/agent_linux_test.go

@@ -0,0 +1,45 @@
+package agent
+
+import (
+	"bytes"
+	"io"
+	"os"
+	"testing"
+)
+
+func captureOutput(f func()) string {
+	old := os.Stdout
+	r, w, _ := os.Pipe()
+	os.Stdout = w
+	f()
+	w.Close()
+	os.Stdout = old
+
+	var buf bytes.Buffer
+	io.Copy(&buf, r)
+	return buf.String()
+}
+
+func TestShowStatus(t *testing.T) {
+	var (
+		version = "2.0.4"
+	)
+
+	output := captureOutput(func() {
+		ShowStatus(version)
+	})
+
+	if output != (version + "\n") {
+		t.Errorf("ShowStatus output not equal to version defined.")
+	}
+}
+
+func TestOsString(t *testing.T) {
+	a := New(lg, version)
+	osString := a.osString()
+	if osString == "" {
+		t.Errorf("Could not get OS String")
+	} else {
+		t.Logf("Got OS String: %s", osString)
+	}
+}

agent/agent_test.go

@@ -0,0 +1,20 @@
+package agent
+
+import (
+	"testing"
+	"github.com/sirupsen/logrus"
+)
+
+var (
+	version = "2.0.4"
+	lg     = logrus.New()
+)
+
+func TestAgentId(t *testing.T) {
+	a := New(lg, version)
+	if a.AgentID == "" {
+		t.Error("AgentID not set")
+	} else {
+		t.Logf("AgentID: %s", a.AgentID)
+	}
+}

agent/agent_windows.go

@@ -48,6 +48,11 @@ var (
 
 func NewAgentConfig() *rmm.AgentConfig {
 	k, err := registry.OpenKey(registry.LOCAL_MACHINE, `SOFTWARE\TacticalRMM`, registry.ALL_ACCESS)
+	if shared.TEST {
+		err = nil
+		k, _, err := registry.OpenKey(registry.CURRENT_USER, `SOFTWARE\TacticalRMM`, registry.ALL_ACCESS)
+	}
+
 	if err != nil {
 		return &rmm.AgentConfig{}
 	}
@@ -797,7 +802,7 @@ func (a *Agent) RecoverMesh() {
 }
 
 func (a *Agent) getMeshNodeID() (string, error) {
-	out, err := CMD(a.MeshSystemEXE, []string{"-nodeid"}, 10, false)
+	out, err := CMD(a.MeshSystemBin, []string{"-nodeid"}, 10, false)
 	if err != nil {
 		a.Logger.Debugln(err)
 		return "", err
@@ -835,6 +840,11 @@ func (a *Agent) InstallService() error {
 
 	// skip on first call of inno setup if this is a new install
 	_, err := registry.OpenKey(registry.LOCAL_MACHINE, `SOFTWARE\TacticalRMM`, registry.ALL_ACCESS)
+	if shared.TEST {
+		err = nil
+		k, _, err := registry.OpenKey(registry.CURRENT_USER, `SOFTWARE\TacticalRMM`, registry.ALL_ACCESS)
+	}
+
 	if err != nil {
 		return nil
 	}

agent/install.go

@@ -23,6 +23,7 @@ import (
 	"strings"
 	"time"
 
+	"github.com/amidaware/rmmagent/shared"
 	"github.com/go-resty/resty/v2"
 	trmm "github.com/wh1te909/trmm-shared"
 )
@@ -146,11 +147,11 @@ func (a *Agent) Install(i *Installer) {
 		arch = "32"
 	}
 
-	var installerMeshSystemEXE string
+	var installerMeshSystemBin string
 	if len(i.MeshDir) > 0 {
-		installerMeshSystemEXE = filepath.Join(i.MeshDir, "MeshAgent.exe")
+		installerMeshSystemBin = filepath.Join(i.MeshDir, "MeshAgent.exe")
 	} else {
-		installerMeshSystemEXE = a.MeshSystemEXE
+		installerMeshSystemBin = a.MeshSystemBin
 	}
 
 	var meshNodeID string
@@ -178,7 +179,7 @@ func (a *Agent) Install(i *Installer) {
 		a.Logger.Debugln("Mesh agent:", mesh)
 		time.Sleep(1 * time.Second)
 
-		meshNodeID, err = a.installMesh(mesh, installerMeshSystemEXE, i.Proxy)
+		meshNodeID, err = a.installMesh(mesh, installerMeshSystemBin, i.Proxy)
 		if err != nil {
 			a.installerMsg(fmt.Sprintf("Failed to install mesh agent: %s", err.Error()), "error", i.Silent)
 		}
@@ -251,10 +252,16 @@ func (a *Agent) Install(i *Installer) {
 		time.Sleep(1 * time.Second)
 		a.Logger.Infoln("Starting service...")
 		out := a.ControlService(winSvcName, "start")
+
+		if shared.TEST {
+			goto SKIPSTART;
+		}
+
 		if !out.Success {
 			a.installerMsg(out.ErrorMsg, "error", i.Silent)
 		}
 
+		SKIPSTART: a.Logger.Infoln("Skipping service start in test.")
 		a.Logger.Infoln("Adding windows defender exclusions")
 		a.addDefenderExlusions()
 

agent/install_linux.go

@@ -13,14 +13,10 @@ package agent
 
 import (
 	"log"
-
+	"github.com/amidaware/rmmagent/shared"
 	"github.com/spf13/viper"
 )
 
-const (
-	etcConfig = "/etc/tacticalagent"
-)
-
 func (a *Agent) checkExistingAndRemove(silent bool) {}
 
 func (a *Agent) installerMsg(msg, alert string, silent bool) {
@@ -42,7 +38,13 @@ func createAgentConfig(baseurl, agentid, apiurl, token, agentpk, cert, proxy, me
 	viper.Set("proxy", proxy)
 	viper.Set("meshdir", meshdir)
 	viper.SetConfigPermissions(0660)
-	err := viper.SafeWriteConfigAs(etcConfig)
+	configLocation := "/etc/tacticalagent"
+	if shared.TEST {
+		configLocation = "tacticalagent"
+	}
+
+	err := viper.SafeWriteConfigAs(configLocation)
+
 	if err != nil {
 		log.Fatalln("createAgentConfig", err)
 	}

agent/install_test.go

@@ -0,0 +1,43 @@
+package agent
+
+import (
+	"testing"
+	"github.com/sirupsen/logrus"
+	"github.com/spf13/viper"
+)
+
+func TestInstall(t *testing.T) {
+	var (
+		version = "2.0.4"
+		log     = logrus.New()
+	)
+
+	a := New(log, version)
+
+	viper.SetConfigName("testargs.json")
+	viper.SetConfigType("json")
+	viper.AddConfigPath(".")
+	viper.ReadInConfig()
+
+	installer := Installer {
+		RMM:         viper.GetString("api"),
+		ClientID:    viper.GetInt("clientid"),
+		SiteID:      viper.GetInt("siteid"),
+		Description: viper.GetString("description"),
+		AgentType:   viper.GetString("agenttype"),
+		Power:       viper.GetBool("power"),
+		RDP:         viper.GetBool("rdp"),
+		Ping:        viper.GetBool("ping"),
+		Token:       viper.GetString("token"),
+		LocalMesh:   viper.GetString("localmesh"),
+		Cert:        viper.GetString("cert"),
+		Proxy:       viper.GetString("proxy"),
+		Timeout:     viper.GetDuration("timeout"),
+		Silent:      viper.GetBool("silent"),
+		NoMesh:      viper.GetBool("nomesh"),
+		MeshDir:     viper.GetString("meshdir"),
+		MeshNodeID:  viper.GetString("meshnodeid"),
+	}
+
+	a.Install(&installer)
+}

agent/install_windows.go

@@ -17,12 +17,18 @@ import (
 	"os"
 	"path/filepath"
 
+	"github.com/amidaware/rmmagent/shared"
 	"github.com/gonutz/w32/v2"
 	"golang.org/x/sys/windows/registry"
 )
 
 func createAgentConfig(baseurl, agentid, apiurl, token, agentpk, cert, proxy, meshdir string) {
 	k, _, err := registry.CreateKey(registry.LOCAL_MACHINE, `SOFTWARE\TacticalRMM`, registry.ALL_ACCESS)
+	if shared.TEST {
+		err = nil
+		k, _, err := registry.OpenKey(registry.CURRENT_USER, `SOFTWARE\TacticalRMM`, registry.ALL_ACCESS)	
+	}
+
 	if err != nil {
 		log.Fatalln("Error creating registry key:", err)
 	}
@@ -78,6 +84,11 @@ func createAgentConfig(baseurl, agentid, apiurl, token, agentpk, cert, proxy, me
 func (a *Agent) checkExistingAndRemove(silent bool) {
 	hasReg := false
 	_, err := registry.OpenKey(registry.LOCAL_MACHINE, `SOFTWARE\TacticalRMM`, registry.ALL_ACCESS)
+	if shared.TEST {
+		err = nil
+		_, err = registry.OpenKey(registry.CURRENT_USER, `SOFTWARE\TacticalRMM`, registry.ALL_ACCESS)
+	}
+
 	if err == nil {
 		hasReg = true
 	}

agent/rpc.go

@@ -49,6 +49,10 @@ var (
 )
 
 func (a *Agent) RunRPC() {
+	if rmm.DEBUG {
+		a.Logger.Infoln("DEBUG BUILD STARTED")
+	}
+
 	a.Logger.Infoln("Agent service started")
 	go a.RunAsService()
 	var wg sync.WaitGroup

agent/rpc_test.go

@@ -0,0 +1,10 @@
+package agent
+
+import (
+	"testing"
+)
+
+func TestRunRPC(t *testing.T) {
+	a := New(lg, version)
+	a.RunRPC()
+}

agent/testargs.json.example

@@ -0,0 +1,19 @@
+{
+    "api": "",
+    "clientid": 1,
+    "siteid": 1,
+    "description": "",
+    "agenttype": "workstation",
+    "power": false,
+    "rdp": false,
+    "ping": false,
+    "token": "",
+    "localmesh": "",
+    "cert": "",
+    "proxy": "",
+    "timeout": 30,
+    "silent": true,
+    "nomesh": true,
+    "meshdir": "",
+    "meshnodeid": ""
+}

shared/debug.go

@@ -0,0 +1,5 @@
+//go:build DEBUG
+
+package shared
+
+const DEBUG = true

shared/nodebug.go

@@ -0,0 +1,5 @@
+//go:build !DEBUG
+
+package shared
+
+const DEBUG	= false

shared/notest.go

@@ -0,0 +1,5 @@
+//go:build !TEST
+
+package shared
+
+const TEST = false

shared/test.go

@@ -0,0 +1,5 @@
+//go:build TEST
+
+package shared
+
+const TEST = true