129 lines
2.8 KiB
Go
129 lines
2.8 KiB
Go
package api
|
|
|
|
import (
|
|
"encoding/json"
|
|
"io/ioutil"
|
|
"log"
|
|
"math/rand"
|
|
"sync"
|
|
"time"
|
|
|
|
nats "github.com/nats-io/nats.go"
|
|
"github.com/ugorji/go/codec"
|
|
)
|
|
|
|
type JsonFile struct {
|
|
Agents []string `json:"agents"`
|
|
Key string `json:"key"`
|
|
NatsURL string `json:"natsurl"`
|
|
}
|
|
|
|
type Recovery struct {
|
|
Func string `json:"func"`
|
|
Data map[string]string `json:"payload"`
|
|
}
|
|
|
|
func setupNatsOptions(key string) []nats.Option {
|
|
opts := []nats.Option{
|
|
nats.Name("TacticalRMM"),
|
|
nats.UserInfo("tacticalrmm", key),
|
|
nats.ReconnectWait(time.Second * 2),
|
|
nats.RetryOnFailedConnect(true),
|
|
nats.MaxReconnects(3),
|
|
nats.ReconnectBufSize(-1),
|
|
}
|
|
return opts
|
|
}
|
|
|
|
func MonitorAgents(file string) {
|
|
var result JsonFile
|
|
var payload, recPayload []byte
|
|
var mh codec.MsgpackHandle
|
|
mh.RawToString = true
|
|
ret := codec.NewEncoderBytes(&payload, new(codec.MsgpackHandle))
|
|
ret.Encode(map[string]string{"func": "ping"})
|
|
|
|
rec := codec.NewEncoderBytes(&recPayload, new(codec.MsgpackHandle))
|
|
rec.Encode(Recovery{
|
|
Func: "recover",
|
|
Data: map[string]string{"mode": "tacagent"},
|
|
})
|
|
|
|
jret, _ := ioutil.ReadFile(file)
|
|
err := json.Unmarshal(jret, &result)
|
|
if err != nil {
|
|
log.Fatalln(err)
|
|
}
|
|
|
|
opts := setupNatsOptions(result.Key)
|
|
|
|
nc, err := nats.Connect(result.NatsURL, opts...)
|
|
if err != nil {
|
|
log.Fatalln(err)
|
|
}
|
|
defer nc.Close()
|
|
|
|
var wg sync.WaitGroup
|
|
var resp string
|
|
wg.Add(len(result.Agents))
|
|
|
|
for _, id := range result.Agents {
|
|
go func(id string, nc *nats.Conn, wg *sync.WaitGroup) {
|
|
defer wg.Done()
|
|
out, err := nc.Request(id, payload, 1*time.Second)
|
|
if err != nil {
|
|
return
|
|
}
|
|
dec := codec.NewDecoderBytes(out.Data, &mh)
|
|
if err := dec.Decode(&resp); err == nil {
|
|
// if the agent is respoding to pong from the rpc service but is not showing as online (handled by tacticalagent service)
|
|
// then tacticalagent service is hung. forcefully restart it
|
|
if resp == "pong" {
|
|
nc.Publish(id, recPayload)
|
|
}
|
|
}
|
|
}(id, nc, &wg)
|
|
}
|
|
wg.Wait()
|
|
}
|
|
|
|
func GetWMI(file string) {
|
|
var result JsonFile
|
|
var payload []byte
|
|
var mh codec.MsgpackHandle
|
|
mh.RawToString = true
|
|
ret := codec.NewEncoderBytes(&payload, new(codec.MsgpackHandle))
|
|
ret.Encode(map[string]string{"func": "wmi"})
|
|
|
|
jret, _ := ioutil.ReadFile(file)
|
|
err := json.Unmarshal(jret, &result)
|
|
if err != nil {
|
|
log.Fatalln(err)
|
|
}
|
|
|
|
opts := setupNatsOptions(result.Key)
|
|
|
|
nc, err := nats.Connect(result.NatsURL, opts...)
|
|
if err != nil {
|
|
log.Fatalln(err)
|
|
}
|
|
defer nc.Close()
|
|
|
|
var wg sync.WaitGroup
|
|
wg.Add(len(result.Agents))
|
|
|
|
for _, id := range result.Agents {
|
|
go func(id string, nc *nats.Conn, wg *sync.WaitGroup) {
|
|
defer wg.Done()
|
|
time.Sleep(time.Duration(randRange(0, 20)) * time.Second)
|
|
nc.Publish(id, payload)
|
|
}(id, nc, &wg)
|
|
}
|
|
wg.Wait()
|
|
}
|
|
|
|
func randRange(min, max int) int {
|
|
rand.Seed(time.Now().UnixNano())
|
|
return rand.Intn(max-min) + min
|
|
}
|