From 6d8cd81e197f02f73574476eb6d7fd5421b50e9a Mon Sep 17 00:00:00 2001
From: wh1te909 <dcparsi@gmail.com>
Date: Tue, 17 Nov 2020 00:14:56 -0800
Subject: [PATCH] start rpc, add procs

---
 agent/process_windows.go |  65 ++++++++++++++++++++++++
 agent/rpc_windows.go     | 103 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 168 insertions(+)
 create mode 100644 agent/process_windows.go
 create mode 100644 agent/rpc_windows.go

diff --git a/agent/process_windows.go b/agent/process_windows.go
new file mode 100644
index 0000000..f1f705a
--- /dev/null
+++ b/agent/process_windows.go
@@ -0,0 +1,65 @@
+package agent
+
+import (
+	"fmt"
+
+	ps "github.com/elastic/go-sysinfo"
+	gops "github.com/shirou/gopsutil/v3/process"
+)
+
+type ProcessMsg struct {
+	Name     string `json:"name"`
+	Pid      int    `json:"pid"`
+	Mem      string `json:"memory_percent"`
+	Username string `json:"username"`
+	UID      int    `json:"id"`
+	CPU      string `json:"cpu_percent"`
+}
+
+func (a *WindowsAgent) GetProcsRPC() []ProcessMsg {
+	ret := make([]ProcessMsg, 0)
+
+	procs, _ := ps.Processes()
+	for i, process := range procs {
+		p, err := process.Info()
+		if err != nil {
+			continue
+		}
+		if p.PID == 0 {
+			continue
+		}
+
+		m, _ := process.Memory()
+		proc, gerr := gops.NewProcess(int32(p.PID))
+		if gerr != nil {
+			continue
+		}
+		cpu, _ := proc.CPUPercent()
+		user, _ := proc.Username()
+
+		ret = append(ret, ProcessMsg{
+			Name:     p.Name,
+			Pid:      p.PID,
+			Mem:      ByteCountSI(m.Resident),
+			Username: user,
+			UID:      i,
+			CPU:      fmt.Sprintf("%.1f", cpu),
+		})
+	}
+	return ret
+}
+
+// https://yourbasic.org/golang/formatting-byte-size-to-human-readable-format/
+func ByteCountSI(b uint64) string {
+	const unit = 1000
+	if b < unit {
+		return fmt.Sprintf("%d B", b)
+	}
+	div, exp := int64(unit), 0
+	for n := b / unit; n >= unit; n /= unit {
+		div *= unit
+		exp++
+	}
+	return fmt.Sprintf("%.1f %cB",
+		float64(b)/float64(div), "kMGTPE"[exp])
+}
diff --git a/agent/rpc_windows.go b/agent/rpc_windows.go
new file mode 100644
index 0000000..13c83a2
--- /dev/null
+++ b/agent/rpc_windows.go
@@ -0,0 +1,103 @@
+package agent
+
+import (
+	"fmt"
+	"runtime"
+	"strconv"
+	"time"
+
+	nats "github.com/nats-io/nats.go"
+	"github.com/ugorji/go/codec"
+)
+
+type NatsMsg struct {
+	Func    string            `json:"func"`
+	Timeout int               `json:"timeout"`
+	Data    map[string]string `json:"payload"`
+}
+
+func (a *WindowsAgent) RunRPC() {
+	opts := []nats.Option{nats.Name("TacticalRMM"), nats.UserInfo(a.AgentID, a.Token)}
+	opts = setupConnOptions(opts)
+
+	server := fmt.Sprintf("tls://%s:4222", a.SaltMaster)
+	nc, err := nats.Connect(server, opts...)
+	if err != nil {
+		fmt.Println(err)
+	}
+
+	nc.Subscribe(a.AgentID, func(msg *nats.Msg) {
+
+		var payload *NatsMsg
+		var mh codec.MsgpackHandle
+		mh.RawToString = true
+
+		dec := codec.NewDecoderBytes(msg.Data, &mh)
+		if err := dec.Decode(&payload); err != nil {
+			fmt.Println(err)
+			return
+		}
+
+		switch payload.Func {
+		case "ping":
+			go func() {
+				var resp []byte
+				ret := codec.NewEncoderBytes(&resp, new(codec.MsgpackHandle))
+				ret.Encode("pong")
+				msg.Respond(resp)
+			}()
+
+		case "eventlog":
+			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)
+				ret.Encode(evtLog)
+				msg.Respond(resp)
+			}(payload)
+
+		case "procs":
+			go func() {
+				var resp []byte
+				ret := codec.NewEncoderBytes(&resp, new(codec.MsgpackHandle))
+				procs := a.GetProcsRPC()
+				ret.Encode(procs)
+				msg.Respond(resp)
+			}()
+
+		case "rawcmd":
+			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)
+
+				if out[1] != "" {
+					ret.Encode(out[1])
+				} else {
+					ret.Encode(out[0])
+				}
+
+				msg.Respond(resp)
+			}(payload)
+		}
+	})
+	nc.Flush()
+
+	if err := nc.LastError(); err != nil {
+		fmt.Println(err)
+		return
+	}
+
+	runtime.Goexit()
+}
+
+func setupConnOptions(opts []nats.Option) []nats.Option {
+	opts = append(opts, nats.ReconnectWait(time.Second*5))
+	opts = append(opts, nats.RetryOnFailedConnect(true))
+	opts = append(opts, nats.MaxReconnects(-1))
+	return opts
+}