From 6ac14b6d6426bf66ec05d591e657cb66818cd138 Mon Sep 17 00:00:00 2001 From: wh1te909 Date: Fri, 25 Nov 2022 00:23:31 -0800 Subject: [PATCH] add env to scripts --- agent/agent_windows.go | 7 ++++++- agent/checks.go | 2 +- agent/choco_windows.go | 2 +- agent/rpc.go | 5 +++-- agent/tasks_windows.go | 2 +- shared/types.go | 8 +++++--- 6 files changed, 17 insertions(+), 9 deletions(-) diff --git a/agent/agent_windows.go b/agent/agent_windows.go index 26cb7b6..c980ebd 100644 --- a/agent/agent_windows.go +++ b/agent/agent_windows.go @@ -82,7 +82,7 @@ func NewAgentConfig() *rmm.AgentConfig { } } -func (a *Agent) RunScript(code string, shell string, args []string, timeout int, runasuser bool) (stdout, stderr string, exitcode int, e error) { +func (a *Agent) RunScript(code string, shell string, args []string, timeout int, runasuser bool, envs []string) (stdout, stderr string, exitcode int, e error) { content := []byte(code) @@ -158,6 +158,11 @@ func (a *Agent) RunScript(code string, shell string, args []string, timeout int, cmd.Stdout = &outb cmd.Stderr = &errb + if len(envs) > 0 { + cmd.Env = os.Environ() + cmd.Env = append(cmd.Env, envs...) + } + if cmdErr := cmd.Start(); cmdErr != nil { a.Logger.Debugln(cmdErr) return "", cmdErr.Error(), 65, cmdErr diff --git a/agent/checks.go b/agent/checks.go index 531bf74..a680917 100644 --- a/agent/checks.go +++ b/agent/checks.go @@ -169,7 +169,7 @@ type ScriptCheckResult struct { // ScriptCheck runs either bat, powershell or python script func (a *Agent) ScriptCheck(data rmm.Check, r *resty.Client) { start := time.Now() - stdout, stderr, retcode, _ := a.RunScript(data.Script.Code, data.Script.Shell, data.ScriptArgs, data.Timeout, data.Script.RunAsUser) + stdout, stderr, retcode, _ := a.RunScript(data.Script.Code, data.Script.Shell, data.ScriptArgs, data.Timeout, data.Script.RunAsUser, data.Script.Env) payload := ScriptCheckResult{ ID: data.CheckPK, diff --git a/agent/choco_windows.go b/agent/choco_windows.go index c37f025..1ba6058 100644 --- a/agent/choco_windows.go +++ b/agent/choco_windows.go @@ -42,7 +42,7 @@ func (a *Agent) InstallChoco() { return } - _, _, exitcode, err := a.RunScript(string(r.Body()), "powershell", []string{}, 900, false) + _, _, exitcode, err := a.RunScript(string(r.Body()), "powershell", []string{}, 900, false, []string{}) if err != nil { a.Logger.Debugln(err) a.rClient.R().SetBody(result).Post(url) diff --git a/agent/rpc.go b/agent/rpc.go index f5dc454..57f0987 100644 --- a/agent/rpc.go +++ b/agent/rpc.go @@ -41,6 +41,7 @@ type NatsMsg struct { ID int `json:"id"` Code string `json:"code"` RunAsUser bool `json:"run_as_user"` + Env []string `json:"env"` } var ( @@ -261,7 +262,7 @@ func (a *Agent) RunRPC() { var resultData rmm.RunScriptResp ret := codec.NewEncoderBytes(&resp, new(codec.MsgpackHandle)) start := time.Now() - stdout, stderr, retcode, err := a.RunScript(p.Data["code"], p.Data["shell"], p.ScriptArgs, p.Timeout, p.RunAsUser) + stdout, stderr, retcode, err := a.RunScript(p.Data["code"], p.Data["shell"], p.ScriptArgs, p.Timeout, p.RunAsUser, p.Env) resultData.ExecTime = time.Since(start).Seconds() resultData.ID = p.ID @@ -291,7 +292,7 @@ func (a *Agent) RunRPC() { var retData rmm.RunScriptResp ret := codec.NewEncoderBytes(&resp, new(codec.MsgpackHandle)) start := time.Now() - stdout, stderr, retcode, err := a.RunScript(p.Data["code"], p.Data["shell"], p.ScriptArgs, p.Timeout, p.RunAsUser) + stdout, stderr, retcode, err := a.RunScript(p.Data["code"], p.Data["shell"], p.ScriptArgs, p.Timeout, p.RunAsUser, p.Env) retData.ExecTime = time.Since(start).Seconds() if err != nil { diff --git a/agent/tasks_windows.go b/agent/tasks_windows.go index da3d7a2..3144c8f 100644 --- a/agent/tasks_windows.go +++ b/agent/tasks_windows.go @@ -59,7 +59,7 @@ func (a *Agent) RunTask(id int) error { action_start := time.Now() if action.ActionType == "script" { - stdout, stderr, retcode, err := a.RunScript(action.Code, action.Shell, action.Args, action.Timeout, action.RunAsUser) + stdout, stderr, retcode, err := a.RunScript(action.Code, action.Shell, action.Args, action.Timeout, action.RunAsUser, action.Env) if err != nil { a.Logger.Debugln(err) diff --git a/shared/types.go b/shared/types.go index b2cb258..90acdc5 100644 --- a/shared/types.go +++ b/shared/types.go @@ -141,9 +141,10 @@ type AssignedTask struct { } type Script struct { - Shell string `json:"shell"` - Code string `json:"code"` - RunAsUser bool `json:"run_as_user"` + Shell string `json:"shell"` + Code string `json:"code"` + RunAsUser bool `json:"run_as_user"` + Env []string `json:"env"` } type CheckInfo struct { @@ -190,6 +191,7 @@ type TaskAction struct { Args []string `json:"script_args"` Timeout int `json:"timeout"` RunAsUser bool `json:"run_as_user"` + Env []string `json:"env"` } type AutomatedTask struct {