added hostname
This commit is contained in:
parent
ca7da7fc62
commit
4f1ee3ad01
|
@ -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)
|
||||
}
|
||||
|
|
|
@ -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());
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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)
|
||||
|
|
|
@ -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
|
||||
|
|
|
@ -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>
|
||||
|
|
|
@ -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> -->
|
||||
|
|
Loading…
Reference in New Issue