From 2a0537c6eb8c4d2b6665533e23fa5c4b5ba674ac Mon Sep 17 00:00:00 2001
From: wh1te909 <dcparsi@gmail.com>
Date: Sun, 22 Nov 2020 02:34:30 -0800
Subject: [PATCH] add run script

---
 agent/checks_windows.go |  2 +-
 agent/rpc_windows.go    | 35 +++++++++++++++++++----------------
 2 files changed, 20 insertions(+), 17 deletions(-)

diff --git a/agent/checks_windows.go b/agent/checks_windows.go
index 96bb209..7edde1e 100644
--- a/agent/checks_windows.go
+++ b/agent/checks_windows.go
@@ -233,7 +233,7 @@ func (a *WindowsAgent) RunScript(code string, shell string, args []string, timeo
 
 	if timedOut {
 		stdout = outb.String()
-		stderr = fmt.Sprintf("%s\nScript check timed out after %d seconds", errb.String(), timeout)
+		stderr = fmt.Sprintf("%s\nScript timed out after %d seconds", errb.String(), timeout)
 		exitcode = 98
 		a.Logger.Debugln("Script check timeout:", ctx.Err())
 	} else {
diff --git a/agent/rpc_windows.go b/agent/rpc_windows.go
index 70fa45a..0f74ed7 100644
--- a/agent/rpc_windows.go
+++ b/agent/rpc_windows.go
@@ -11,9 +11,10 @@ import (
 )
 
 type NatsMsg struct {
-	Func    string            `json:"func"`
-	Timeout int               `json:"timeout"`
-	Data    map[string]string `json:"payload"`
+	Func       string            `json:"func"`
+	Timeout    int               `json:"timeout"`
+	Data       map[string]string `json:"payload"`
+	ScriptArgs []string          `json:"script_args"`
 }
 
 func (a *WindowsAgent) RunRPC() {
@@ -54,9 +55,8 @@ func (a *WindowsAgent) RunRPC() {
 			go func(p *NatsMsg) {
 				var resp []byte
 				ret := codec.NewEncoderBytes(&resp, new(codec.MsgpackHandle))
-				logName := p.Data["logname"]
 				days, _ := strconv.Atoi(p.Data["days"])
-				evtLog := a.GetEventLog(logName, days)
+				evtLog := a.GetEventLog(p.Data["logname"], days)
 				ret.Encode(evtLog)
 				msg.Respond(resp)
 			}(payload)
@@ -74,9 +74,7 @@ func (a *WindowsAgent) RunRPC() {
 			go func(p *NatsMsg) {
 				var resp []byte
 				ret := codec.NewEncoderBytes(&resp, new(codec.MsgpackHandle))
-				command := p.Data["command"]
-				shell := p.Data["shell"]
-				out, _ := CMDShell(shell, []string{}, command, payload.Timeout, false)
+				out, _ := CMDShell(p.Data["shell"], []string{}, p.Data["command"], p.Timeout, false)
 
 				if out[1] != "" {
 					ret.Encode(out[1])
@@ -100,8 +98,7 @@ func (a *WindowsAgent) RunRPC() {
 			go func(p *NatsMsg) {
 				var resp []byte
 				ret := codec.NewEncoderBytes(&resp, new(codec.MsgpackHandle))
-				svcName := p.Data["name"]
-				svc := a.GetServiceDetail(svcName)
+				svc := a.GetServiceDetail(p.Data["name"])
 				ret.Encode(svc)
 				msg.Respond(resp)
 			}(payload)
@@ -110,9 +107,7 @@ func (a *WindowsAgent) RunRPC() {
 			go func(p *NatsMsg) {
 				var resp []byte
 				ret := codec.NewEncoderBytes(&resp, new(codec.MsgpackHandle))
-				svcName := p.Data["name"]
-				action := p.Data["action"]
-				retData := a.ControlService(svcName, action)
+				retData := a.ControlService(p.Data["name"], p.Data["action"])
 				ret.Encode(retData)
 				msg.Respond(resp)
 			}(payload)
@@ -121,9 +116,17 @@ func (a *WindowsAgent) RunRPC() {
 			go func(p *NatsMsg) {
 				var resp []byte
 				ret := codec.NewEncoderBytes(&resp, new(codec.MsgpackHandle))
-				svcName := p.Data["name"]
-				startType := p.Data["startType"]
-				retData := a.EditService(svcName, startType)
+				retData := a.EditService(p.Data["name"], p.Data["startType"])
+				ret.Encode(retData)
+				msg.Respond(resp)
+			}(payload)
+
+		case "runscript":
+			go func(p *NatsMsg) {
+				var resp []byte
+				ret := codec.NewEncoderBytes(&resp, new(codec.MsgpackHandle))
+				out, err, _, _ := a.RunScript(p.Data["code"], p.Data["shell"], p.ScriptArgs, p.Timeout)
+				retData := out + err
 				ret.Encode(retData)
 				msg.Respond(resp)
 			}(payload)