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