Mirrored_Repos/rmmagent
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

add env to scripts

Browse Source
This commit is contained in:
wh1te909 2022-11-25 00:23:31 -08:00
parent 9565fea27c
commit 6ac14b6d64
6 changed files with 17 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
agent/agent_windows.go
View File

@ -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) content := []byte(code)
@ -158,6 +158,11 @@ func (a *Agent) RunScript(code string, shell string, args []string, timeout int,
cmd.Stdout = &outb cmd.Stdout = &outb
cmd.Stderr = &errb cmd.Stderr = &errb
if len(envs) > 0 {
cmd.Env = os.Environ()
cmd.Env = append(cmd.Env, envs...)
}
if cmdErr := cmd.Start(); cmdErr != nil { if cmdErr := cmd.Start(); cmdErr != nil {
a.Logger.Debugln(cmdErr) a.Logger.Debugln(cmdErr)
return "", cmdErr.Error(), 65, cmdErr return "", cmdErr.Error(), 65, cmdErr

2
agent/checks.go
View File

@ -169,7 +169,7 @@ type ScriptCheckResult struct {
// ScriptCheck runs either bat, powershell or python script // ScriptCheck runs either bat, powershell or python script
func (a *Agent) ScriptCheck(data rmm.Check, r *resty.Client) { func (a *Agent) ScriptCheck(data rmm.Check, r *resty.Client) {
start := time.Now() 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{ payload := ScriptCheckResult{
ID: data.CheckPK, ID: data.CheckPK,

2
agent/choco_windows.go
View File

@ -42,7 +42,7 @@ func (a *Agent) InstallChoco() {
return 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 { if err != nil {
a.Logger.Debugln(err) a.Logger.Debugln(err)
a.rClient.R().SetBody(result).Post(url) a.rClient.R().SetBody(result).Post(url)

5
agent/rpc.go
View File

@ -41,6 +41,7 @@ type NatsMsg struct {
ID int `json:"id"` ID int `json:"id"`
Code string `json:"code"` Code string `json:"code"`
RunAsUser bool `json:"run_as_user"` RunAsUser bool `json:"run_as_user"`
Env []string `json:"env"`
} }
var ( var (
@ -261,7 +262,7 @@ func (a *Agent) RunRPC() {
var resultData rmm.RunScriptResp var resultData rmm.RunScriptResp
ret := codec.NewEncoderBytes(&resp, new(codec.MsgpackHandle)) ret := codec.NewEncoderBytes(&resp, new(codec.MsgpackHandle))
start := time.Now() 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.ExecTime = time.Since(start).Seconds()
resultData.ID = p.ID resultData.ID = p.ID
@ -291,7 +292,7 @@ func (a *Agent) RunRPC() {
var retData rmm.RunScriptResp var retData rmm.RunScriptResp
ret := codec.NewEncoderBytes(&resp, new(codec.MsgpackHandle)) ret := codec.NewEncoderBytes(&resp, new(codec.MsgpackHandle))
start := time.Now() 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() retData.ExecTime = time.Since(start).Seconds()
if err != nil { if err != nil {

2
agent/tasks_windows.go
View File

@ -59,7 +59,7 @@ func (a *Agent) RunTask(id int) error {
action_start := time.Now() action_start := time.Now()
if action.ActionType == "script" { 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 { if err != nil {
a.Logger.Debugln(err) a.Logger.Debugln(err)

8
shared/types.go
View File

@ -141,9 +141,10 @@ type AssignedTask struct {
} }
type Script struct { type Script struct {
Shell string `json:"shell"` Shell string `json:"shell"`
Code string `json:"code"` Code string `json:"code"`
RunAsUser bool `json:"run_as_user"` RunAsUser bool `json:"run_as_user"`
Env []string `json:"env"`
} }
type CheckInfo struct { type CheckInfo struct {
@ -190,6 +191,7 @@ type TaskAction struct {
Args []string `json:"script_args"` Args []string `json:"script_args"`
Timeout int `json:"timeout"` Timeout int `json:"timeout"`
RunAsUser bool `json:"run_as_user"` RunAsUser bool `json:"run_as_user"`
Env []string `json:"env"`
} }
type AutomatedTask struct { type AutomatedTask struct {