added hostname

This commit is contained in:
Stefan Etringer 2025-05-26 15:22:07 +00:00
parent ca7da7fc62
commit 4f1ee3ad01
7 changed files with 36 additions and 23 deletions

View File

@ -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)
}

View File

@ -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());

View File

@ -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

View File

@ -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)

View File

@ -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

View File

@ -7,6 +7,7 @@
<p>Interactive Port: {{.AddPort}}</p>
<p>Initial Contact: {{.InitialContact}}</p>
<p>Last Contact: {{.LastContact}}</p>
<p>Hostname: {{.HostName}}</p>
<!-- <button hx-get="/proxyAgent?ip={{.IPv4Address}}" hx-target="#agentConnect" hx-swap="innerHTML">Open</button> -->
<a href="http://{{.IPv4Address}}:{{.AddPort}}" class="btn btn-info" target="_blank">Open</a>
</div>

View File

@ -4,7 +4,8 @@
<th>ID</th>
<th>Name</th>
<th>Type</th>
<th>IPv4 Address</th>
<th>Hostname</th>
<!-- <th>IPv4 Address</th> -->
<!-- <th>Initial Contact</th> -->
<!-- <th>Last Contact</th> -->
<th>Status</th>
@ -18,7 +19,8 @@
<td>{{.AgentID}}</td>
<td>{{.AgentName}}</td>
<td>{{.AgentType}}</td>
<td>{{.IPv4Address}}</td>
<!-- <td>{{.IPv4Address}}</td> -->
<td>{{.HostName}}</td>
<!-- <td>{{.InitialContact}}</td> -->
<!-- <td>{{.LastContact}}</td> -->
<!-- <td><span class="badge bg-success">Connected</span></td> -->