move checkin to go

natsapi/tasks.go

@@ -2,12 +2,15 @@ package api
 
 import (
 	"encoding/json"
+	"fmt"
 	"io/ioutil"
 	"log"
 	"math/rand"
 	"sync"
 	"time"
 
+	"github.com/jmoiron/sqlx"
+	_ "github.com/lib/pq"
 	nats "github.com/nats-io/nats.go"
 	"github.com/ugorji/go/codec"
 )
@@ -18,6 +21,21 @@ type JsonFile struct {
 	NatsURL string   `json:"natsurl"`
 }
 
+type DjangoConfig struct {
+	Key     string `json:"key"`
+	NatsURL string `json:"natsurl"`
+	User    string `json:"user"`
+	Pass    string `json:"pass"`
+	Host    string `json:"host"`
+	Port    int    `json:"port"`
+	DBName  string `json:"dbname"`
+}
+
+type Agent struct {
+	ID      int    `db:"id"`
+	AgentID string `db:"agent_id"`
+}
+
 type Recovery struct {
 	Func string            `json:"func"`
 	Data map[string]string `json:"payload"`
@@ -87,6 +105,90 @@ func MonitorAgents(file string) {
 	wg.Wait()
 }
 
+func CheckIn(file string) {
+	var r DjangoConfig
+
+	jret, _ := ioutil.ReadFile(file)
+	err := json.Unmarshal(jret, &r)
+	if err != nil {
+		log.Fatalln(err)
+	}
+
+	psqlInfo := fmt.Sprintf("host=%s port=%d user=%s "+
+		"password=%s dbname=%s sslmode=disable",
+		r.Host, r.Port, r.User, r.Pass, r.DBName)
+
+	db, err := sqlx.Connect("postgres", psqlInfo)
+	if err != nil {
+		log.Fatalln(err)
+	}
+	db.SetMaxOpenConns(15)
+
+	agent := Agent{}
+	agents := make([]Agent, 0)
+	rows, err := db.Queryx("SELECT agents_agent.id, agents_agent.agent_id FROM agents_agent")
+	if err != nil {
+		log.Fatalln(err)
+	}
+
+	for rows.Next() {
+		err := rows.StructScan(&agent)
+		if err != nil {
+			continue
+		}
+		agents = append(agents, agent)
+	}
+
+	var payload []byte
+	ret := codec.NewEncoderBytes(&payload, new(codec.MsgpackHandle))
+	ret.Encode(map[string]string{"func": "ping"})
+
+	opts := setupNatsOptions(r.Key)
+
+	nc, err := nats.Connect(r.NatsURL, opts...)
+	if err != nil {
+		log.Fatalln(err)
+	}
+	defer nc.Close()
+
+	var wg sync.WaitGroup
+	wg.Add(len(agents))
+
+	loc, _ := time.LoadLocation("UTC")
+	now := time.Now().In(loc)
+
+	for _, a := range agents {
+		go func(id string, pk int, nc *nats.Conn, wg *sync.WaitGroup, db *sqlx.DB, now time.Time) {
+			defer wg.Done()
+
+			var resp string
+			var mh codec.MsgpackHandle
+			mh.RawToString = true
+
+			time.Sleep(time.Duration(randRange(100, 1500)) * time.Millisecond)
+			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 resp == "pong" {
+					_, err = db.NamedExec(
+						`UPDATE agents_agent SET last_seen=:lastSeen WHERE agents_agent.id=:pk`,
+						map[string]interface{}{"lastSeen": now, "pk": pk},
+					)
+					if err != nil {
+						fmt.Println(err)
+					}
+				}
+			}
+		}(a.AgentID, a.ID, nc, &wg, db, now)
+	}
+	wg.Wait()
+	db.Close()
+}
+
 func GetWMI(file string) {
 	var result JsonFile
 	var payload []byte
@@ -115,7 +217,7 @@ func GetWMI(file string) {
 	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)
+			time.Sleep(time.Duration(randRange(0, 28)) * time.Second)
 			nc.Publish(id, payload)
 		}(id, nc, &wg)
 	}