refactor natsapi
This commit is contained in:
wh1te909
@@ -1,82 +0,0 @@
|
||||
package api
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"errors"
|
||||
"fmt"
|
||||
"log"
|
||||
"os"
|
||||
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/go-resty/resty/v2"
|
||||
nats "github.com/nats-io/nats.go"
|
||||
)
|
||||
|
||||
var rClient = resty.New()
|
||||
|
||||
func getAPI(apihost, natshost string) (string, string, error) {
|
||||
if apihost != "" && natshost != "" {
|
||||
return apihost, natshost, nil
|
||||
}
|
||||
|
||||
f, err := os.Open(`/etc/nginx/sites-available/rmm.conf`)
|
||||
if err != nil {
|
||||
return "", "", err
|
||||
}
|
||||
defer f.Close()
|
||||
|
||||
scanner := bufio.NewScanner(f)
|
||||
for scanner.Scan() {
|
||||
if strings.Contains(scanner.Text(), "server_name") && !strings.Contains(scanner.Text(), "301") {
|
||||
r := strings.NewReplacer("server_name", "", ";", "")
|
||||
s := strings.ReplaceAll(r.Replace(scanner.Text()), " ", "")
|
||||
return fmt.Sprintf("https://%s/natsapi", s), fmt.Sprintf("tls://%s:4222", s), nil
|
||||
}
|
||||
}
|
||||
return "", "", errors.New("unable to parse api from nginx conf")
|
||||
}
|
||||
|
||||
func Listen(apihost, natshost, version string, debug bool) {
|
||||
api, natsurl, err := getAPI(apihost, natshost)
|
||||
if err != nil {
|
||||
log.Fatalln(err)
|
||||
}
|
||||
|
||||
log.Printf("Tactical Nats API Version %s\n", version)
|
||||
log.Println("Api base url: ", api)
|
||||
log.Println("Nats connection url: ", natsurl)
|
||||
|
||||
rClient.SetHostURL(api)
|
||||
rClient.SetTimeout(10 * time.Second)
|
||||
natsinfo, err := rClient.R().SetResult(&NatsInfo{}).Get("/natsinfo/")
|
||||
if err != nil {
|
||||
log.Fatalln(err)
|
||||
}
|
||||
if natsinfo.IsError() {
|
||||
log.Fatalln(natsinfo.String())
|
||||
}
|
||||
|
||||
opts := []nats.Option{
|
||||
nats.Name("TacticalRMM"),
|
||||
nats.UserInfo(natsinfo.Result().(*NatsInfo).User,
|
||||
natsinfo.Result().(*NatsInfo).Password),
|
||||
nats.ReconnectWait(time.Second * 5),
|
||||
nats.RetryOnFailedConnect(true),
|
||||
nats.MaxReconnects(-1),
|
||||
nats.ReconnectBufSize(-1),
|
||||
}
|
||||
|
||||
nc, err := nats.Connect(natsurl, opts...)
|
||||
if err != nil {
|
||||
log.Fatalln(err)
|
||||
}
|
||||
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(1)
|
||||
go getWMI(rClient, nc)
|
||||
go monitorAgents(rClient, nc)
|
||||
wg.Wait()
|
||||
}
|
||||
Binary file not shown.
139
natsapi/tasks.go
139
natsapi/tasks.go
@@ -1,15 +1,42 @@
|
||||
package api
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"io/ioutil"
|
||||
"log"
|
||||
"math/rand"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/go-resty/resty/v2"
|
||||
nats "github.com/nats-io/nats.go"
|
||||
"github.com/ugorji/go/codec"
|
||||
)
|
||||
|
||||
func monitorAgents(c *resty.Client, nc *nats.Conn) {
|
||||
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
|
||||
@@ -22,56 +49,80 @@ func monitorAgents(c *resty.Client, nc *nats.Conn) {
|
||||
Data: map[string]string{"mode": "tacagent"},
|
||||
})
|
||||
|
||||
tick := time.NewTicker(7 * time.Minute)
|
||||
for range tick.C {
|
||||
var wg sync.WaitGroup
|
||||
agentids, _ := c.R().SetResult(&AgentIDS{}).Get("/offline/agents/")
|
||||
ids := agentids.Result().(*AgentIDS).IDs
|
||||
wg.Add(len(ids))
|
||||
var resp string
|
||||
|
||||
for _, id := range ids {
|
||||
go func(id string, nc *nats.Conn, wg *sync.WaitGroup, c *resty.Client) {
|
||||
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)
|
||||
p := map[string]string{"agentid": id}
|
||||
c.R().SetBody(p).Post("/logcrash/")
|
||||
}
|
||||
}
|
||||
}(id, nc, &wg, c)
|
||||
}
|
||||
wg.Wait()
|
||||
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(c *resty.Client, nc *nats.Conn) {
|
||||
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"})
|
||||
|
||||
tick := time.NewTicker(18 * time.Minute)
|
||||
for range tick.C {
|
||||
agentids, _ := c.R().SetResult(&AgentIDS{}).Get("/online/agents/")
|
||||
ids := agentids.Result().(*AgentIDS).IDs
|
||||
chunks := makeChunks(ids, 40)
|
||||
|
||||
for _, id := range chunks {
|
||||
for _, chunk := range id {
|
||||
nc.Publish(chunk, payload)
|
||||
time.Sleep(time.Duration(randRange(50, 400)) * time.Millisecond)
|
||||
}
|
||||
time.Sleep(15 * time.Second)
|
||||
}
|
||||
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
|
||||
}
|
||||
|
||||
@@ -1,15 +0,0 @@
|
||||
package api
|
||||
|
||||
type NatsInfo struct {
|
||||
User string `json:"user"`
|
||||
Password string `json:"password"`
|
||||
}
|
||||
|
||||
type AgentIDS struct {
|
||||
IDs []string `json:"agent_ids"`
|
||||
}
|
||||
|
||||
type Recovery struct {
|
||||
Func string `json:"func"`
|
||||
Data map[string]string `json:"payload"`
|
||||
}
|
||||
@@ -1,23 +0,0 @@
|
||||
package api
|
||||
|
||||
import (
|
||||
"math/rand"
|
||||
"time"
|
||||
)
|
||||
|
||||
func makeChunks(ids []string, chunkSize int) [][]string {
|
||||
var chunks [][]string
|
||||
for i := 0; i < len(ids); i += chunkSize {
|
||||
end := i + chunkSize
|
||||
if end > len(ids) {
|
||||
end = len(ids)
|
||||
}
|
||||
chunks = append(chunks, ids[i:end])
|
||||
}
|
||||
return chunks
|
||||
}
|
||||
|
||||
func randRange(min, max int) int {
|
||||
rand.Seed(time.Now().UnixNano())
|
||||
return rand.Intn(max-min) + min
|
||||
}
|
||||
Reference in New Issue
Block a user