From 9ede622837658288f543724f7b79708afa053dde Mon Sep 17 00:00:00 2001
From: wh1te909 <dcparsi@gmail.com>
Date: Fri, 20 Nov 2020 01:20:19 -0800
Subject: [PATCH] add agent update

---
 agent/agent_windows.go       | 27 +++++++++++++++++++++++++++
 agent/winagentsvc_windows.go |  9 +++++++++
 2 files changed, 36 insertions(+)

diff --git a/agent/agent_windows.go b/agent/agent_windows.go
index 8178339..31da6b5 100644
--- a/agent/agent_windows.go
+++ b/agent/agent_windows.go
@@ -706,6 +706,33 @@ func (a *WindowsAgent) installerMsg(msg, alert string) {
 	}
 }
 
+func (a *WindowsAgent) AgentUpdate(url, inno, version string) {
+	updater := filepath.Join(a.ProgramDir, inno)
+	a.Logger.Infof("Agent updating from %s to %s", a.Version, version)
+	a.Logger.Infoln("Downloading agent update from", url)
+
+	rClient := resty.New()
+	rClient.SetCloseConnection(true)
+	rClient.SetTimeout(15 * time.Minute)
+	rClient.SetDebug(a.Debug)
+	r, err := rClient.R().SetOutput(updater).Get(url)
+	if err != nil {
+		a.Logger.Errorln(err)
+		return
+	}
+	if r.IsError() {
+		a.Logger.Errorln("Download failed with status code", r.StatusCode())
+		return
+	}
+
+	args := []string{"/C", updater, "/VERYSILENT", "/SUPPRESSMSGBOXES"}
+	cmd := exec.Command("cmd.exe", args...)
+	cmd.SysProcAttr = &windows.SysProcAttr{
+		CreationFlags: windows.DETACHED_PROCESS | windows.CREATE_NEW_PROCESS_GROUP,
+	}
+	cmd.Start()
+}
+
 // CleanupPythonAgent cleans up files from the old python agent if this is an upgrade
 func (a *WindowsAgent) CleanupPythonAgent() {
 	cderr := os.Chdir(a.ProgramDir)
diff --git a/agent/winagentsvc_windows.go b/agent/winagentsvc_windows.go
index a8dcce8..a6a3504 100644
--- a/agent/winagentsvc_windows.go
+++ b/agent/winagentsvc_windows.go
@@ -105,6 +105,7 @@ func (a *WindowsAgent) WinAgentSvc() {
 				if err := json.Unmarshal(r.Body(), &data); err != nil {
 					a.Logger.Debugln(err)
 				} else {
+					// recovery
 					if action, ok := data["recovery"].(string); ok {
 						switch action {
 						case "salt":
@@ -117,6 +118,14 @@ func (a *WindowsAgent) WinAgentSvc() {
 							}
 						}
 					}
+					// agent update
+					if version, ok := data["version"].(string); ok {
+						if inno, iok := data["inno"].(string); iok {
+							if url, uok := data["url"].(string); uok {
+								a.AgentUpdate(url, inno, version)
+							}
+						}
+					}
 				}
 			}
 		}