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 +}