Skip to content

Commit

Permalink
allow setting custom defined tmpdir
Browse files Browse the repository at this point in the history
  • Loading branch information
@wh1te909
wh1te909 committed Apr 29, 2023
1 parent ec49d49 commit 588a4bc
Show file tree
Hide file tree
Showing 3 changed files with 122 additions and 95 deletions.
142 changes: 78 additions & 64 deletions agent/agent.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,37 +40,39 @@ import (

// Agent struct
type Agent struct {
Hostname string
Arch string
AgentID string
BaseURL string
ApiURL string
Token string
AgentPK int
Cert string
ProgramDir string
EXE string
SystemDrive string
MeshInstaller string
MeshSystemEXE string
MeshSVC string
PyBin string
Headers map[string]string
Logger *logrus.Logger
Version string
Debug bool
rClient *resty.Client
Proxy string
LogTo string
LogFile *os.File
Platform string
GoArch string
ServiceConfig *service.Config
NatsServer string
NatsProxyPath string
NatsProxyPort string
NatsPingInterval int
NatsWSCompression bool
Hostname string
Arch string
AgentID string
BaseURL string
ApiURL string
Token string
AgentPK int
Cert string
ProgramDir string
EXE string
SystemDrive string
WinTmpDir string
WinRunAsUserTmpDir string
MeshInstaller string
MeshSystemEXE string
MeshSVC string
PyBin string
Headers map[string]string
Logger *logrus.Logger
Version string
Debug bool
rClient *resty.Client
Proxy string
LogTo string
LogFile *os.File
Platform string
GoArch string
ServiceConfig *service.Config
NatsServer string
NatsProxyPath string
NatsProxyPort string
NatsPingInterval int
NatsWSCompression bool
}

const (
Expand All @@ -88,7 +90,7 @@ const (
defaultMacMeshSvcDir = "/usr/local/mesh_services"
)

var winTempDir = filepath.Join(os.Getenv("PROGRAMDATA"), "TacticalRMM")
var defaultWinTmpDir = filepath.Join(os.Getenv("PROGRAMDATA"), "TacticalRMM")
var winMeshDir = filepath.Join(os.Getenv("PROGRAMFILES"), "Mesh Agent")
var natsCheckin = []string{"agent-hello", "agent-agentinfo", "agent-disks", "agent-winsvc", "agent-publicip", "agent-wmi"}
var limitNatsData = []string{"agent-winsvc", "agent-wmi"}
Expand All @@ -99,6 +101,8 @@ func New(logger *logrus.Logger, version string) *Agent {
pd := filepath.Join(os.Getenv("ProgramFiles"), progFilesName)
exe := filepath.Join(pd, winExeName)
sd := os.Getenv("SystemDrive")
winTempDir := defaultWinTmpDir
winRunAsUserTmpDir := defaultWinTmpDir

var pybin string
switch runtime.GOARCH {
Expand Down Expand Up @@ -130,6 +134,14 @@ func New(logger *logrus.Logger, version string) *Agent {
restyC.SetRootCertificate(ac.Cert)
}

if len(ac.WinTmpDir) > 0 {
winTempDir = ac.WinTmpDir
}

if len(ac.WinRunAsUserTmpDir) > 0 {
winRunAsUserTmpDir = ac.WinRunAsUserTmpDir
}

var MeshSysExe string
switch runtime.GOOS {
case "windows":
Expand Down Expand Up @@ -189,34 +201,36 @@ func New(logger *logrus.Logger, version string) *Agent {
}

return &Agent{
Hostname: info.Hostname,
BaseURL: ac.BaseURL,
AgentID: ac.AgentID,
ApiURL: ac.APIURL,
Token: ac.Token,
AgentPK: ac.PK,
Cert: ac.Cert,
ProgramDir: pd,
EXE: exe,
SystemDrive: sd,
MeshInstaller: "meshagent.exe",
MeshSystemEXE: MeshSysExe,
MeshSVC: meshSvcName,
PyBin: pybin,
Headers: headers,
Logger: logger,
Version: version,
Debug: logger.IsLevelEnabled(logrus.DebugLevel),
rClient: restyC,
Proxy: ac.Proxy,
Platform: runtime.GOOS,
GoArch: runtime.GOARCH,
ServiceConfig: svcConf,
NatsServer: natsServer,
NatsProxyPath: natsProxyPath,
NatsProxyPort: natsProxyPort,
NatsPingInterval: natsPingInterval,
NatsWSCompression: natsWsCompression,
Hostname: info.Hostname,
BaseURL: ac.BaseURL,
AgentID: ac.AgentID,
ApiURL: ac.APIURL,
Token: ac.Token,
AgentPK: ac.PK,
Cert: ac.Cert,
ProgramDir: pd,
EXE: exe,
SystemDrive: sd,
WinTmpDir: winTempDir,
WinRunAsUserTmpDir: winRunAsUserTmpDir,
MeshInstaller: "meshagent.exe",
MeshSystemEXE: MeshSysExe,
MeshSVC: meshSvcName,
PyBin: pybin,
Headers: headers,
Logger: logger,
Version: version,
Debug: logger.IsLevelEnabled(logrus.DebugLevel),
rClient: restyC,
Proxy: ac.Proxy,
Platform: runtime.GOOS,
GoArch: runtime.GOARCH,
ServiceConfig: svcConf,
NatsServer: natsServer,
NatsProxyPath: natsProxyPath,
NatsProxyPort: natsProxyPort,
NatsPingInterval: natsPingInterval,
NatsWSCompression: natsWsCompression,
}
}

Expand Down Expand Up @@ -457,7 +471,7 @@ func (a *Agent) GetUninstallExe() string {

func (a *Agent) CleanupAgentUpdates() {
// TODO remove a.ProgramDir, updates are now in winTempDir
dirs := [3]string{winTempDir, os.Getenv("TMP"), a.ProgramDir}
dirs := [3]string{a.WinTmpDir, os.Getenv("TMP"), a.ProgramDir}
for _, dir := range dirs {
err := os.Chdir(dir)
if err != nil {
Expand Down Expand Up @@ -491,7 +505,7 @@ func (a *Agent) CleanupAgentUpdates() {

func (a *Agent) RunPythonCode(code string, timeout int, args []string) (string, error) {
content := []byte(code)
tmpfn, _ := ioutil.TempFile(winTempDir, "*.py")
tmpfn, _ := ioutil.TempFile(a.WinTmpDir, "*.py")
if _, err := tmpfn.Write(content); err != nil {
a.Logger.Debugln(err)
return "", err
Expand Down Expand Up @@ -537,8 +551,8 @@ func (a *Agent) RunPythonCode(code string, timeout int, args []string) (string,
}

func createWinTempDir() error {
if !trmm.FileExists(winTempDir) {
err := os.Mkdir(winTempDir, 0775)
if !trmm.FileExists(defaultWinTmpDir) {
err := os.Mkdir(defaultWinTmpDir, 0775)
if err != nil {
return err
}
Expand Down
47 changes: 29 additions & 18 deletions agent/agent_windows.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -62,26 +62,30 @@ func NewAgentConfig() *rmm.AgentConfig {
cert, _, _ := k.GetStringValue("Cert")
proxy, _, _ := k.GetStringValue("Proxy")
customMeshDir, _, _ := k.GetStringValue("MeshDir")
winTmpDir, _, _ := k.GetStringValue("WinTmpDir")
winRunAsUserTmpDir, _, _ := k.GetStringValue("WinRunAsUserTmpDir")
natsProxyPath, _, _ := k.GetStringValue("NatsProxyPath")
natsProxyPort, _, _ := k.GetStringValue("NatsProxyPort")
natsStandardPort, _, _ := k.GetStringValue("NatsStandardPort")
natsPingInterval, _, _ := k.GetStringValue("NatsPingInterval")
npi, _ := strconv.Atoi(natsPingInterval)

return &rmm.AgentConfig{
BaseURL: baseurl,
AgentID: agentid,
APIURL: apiurl,
Token: token,
AgentPK: agentpk,
PK: pk,
Cert: cert,
Proxy: proxy,
CustomMeshDir: customMeshDir,
NatsProxyPath: natsProxyPath,
NatsProxyPort: natsProxyPort,
NatsStandardPort: natsStandardPort,
NatsPingInterval: npi,
BaseURL: baseurl,
AgentID: agentid,
APIURL: apiurl,
Token: token,
AgentPK: agentpk,
PK: pk,
Cert: cert,
Proxy: proxy,
CustomMeshDir: customMeshDir,
WinTmpDir: winTmpDir,
WinRunAsUserTmpDir: winRunAsUserTmpDir,
NatsProxyPath: natsProxyPath,
NatsProxyPort: natsProxyPort,
NatsStandardPort: natsStandardPort,
NatsPingInterval: npi,
}
}

Expand Down Expand Up @@ -114,7 +118,13 @@ func (a *Agent) RunScript(code string, shell string, args []string, timeout int,
ext = "*.bat"
}

tmpfn, err := ioutil.TempFile(winTempDir, ext)
tmpDir := a.WinTmpDir

if runasuser {
tmpDir = a.WinRunAsUserTmpDir
}

tmpfn, err := ioutil.TempFile(tmpDir, ext)
if err != nil {
a.Logger.Errorln(err)
return "", err.Error(), 85, err
Expand All @@ -133,7 +143,7 @@ func (a *Agent) RunScript(code string, shell string, args []string, timeout int,
switch shell {
case "powershell":
exe = getPowershellExe()
cmdArgs = []string{"-NonInteractive", "-NoProfile", "-ExecutionPolicy", "Bypass", tmpfn.Name()}
cmdArgs = []string{"-NonInteractive", "-NoProfile", "-ExecutionPolicy", "Bypass", "-File", tmpfn.Name()}
case "python":
exe = a.PyBin
cmdArgs = []string{tmpfn.Name()}
Expand Down Expand Up @@ -587,15 +597,16 @@ func (a *Agent) UninstallCleanup() {
a.PatchMgmnt(false)
a.CleanupAgentUpdates()
CleanupSchedTasks()
os.RemoveAll(winTempDir)
os.RemoveAll(a.WinTmpDir)
os.RemoveAll(a.WinRunAsUserTmpDir)
}

func (a *Agent) AgentUpdate(url, inno, version string) {
time.Sleep(time.Duration(randRange(1, 15)) * time.Second)
a.KillHungUpdates()
time.Sleep(1 * time.Second)
a.CleanupAgentUpdates()
updater := filepath.Join(winTempDir, inno)
updater := filepath.Join(a.WinTmpDir, inno)
a.Logger.Infof("Agent updating from %s to %s", a.Version, version)
a.Logger.Debugln("Downloading agent update from", url)

Expand All @@ -618,7 +629,7 @@ func (a *Agent) AgentUpdate(url, inno, version string) {
return
}

innoLogFile := filepath.Join(winTempDir, fmt.Sprintf("tacticalagent_update_v%s.txt", version))
innoLogFile := filepath.Join(a.WinTmpDir, fmt.Sprintf("tacticalagent_update_v%s.txt", version))

args := []string{"/C", updater, "/VERYSILENT", fmt.Sprintf("/LOG=%s", innoLogFile)}
cmd := exec.Command("cmd.exe", args...)
Expand Down
28 changes: 15 additions & 13 deletions shared/types.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -33,19 +33,21 @@ type ProcessMsg struct {
}

type AgentConfig struct {
BaseURL string
AgentID string
APIURL string
Token string
AgentPK string
PK int
Cert string
Proxy string
CustomMeshDir string
NatsProxyPath string
NatsProxyPort string
NatsStandardPort string
NatsPingInterval int
BaseURL string
AgentID string
APIURL string
Token string
AgentPK string
PK int
Cert string
Proxy string
CustomMeshDir string
WinTmpDir string
WinRunAsUserTmpDir string
NatsProxyPath string
NatsProxyPort string
NatsStandardPort string
NatsPingInterval int
}

type RunScriptResp struct {
Expand Down

0 comments on commit 588a4bc

Please sign in to comment.