From f1db416d565dfce8c48534be35b9fb07b21457b7 Mon Sep 17 00:00:00 2001 From: wh1te909 Date: Thu, 24 Nov 2022 23:45:59 -0800 Subject: [PATCH] make checkin intervals configurable and remove extra nats conn --- agent/agent.go | 1 + agent/rpc.go | 9 ++++-- agent/svc.go | 84 +++++++++++++++++++++++++++++++++++++------------- agent/utils.go | 9 ++++++ 4 files changed, 79 insertions(+), 24 deletions(-) diff --git a/agent/agent.go b/agent/agent.go index c9b9b04..daa3503 100644 --- a/agent/agent.go +++ b/agent/agent.go @@ -89,6 +89,7 @@ const ( var winTempDir = filepath.Join(os.Getenv("PROGRAMDATA"), "TacticalRMM") var winMeshDir = filepath.Join(os.Getenv("PROGRAMFILES"), "Mesh Agent") var natsCheckin = []string{"agent-hello", "agent-agentinfo", "agent-disks", "agent-winsvc", "agent-publicip", "agent-wmi"} +var limitNatsData = []string{"agent-winsvc", "agent-wmi"} func New(logger *logrus.Logger, version string) *Agent { host, _ := ps.Host() diff --git a/agent/rpc.go b/agent/rpc.go index b141be0..f5dc454 100644 --- a/agent/rpc.go +++ b/agent/rpc.go @@ -51,15 +51,18 @@ var ( func (a *Agent) RunRPC() { a.Logger.Infoln("Agent service started") - go a.RunAsService() - var wg sync.WaitGroup - wg.Add(1) + opts := a.setupNatsOptions() nc, err := nats.Connect(a.NatsServer, opts...) if err != nil { a.Logger.Fatalln("RunRPC() nats.Connect()", err) } + go a.RunAsService(nc) + + var wg sync.WaitGroup + wg.Add(1) + nc.Subscribe(a.AgentID, func(msg *nats.Msg) { var payload *NatsMsg var mh codec.MsgpackHandle diff --git a/agent/svc.go b/agent/svc.go index e695ffd..24d72d8 100644 --- a/agent/svc.go +++ b/agent/svc.go @@ -12,6 +12,7 @@ https://license.tacticalrmm.com package agent import ( + "fmt" "runtime" "sync" "time" @@ -19,15 +20,27 @@ import ( nats "github.com/nats-io/nats.go" ) -func (a *Agent) RunAsService() { +func (a *Agent) RunAsService(nc *nats.Conn) { var wg sync.WaitGroup wg.Add(1) - go a.AgentSvc() + go a.AgentSvc(nc) go a.CheckRunner() wg.Wait() } -func (a *Agent) AgentSvc() { +type AgentCheckInConfig struct { + Hello int `json:"checkin_hello"` + AgentInfo int `json:"checkin_agentinfo"` + WinSvc int `json:"checkin_winsvc"` + PubIP int `json:"checkin_pubip"` + Disks int `json:"checkin_disks"` + SW int `json:"checkin_sw"` + WMI int `json:"checkin_wmi"` + SyncMesh int `json:"checkin_syncmesh"` + LimitData bool `json:"limit_data"` +} + +func (a *Agent) AgentSvc(nc *nats.Conn) { if runtime.GOOS == "windows" { go a.GetPython(false) @@ -38,37 +51,36 @@ func (a *Agent) AgentSvc() { } a.RunMigrations() - sleepDelay := randRange(14, 22) + sleepDelay := randRange(7, 25) a.Logger.Debugf("AgentSvc() sleeping for %v seconds", sleepDelay) time.Sleep(time.Duration(sleepDelay) * time.Second) - opts := a.setupNatsOptions() - nc, err := nats.Connect(a.NatsServer, opts...) - if err != nil { - a.Logger.Fatalln("AgentSvc() nats.Connect()", err) - } - + conf := a.GetAgentCheckInConfig() for _, s := range natsCheckin { - a.NatsMessage(nc, s) - time.Sleep(time.Duration(randRange(100, 400)) * time.Millisecond) + if conf.LimitData && stringInSlice(s, limitNatsData) { + continue + } else { + a.NatsMessage(nc, s) + time.Sleep(time.Duration(randRange(100, 400)) * time.Millisecond) + } } go a.SyncMeshNodeID() time.Sleep(time.Duration(randRange(1, 3)) * time.Second) - if runtime.GOOS == "windows" { + if runtime.GOOS == "windows" && !conf.LimitData { a.AgentStartup() a.SendSoftware() } - checkInHelloTicker := time.NewTicker(time.Duration(randRange(30, 60)) * time.Second) - checkInAgentInfoTicker := time.NewTicker(time.Duration(randRange(200, 400)) * time.Second) - checkInWinSvcTicker := time.NewTicker(time.Duration(randRange(2400, 3000)) * time.Second) - checkInPubIPTicker := time.NewTicker(time.Duration(randRange(300, 500)) * time.Second) - checkInDisksTicker := time.NewTicker(time.Duration(randRange(1000, 2000)) * time.Second) - checkInSWTicker := time.NewTicker(time.Duration(randRange(2800, 3500)) * time.Second) - checkInWMITicker := time.NewTicker(time.Duration(randRange(3000, 4000)) * time.Second) - syncMeshTicker := time.NewTicker(time.Duration(randRange(800, 1200)) * time.Second) + checkInHelloTicker := time.NewTicker(time.Duration(conf.Hello) * time.Second) + checkInAgentInfoTicker := time.NewTicker(time.Duration(conf.AgentInfo) * time.Second) + checkInWinSvcTicker := time.NewTicker(time.Duration(conf.WinSvc) * time.Second) + checkInPubIPTicker := time.NewTicker(time.Duration(conf.PubIP) * time.Second) + checkInDisksTicker := time.NewTicker(time.Duration(conf.Disks) * time.Second) + checkInSWTicker := time.NewTicker(time.Duration(conf.SW) * time.Second) + checkInWMITicker := time.NewTicker(time.Duration(conf.WMI) * time.Second) + syncMeshTicker := time.NewTicker(time.Duration(conf.SyncMesh) * time.Second) for { select { @@ -100,3 +112,33 @@ func (a *Agent) AgentStartup() { a.Logger.Debugln("AgentStartup()", err) } } + +func (a *Agent) GetAgentCheckInConfig() AgentCheckInConfig { + ret := AgentCheckInConfig{} + url := fmt.Sprintf("/api/v3/%s/config/", a.AgentID) + r, err := a.rClient.R().SetResult(&AgentCheckInConfig{}).Get(url) + if err != nil { + a.Logger.Debugln("GetAgentCheckInConfig()", err) + ret.Hello = randRange(30, 60) + ret.AgentInfo = randRange(200, 400) + ret.WinSvc = randRange(2400, 3000) + ret.PubIP = randRange(300, 500) + ret.Disks = randRange(1000, 2000) + ret.SW = randRange(2800, 3500) + ret.WMI = randRange(3000, 4000) + ret.SyncMesh = randRange(800, 1200) + ret.LimitData = false + return ret + } + ret.Hello = r.Result().(*AgentCheckInConfig).Hello + ret.AgentInfo = r.Result().(*AgentCheckInConfig).AgentInfo + ret.WinSvc = r.Result().(*AgentCheckInConfig).WinSvc + ret.PubIP = r.Result().(*AgentCheckInConfig).PubIP + ret.Disks = r.Result().(*AgentCheckInConfig).Disks + ret.SW = r.Result().(*AgentCheckInConfig).SW + ret.WMI = r.Result().(*AgentCheckInConfig).WMI + ret.SyncMesh = r.Result().(*AgentCheckInConfig).SyncMesh + ret.LimitData = r.Result().(*AgentCheckInConfig).LimitData + a.Logger.Debugf("%+v\n", r) + return ret +} diff --git a/agent/utils.go b/agent/utils.go index f5916f8..ede7753 100644 --- a/agent/utils.go +++ b/agent/utils.go @@ -322,3 +322,12 @@ func createTmpFile() (*os.File, error) { } return f, nil } + +func stringInSlice(a string, list []string) bool { + for _, b := range list { + if b == a { + return true + } + } + return false +}