diff --git a/agents/agent.go b/agents/agent.go index 2f04dd3..9ea60f1 100644 --- a/agents/agent.go +++ b/agents/agent.go @@ -4,13 +4,15 @@ import ( "encoding/json" "fmt" "log" + "math" + "math/rand" "net/http" "net/url" + "os" "os/exec" - "time" - "math/rand" - "math" "strconv" + "time" + // "gontrol/src/logger" "github.com/gorilla/websocket" @@ -30,6 +32,7 @@ type Agent struct { AgentIP string `json:"agentIp"` InitialContact string `json:"initialContact"` LastContact string `json:"lastContact"` + HostName string `json:"hostName"` } type Message struct { @@ -63,8 +66,8 @@ func registerAgent(agentName, agentId, agentIp, agentType, addPort string) error return nil } -func connectToWebSocket(agentName, agentId, agentIp, agentType string) error { - wsURL := fmt.Sprintf("ws://%s/data?agentName=%s&agentId=%s&IPv4Address=%s&agentType=%s", webSocketAddr, url.QueryEscape(agentName), url.QueryEscape(agentId), url.QueryEscape(agentIp), url.QueryEscape(agentType)) +func connectToWebSocket(agentName, agentId, agentIp, agentType, hostName string) error { + wsURL := fmt.Sprintf("ws://%s/data?agentName=%s&agentId=%s&IPv4Address=%s&agentType=%s&hostname=%s", webSocketAddr, url.QueryEscape(agentName), url.QueryEscape(agentId), url.QueryEscape(agentIp), url.QueryEscape(agentType), url.QueryEscape(hostName)) var err error for { conn, _, err = websocket.DefaultDialer.Dial(wsURL, nil) @@ -79,13 +82,13 @@ func connectToWebSocket(agentName, agentId, agentIp, agentType string) error { } } -func reconnectToWebSocket(agentName, agentId, agentIp, agentType string) error { +func reconnectToWebSocket(agentName, agentId, agentIp, agentType, hostName string) error { backoff := 2 * time.Second maxBackoff := 1 * time.Minute for { log.Println("Attempting to reconnect to WebSocket...") - err := connectToWebSocket(agentName, agentId, agentIp, agentType) + err := connectToWebSocket(agentName, agentId, agentIp, agentType, hostName) if err == nil { log.Println("Reconnection succesful.") return nil @@ -101,14 +104,14 @@ func reconnectToWebSocket(agentName, agentId, agentIp, agentType string) error { } } -func listenForCommands(agentName, agentId, agentIp, agentType string) { +func listenForCommands(agentName, agentId, agentIp, agentType, hostName string) { defer conn.Close() for { _, rawMessage, err := conn.ReadMessage() if err != nil { log.Printf("Connection lost: %v", err) - if reconnectErr := reconnectToWebSocket(agentName, agentId, agentIp, agentType); reconnectErr != nil { + if reconnectErr := reconnectToWebSocket(agentName, agentId, agentIp, agentType, hostName); reconnectErr != nil { log.Printf("Critical error during reconnection: %v", reconnectErr) } continue @@ -163,6 +166,7 @@ func main() { agentId := strconv.Itoa(randomInt(5)) agentIp := "127.0.0.1" agentType := "BaseAgent" + hostName, _ := os.Hostname() log.Printf("AgentId: %s", agentId) @@ -170,9 +174,9 @@ func main() { // log.Fatalf("Agent registration failed: %v", err) // } - if err := connectToWebSocket(agentName, agentId, agentIp, agentType); err != nil { + if err := connectToWebSocket(agentName, agentId, agentIp, agentType, hostName); err != nil { log.Fatalf("Websocket connection failed: %v", err) } - listenForCommands(agentName, agentId, agentIp, agentType) + listenForCommands(agentName, agentId, agentIp, agentType, hostName) } diff --git a/gomatic.sql b/gomatic.sql index dd7c6f3..df94bc7 100644 --- a/gomatic.sql +++ b/gomatic.sql @@ -10,7 +10,8 @@ create table agents ( initialContact timestamp, lastContact timestamp, status Boolean, - addPort varchar(255) + addPort varchar(255), + hostName varchar(255), ); insert into agents (IPv4Address, agentId, agentName, initialContact, lastContact) values ( '127.0.0.1', 'testAgent', NOW(), NOW()); diff --git a/src/server/api/agentApi.go b/src/server/api/agentApi.go index a3541f8..a2713d5 100644 --- a/src/server/api/agentApi.go +++ b/src/server/api/agentApi.go @@ -20,6 +20,7 @@ type Agent struct { IPv4Address string `json:"IPv4Address"` Status string `json:"status"` AddPort string `json:"addPort"` + HostName string `json:"HostName"` } // var db *sql.DB @@ -52,12 +53,13 @@ func CreateAgent(db *sql.DB, w http.ResponseWriter, r * http.Request) (http.Resp agentType := r.FormValue("agentType") IPv4Address := r.FormValue("IPv4Address") addPort :=r.FormValue("addPort") + hostname := r.FormValue("hostname") // initalContact := r.FormValue("initialContact") // lastContact := r.FormValue("lastContact") // query := "INSERT INTO agents (agentId, agentName, agentType, IPv4Address, addPort, initialContact, lastContact) VALUES (?, ?, ?, ?, ?, NOW(), NOW())" - query := "INSERT INTO agents (agentId, agentName, agentType, IPv4Address, addPort, initialContact, lastContact) VALUES (?, ?, ?, ?, ?, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)" - _, err = db.Exec(query, agentId, agentName, agentType, IPv4Address, addPort) + query := "INSERT INTO agents (agentId, agentName, agentType, IPv4Address, addPort, hostName, initialContact, lastContact) VALUES (?, ?, ?, ?, ?, ?, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)" + _, err = db.Exec(query, agentId, agentName, agentType, IPv4Address, addPort, hostname) if err != nil { log.Printf("Database err is: %s", err) http.Error(w, "Failed to create agent", http.StatusInternalServerError) @@ -90,7 +92,7 @@ func UpdateAgent(db *sql.DB, w http.ResponseWriter, r *http.Request, agentId str } func GetAgents(db *sql.DB) ([]Agent, error) { - query := "SELECT agentId, agentName, agentType, IPv4Address, addPort, initialContact, lastContact FROM agents" + query := "SELECT agentId, agentName, agentType, IPv4Address, addPort, hostName, initialContact, lastContact FROM agents" rows, err := db.Query(query) if err != nil { return nil, err @@ -100,7 +102,7 @@ func GetAgents(db *sql.DB) ([]Agent, error) { var agents []Agent for rows.Next() { var agent Agent - err := rows.Scan(&agent.AgentID, &agent.AgentName, &agent.AgentType, &agent.IPv4Address, &agent.AddPort, &agent.InitialContact, &agent.LastContact) + err := rows.Scan(&agent.AgentID, &agent.AgentName, &agent.AgentType, &agent.IPv4Address, &agent.AddPort, &agent.HostName, &agent.InitialContact, &agent.LastContact) if err != nil { return nil, err } @@ -110,9 +112,9 @@ func GetAgents(db *sql.DB) ([]Agent, error) { } func GetAgent(db *sql.DB, w http.ResponseWriter, r *http.Request, agentId string) (Agent, error) { - query := "Select agentId, agentName, agentType, IPv4Address, addPort, initialContact, lastContact from agents where agentId = ?" + query := "Select agentId, agentName, agentType, IPv4Address, addPort, hostName, initialContact, lastContact from agents where agentId = ?" var agent Agent - err := db.QueryRow(query, agentId).Scan(&agent.AgentID, &agent.AgentName, &agent.AgentType, &agent.IPv4Address, &agent.AddPort, &agent.InitialContact, &agent.LastContact) + err := db.QueryRow(query, agentId).Scan(&agent.AgentID, &agent.AgentName, &agent.AgentType, &agent.IPv4Address, &agent.AddPort, &agent.HostName, &agent.InitialContact, &agent.LastContact) if err == sql.ErrNoRows { http.Error(w, "Agent not found", http.StatusNotFound) return Agent{} , err diff --git a/src/server/database/database.go b/src/server/database/database.go index aff92a0..bc06e85 100644 --- a/src/server/database/database.go +++ b/src/server/database/database.go @@ -55,7 +55,8 @@ func InitSQLiteDB(dbPath string) *sql.DB { IPv4Address TEXT, initialContact datetime, lastContact datetime, - addPort TEXT + addPort TEXT, + hostname TEXT );` _, err = Agent_db.Exec(CreateTableQuery) diff --git a/src/server/websocket/websocketServer.go b/src/server/websocket/websocketServer.go index c612165..6aa8a21 100644 --- a/src/server/websocket/websocketServer.go +++ b/src/server/websocket/websocketServer.go @@ -47,7 +47,7 @@ var getAgentNames http.HandlerFunc = func(w http.ResponseWriter, r *http.Request } -func registerAgent(agentName, agentId, agentIp, agentType, addPort string) error { +func registerAgent(agentName, agentId, agentIp, agentType, addPort, hostname string) error { registerURL := "http://localhost:3333/agents" @@ -57,6 +57,7 @@ func registerAgent(agentName, agentId, agentIp, agentType, addPort string) error form.Add("agentType", agentType) form.Add("IPv4Address", agentIp) form.Add("addPort", addPort) + form.Add("hostname", hostname) resp, err := http.PostForm(registerURL, form) if err != nil { @@ -175,6 +176,7 @@ func (wsh webSocketHandler) ServeHTTP(w http.ResponseWriter, r *http.Request){ agentIP := r.URL.Query().Get("IPv4Address") agentId := r.URL.Query().Get("agentId") agentType := r.URL.Query().Get("agentType") + hostName := r.URL.Query().Get("hostname") agentName := "" log.Println("addPort is the following") log.Println(r.URL.Query().Get("addPort")) @@ -194,7 +196,7 @@ func (wsh webSocketHandler) ServeHTTP(w http.ResponseWriter, r *http.Request){ if !containsId(agentIds, agentId) { agentName = randomname.GenerateRandomName() - registerAgent(agentName, agentId, agentIP, agentType, addPort) + registerAgent(agentName, agentId, agentIP, agentType, addPort, hostName) } else { agentDetails, _ := getAgentDetails(agentId) agentName = agentDetails.AgentName diff --git a/templates/partials/agent_detail.html b/templates/partials/agent_detail.html index 7dfb50c..5cd282a 100644 --- a/templates/partials/agent_detail.html +++ b/templates/partials/agent_detail.html @@ -7,6 +7,7 @@
Interactive Port: {{.AddPort}}
Initial Contact: {{.InitialContact}}
Last Contact: {{.LastContact}}
+Hostname: {{.HostName}}
Open diff --git a/templates/partials/agent_list.html b/templates/partials/agent_list.html index 8060221..0696a22 100644 --- a/templates/partials/agent_list.html +++ b/templates/partials/agent_list.html @@ -4,7 +4,8 @@