A bit more refactor to clean up the code.

This commit is contained in:
James Wells 2021-11-06 20:18:49 -07:00
parent 7e51584832
commit de41740bac
Signed by: jwells
GPG key ID: 73196D10B8E65666
5 changed files with 40 additions and 35 deletions

View file

@ -1,9 +1,7 @@
package main package main
import ( import (
sha "crypto/sha256"
fmt "fmt" fmt "fmt"
ioutil "io/ioutil"
log "log" log "log"
http "net/http" http "net/http"
os "os" os "os"
@ -13,7 +11,7 @@ import (
autorestart "github.com/slayer/autorestart" autorestart "github.com/slayer/autorestart"
gstruct "git.dragonheim.net/dragonheim/gagent/internal/gstructs" gstructs "git.dragonheim.net/dragonheim/gagent/internal/gstructs"
gc "git.dragonheim.net/dragonheim/gagent/internal/client" gc "git.dragonheim.net/dragonheim/gagent/internal/client"
gr "git.dragonheim.net/dragonheim/gagent/internal/router" gr "git.dragonheim.net/dragonheim/gagent/internal/router"
@ -43,8 +41,8 @@ var exitCodes = struct {
m map[string]int m map[string]int
}{m: map[string]int{}} }{m: map[string]int{}}
var config gstruct.GagentConfig var config gstructs.GagentConfig
var agent gstruct.AgentDetails var agent gstructs.AgentDetails
func main() { func main() {
filter := &logutils.LevelFilter{ filter := &logutils.LevelFilter{
@ -65,20 +63,8 @@ func main() {
os.Exit(exitCodes.m["NO_ROUTERS_DEFINED"]) os.Exit(exitCodes.m["NO_ROUTERS_DEFINED"])
} }
var err error
if config.CMode {
agent.ScriptCode, err = ioutil.ReadFile(config.File)
if err != nil {
log.Printf("[ERROR] No such file or directory: %s", config.File)
os.Exit(exitCodes.m["AGENT_LOAD_FAILED"])
}
agent.Shasum = fmt.Sprintf("%x", sha.Sum256(agent.ScriptCode))
agent.Status = "loaded"
log.Printf("[DEBUG] SHA256 of Agent file: %s", agent.Shasum)
}
wg.Add(1) wg.Add(1)
go gc.Main(&wg, config, string(agent.ScriptCode)) go gc.Main(&wg, config)
case "router": case "router":
log.Printf("[INFO] Running in router mode\n") log.Printf("[INFO] Running in router mode\n")
@ -103,6 +89,8 @@ func main() {
go gw.Main(&wg, config) go gw.Main(&wg, config)
case "setup": case "setup":
log.Printf("[INFO] Running in setup mode\n")
wg.Add(1) wg.Add(1)
go gs.Main(&wg, config) go gs.Main(&wg, config)
@ -185,9 +173,9 @@ func init() {
*/ */
config.WorkerPort = 35572 config.WorkerPort = 35572
config.Clients = make([]*gstruct.ClientDetails, 0) config.Clients = make([]*gstructs.ClientDetails, 0)
config.Routers = make([]*gstruct.RouterDetails, 0) config.Routers = make([]*gstructs.RouterDetails, 0)
config.Workers = make([]*gstruct.WorkerDetails, 0) config.Workers = make([]*gstructs.WorkerDetails, 0)
/* /*
* Create a usage variable and then use that to declare the arguments and * Create a usage variable and then use that to declare the arguments and

View file

@ -1,12 +1,15 @@
package client package client
import ( import (
sha "crypto/sha256"
fmt "fmt" fmt "fmt"
ioutil "io/ioutil"
log "log" log "log"
os "os"
sync "sync" sync "sync"
time "time" time "time"
gs "git.dragonheim.net/dragonheim/gagent/internal/gstructs" gstructs "git.dragonheim.net/dragonheim/gagent/internal/gstructs"
zmq "github.com/pebbe/zmq4" zmq "github.com/pebbe/zmq4"
) )
@ -18,9 +21,23 @@ import (
will contact the router and attempt to retrieve the results will contact the router and attempt to retrieve the results
of it's most recent request. of it's most recent request.
*/ */
func Main(wg *sync.WaitGroup, config gs.GagentConfig, agent string) { func Main(wg *sync.WaitGroup, config gstructs.GagentConfig) {
defer wg.Done()
log.Printf("[INFO] Starting client\n") log.Printf("[INFO] Starting client\n")
defer wg.Done()
var agent gstructs.AgentDetails
var err error
if config.CMode {
agent.ScriptCode, err = ioutil.ReadFile(config.File)
if err != nil {
log.Printf("[ERROR] No such file or directory: %s", config.File)
os.Exit(6)
}
agent.Shasum = fmt.Sprintf("%x", sha.Sum256(agent.ScriptCode))
agent.Status = "loaded"
log.Printf("[DEBUG] SHA256 of Agent file: %s", agent.Shasum)
}
for key := range config.Routers { for key := range config.Routers {
/* /*
@ -33,14 +50,15 @@ func Main(wg *sync.WaitGroup, config gs.GagentConfig, agent string) {
connectString := fmt.Sprintf("tcp://%s:%d", config.Routers[key].RouterAddr, rport) connectString := fmt.Sprintf("tcp://%s:%d", config.Routers[key].RouterAddr, rport)
wg.Add(1) wg.Add(1)
go sendAgent(wg, config.UUID, connectString, agent) go sendAgent(wg, config.UUID, connectString, agent.ScriptCode)
time.Sleep(10 * time.Millisecond) time.Sleep(10 * time.Millisecond)
} }
} }
func sendAgent(wg *sync.WaitGroup, uuid string, connectString string, agent string) { func sendAgent(wg *sync.WaitGroup, uuid string, connectString string, agent []byte) {
log.Printf("[DEBUG] Attempting to connect to %s\n", connectString) log.Printf("[DEBUG] Attempting to connect to %s\n", connectString)
defer wg.Done() defer wg.Done()
var mu sync.Mutex var mu sync.Mutex
mu.Lock() mu.Lock()

View file

@ -6,7 +6,7 @@ import (
http "net/http" http "net/http"
sync "sync" sync "sync"
gs "git.dragonheim.net/dragonheim/gagent/internal/gstructs" gstructs "git.dragonheim.net/dragonheim/gagent/internal/gstructs"
prometheus "github.com/prometheus/client_golang/prometheus" prometheus "github.com/prometheus/client_golang/prometheus"
promauto "github.com/prometheus/client_golang/prometheus/promauto" promauto "github.com/prometheus/client_golang/prometheus/promauto"
@ -27,12 +27,11 @@ var (
or client node. Tags are used by the agent to give hints as to where or client node. Tags are used by the agent to give hints as to where
it should be routed. it should be routed.
*/ */
func Main(wg *sync.WaitGroup, config gs.GagentConfig) { func Main(wg *sync.WaitGroup, config gstructs.GagentConfig) {
defer wg.Done()
http.HandleFunc("/hello", answerClient)
log.Printf("[INFO] Starting router\n") log.Printf("[INFO] Starting router\n")
defer wg.Done()
http.HandleFunc("/hello", answerClient)
clientSock, _ := zmq.NewSocket(zmq.ROUTER) clientSock, _ := zmq.NewSocket(zmq.ROUTER)
defer clientSock.Close() defer clientSock.Close()

View file

@ -12,7 +12,7 @@ import (
) )
func Main(wg *sync.WaitGroup, config gs.GagentConfig) { func Main(wg *sync.WaitGroup, config gs.GagentConfig) {
log.Printf("[INFO] Running in setup mode\n") log.Printf("[INFO] setup mode\n")
defer wg.Done() defer wg.Done()
f := hclwrite.NewEmptyFile() f := hclwrite.NewEmptyFile()

View file

@ -5,7 +5,7 @@ import (
log "log" log "log"
sync "sync" sync "sync"
gs "git.dragonheim.net/dragonheim/gagent/internal/gstructs" gstructs "git.dragonheim.net/dragonheim/gagent/internal/gstructs"
// picol "git.dragonheim.net/dragonheim/gagent/src/picol" // picol "git.dragonheim.net/dragonheim/gagent/src/picol"
@ -27,9 +27,9 @@ var (
router(s) they are connected. The worker will execute the agent code and router(s) they are connected. The worker will execute the agent code and
pass the agent and it's results to a router. pass the agent and it's results to a router.
*/ */
func Main(wg *sync.WaitGroup, config gs.GagentConfig) { func Main(wg *sync.WaitGroup, config gstructs.GagentConfig) {
defer wg.Done()
log.Printf("[INFO] Starting worker\n") log.Printf("[INFO] Starting worker\n")
defer wg.Done()
for key := range config.Routers { for key := range config.Routers {
rport := config.WorkerPort rport := config.WorkerPort