package api import ( "encoding/json" "reflect" "runtime" "time" _ "github.com/lib/pq" nats "github.com/nats-io/nats.go" "github.com/sirupsen/logrus" "github.com/ugorji/go/codec" trmm "github.com/wh1te909/trmm-shared" ) func Svc(logger *logrus.Logger, cfg string) { logger.Debugln("Starting Svc()") db, r, err := GetConfig(cfg) if err != nil { logger.Fatalln(err) } opts := setupNatsOptions(r.Key) nc, err := nats.Connect(r.NatsURL, opts...) if err != nil { logger.Fatalln(err) } nc.Subscribe("*", func(msg *nats.Msg) { var mh codec.MsgpackHandle mh.MapType = reflect.TypeOf(map[string]interface{}(nil)) mh.RawToString = true dec := codec.NewDecoderBytes(msg.Data, &mh) switch msg.Reply { case "agent-hello": go func() { var p trmm.CheckInNats if err := dec.Decode(&p); err == nil { loc, _ := time.LoadLocation("UTC") now := time.Now().In(loc) logger.Debugln("Hello", p, now) stmt := ` UPDATE agents_agent SET last_seen=$1, version=$2 WHERE agents_agent.agent_id=$3; ` _, err = db.Exec(stmt, now, p.Version, p.Agentid) if err != nil { logger.Errorln(err) } } }() case "agent-publicip": go func() { var p trmm.PublicIPNats if err := dec.Decode(&p); err == nil { logger.Debugln("Public IP", p) stmt := ` UPDATE agents_agent SET public_ip=$1 WHERE agents_agent.agent_id=$2;` _, err = db.Exec(stmt, p.PublicIP, p.Agentid) if err != nil { logger.Errorln(err) } } }() case "agent-agentinfo": go func() { var r trmm.AgentInfoNats if err := dec.Decode(&r); err == nil { stmt := ` UPDATE agents_agent SET hostname=$1, operating_system=$2, plat=$3, total_ram=$4, boot_time=$5, needs_reboot=$6, logged_in_username=$7, goarch=$8 WHERE agents_agent.agent_id=$9;` logger.Debugln("Info", r) _, err = db.Exec(stmt, r.Hostname, r.OS, r.Platform, r.TotalRAM, r.BootTime, r.RebootNeeded, r.Username, r.GoArch, r.Agentid) if err != nil { logger.Errorln(err) } if r.Username != "None" { stmt = `UPDATE agents_agent SET last_logged_in_user=$1 WHERE agents_agent.agent_id=$2;` logger.Debugln("Updating last logged in user:", r.Username) _, err = db.Exec(stmt, r.Username, r.Agentid) if err != nil { logger.Errorln(err) } } } }() case "agent-disks": go func() { var r trmm.WinDisksNats if err := dec.Decode(&r); err == nil { logger.Debugln("Disks", r) b, err := json.Marshal(r.Disks) if err != nil { logger.Errorln(err) return } stmt := ` UPDATE agents_agent SET disks=$1 WHERE agents_agent.agent_id=$2;` _, err = db.Exec(stmt, b, r.Agentid) if err != nil { logger.Errorln(err) } } }() case "agent-winsvc": go func() { var r trmm.WinSvcNats if err := dec.Decode(&r); err == nil { logger.Debugln("WinSvc", r) b, err := json.Marshal(r.WinSvcs) if err != nil { logger.Errorln(err) return } stmt := ` UPDATE agents_agent SET services=$1 WHERE agents_agent.agent_id=$2;` _, err = db.Exec(stmt, b, r.Agentid) if err != nil { logger.Errorln(err) } } }() case "agent-wmi": go func() { var r trmm.WinWMINats if err := dec.Decode(&r); err == nil { logger.Debugln("WMI", r) b, err := json.Marshal(r.WMI) if err != nil { logger.Errorln(err) return } stmt := ` UPDATE agents_agent SET wmi_detail=$1 WHERE agents_agent.agent_id=$2;` _, err = db.Exec(stmt, b, r.Agentid) if err != nil { logger.Errorln(err) } } }() } }) nc.Flush() if err := nc.LastError(); err != nil { logger.Fatalln(err) } runtime.Goexit() }