From 205647922475f3cd0f25436522deab85c67210f4 Mon Sep 17 00:00:00 2001 From: Stefan Etringer Date: Tue, 6 May 2025 15:46:43 +0000 Subject: [PATCH] bugfixing --- main.go | 105 +++++++---- src/logger/logger.go | 15 +- src/server/websocket/websocketServer.go | 45 +++-- templates/index.html | 221 +++++++++++++++++------- templates/logs.html | 60 ++++++- templates/partials/agent_list.html | 2 +- templates/partials/logs_partial.html | 5 + 7 files changed, 326 insertions(+), 127 deletions(-) create mode 100644 templates/partials/logs_partial.html diff --git a/main.go b/main.go index 765df01..895b791 100644 --- a/main.go +++ b/main.go @@ -2,22 +2,23 @@ package main import ( "context" + "encoding/json" "os" "strings" "time" - "fmt" "database/sql" + "fmt" "html/template" "log" "net/http" "sync" "syscall" - "gontrol/src/randomname" - "gontrol/src/server/database" "gontrol/src/logger" + "gontrol/src/randomname" api "gontrol/src/server/api" + "gontrol/src/server/database" websocketserver "gontrol/src/server/websocket" "os/signal" @@ -31,6 +32,8 @@ var ( db *sql.DB ) + + type Config struct { Database struct { Username string `envconfig:"DB_USERNAME"` @@ -55,22 +58,14 @@ func processError(err error) { func init() { tmpl, _ = template.ParseGlob("templates/*.html") + + // Sqlite3 + err := logger.InitDB("/tmp/gontrol_logs.db") + if err != nil { + log.Fatal(err) + } } -// func initDB (dbUser string, dbPassword string, dbHost string, dbPort int16, dbName string ) { -// var err error -// dbOpen := fmt.Sprintf("%s:%s@(%s:%d)/%s?parseTime=true", dbUser,dbPassword, dbHost, dbPort, dbName) -// db, err = sql.Open("mysql", dbOpen) -// if err != nil { -// log.Fatal(err) -// } - -// if err = db.Ping(); err != nil { -// log.Fatal(err) -// } -// log.Println("Connected to database") -// } - func renderTemplate(w http.ResponseWriter, tmpl string, data interface{}) { t, err := template.ParseFiles(tmpl) if err != nil { @@ -119,7 +114,24 @@ func getHomepage(w http.ResponseWriter, r *http.Request) { } func listAgents(w http.ResponseWriter, r *http.Request) { + var agents []api.Agent agents, err := api.GetAgents(db) + currentAgents := getAgentsStatus() + + for _, currAgent := range currentAgents { + for i, agent := range agents { + if currAgent == agent.AgentName { + // log.Printf("%s online", agent.AgentName) + // logger.InsertLog(logger.Debug, fmt.Sprintf("%s online after page refresh", agent.AgentName)) + // agents[i].Status = fmt.Sprint("Connected") + agents[i].Status = "Connected" + } else { + // agent.Status = fmt.Sprintf("Disconnected") + agents[i].Status = "Disconnected" + } + } + } + if err != nil { http.Error(w, "Failed to fetch agents", http.StatusInternalServerError) return @@ -128,6 +140,23 @@ func listAgents(w http.ResponseWriter, r *http.Request) { renderTemplate(w, "templates/partials/agent_list.html", agents) } +func getAgentsStatus() []string { + resp, err := http.Get("http://localhost:5555/agentNames") + if err != nil { + log.Println("Error fetching agent names:", err) + logger.InsertLog(logger.Error, "Error fetching agent names from websocketServer") + } + defer resp.Body.Close() + + var agentNames []string + if err := json.NewDecoder(resp.Body).Decode(&agentNames); err != nil { + log.Println("Error decoding response:", err) + return []string{} + } + + return agentNames +} + func getAgentNames(w http.ResponseWriter, r *http.Request) { api.GetAgentNames(db, w, r) return @@ -139,25 +168,38 @@ func getAgentIds(w http.ResponseWriter, r *http.Request) { } func logsHandler(w http.ResponseWriter, r *http.Request) { - + var logs[] logger.LogEntry logs, err := logger.FetchLogs(10) + + // for i, logLine := range logs { + // logs[i].Message = strings.ReplaceAll(logLine.Message, `"`, `\"`) + // } + if err != nil { http.Error(w, "Error fetching logs", http.StatusInternalServerError) return } - fmt.Fprintf(w, "
") - // for _, logEntry := range logsToSend { - for _, logEntry := range logs { - fmt.Fprintf(w, "

[%s] [%s] %s

", logEntry.Timestamp, logEntry.Level, logEntry.Message) - } - fmt.Fprintf(w, "
") + renderTemplate(w, "templates/partials/logs_partial.html", logs) + +// fmt.Fprintf(w, "
") +// // for _, logEntry := range logsToSend { +// for _, logEntry := range logs { +// fmt.Fprintf(w, "

[%s] [%s] %s

", logEntry.Timestamp, logEntry.Level, logEntry.Message) +// } +// fmt.Fprintf(w, "
") + + // w.Header().Set("Content-Type", "application/json") + // json.NewEncoder(w).Encode(logs) } func main() { + // sqlite3 has been initialized in init() + defer logger.CloseDB() + var cfg Config readEnv(&cfg) @@ -176,18 +218,11 @@ func main() { webMux.HandleFunc("/agents/{agentId}", agentsHandler) webMux.HandleFunc("/logs", logsHandler) - // Sqlite3 - err := logger.InitDB("/tmp/gontrol_logs.db") - if err != nil { - log.Fatal(err) - } - defer logger.CloseDB() - // initDB (cfg.Database.Username, cfg.Database.Password, cfg.Database.Host, cfg.Database.Port, cfg.Database.Name) - // mysql db = database.InitDB (cfg.Database.Username, cfg.Database.Password, cfg.Database.Host, cfg.Database.Port, cfg.Database.Name) defer db.Close() + name := randomname.GenerateRandomName() - fmt.Println(name) + fmt.Sprintf("Server instance: %s", name) webServer := &http.Server { Addr: ":3333", @@ -202,7 +237,7 @@ func main() { if err := websocketServer.ListenAndServe(); err != http.ErrServerClosed { log.Fatalf("Websocket server failed: %s", err) } - err = logger.InsertLog(logger.Info, logLine) + err := logger.InsertLog(logger.Info, logLine) if err != nil { log.Println("Error inserting log:", err) } @@ -218,7 +253,7 @@ func main() { log.Fatalf("Web server failed: %s", err) } - err = logger.InsertLog(logger.Info, logLine) + err := logger.InsertLog(logger.Info, logLine) if err != nil { log.Println("Error inserting log:", err) } diff --git a/src/logger/logger.go b/src/logger/logger.go index c4158c3..5fb7a09 100644 --- a/src/logger/logger.go +++ b/src/logger/logger.go @@ -5,6 +5,7 @@ import ( "fmt" "log" // "net/http" + "strings" "sync" "time" @@ -17,16 +18,16 @@ var ( lite_dbMutex sync.Mutex logMutex sync.Mutex - logLimit = 10 + logLimit = 100 // LogEntries []string ) const ( - Debug LogLevel = "DEBUG" - Info LogLevel = "INFO" - Warning LogLevel = "WARNING" - Error LogLevel = "ERROR" - Fatal LogLevel = "FATAL" + Debug LogLevel = "debug" + Info LogLevel = "info" + Warning LogLevel = "warning" + Error LogLevel = "error" + Fatal LogLevel = "fatal" ) type LogLevel string @@ -74,6 +75,8 @@ func InsertLog(level LogLevel, message string) error { return fmt.Errorf("Error starting transaction: %w", err) } + message = strings.ReplaceAll(message, `"`, `\"`) + // insertQuery := `INSERT INTO logs (message) VALUES (?)` insertQuery := `INSERT INTO logs (message, level) VALUES (?, ?)` // _, err := Lite_db.Exec(insertQuery, message, level) diff --git a/src/server/websocket/websocketServer.go b/src/server/websocket/websocketServer.go index 56a9479..ed127aa 100644 --- a/src/server/websocket/websocketServer.go +++ b/src/server/websocket/websocketServer.go @@ -60,21 +60,22 @@ func registerAgent(agentName, agentId, agentIp, agentType string) error { resp, err := http.PostForm(registerURL, form) if err != nil { log.Printf("Error registering agent: %v", err) + logger.InsertLog(logger.Error, fmt.Sprintf("Error registering agent: %v", err)) return err } defer resp.Body.Close() if resp.StatusCode == http.StatusCreated { - logLine := fmt.Sprintf("Agent %s successfully registered.", agentName) - log.Printf(logLine) - // logLine = logger.ToLog(logLine) - logger.InsertLog(logger.Info, logLine) + log.Printf("Agent %s successfully registered.", agentName) + logger.InsertLog(logger.Info, fmt.Sprintf("Agent %s successfully registered.", agentName)) return nil } else if resp.StatusCode == http.StatusOK { log.Printf("Agent %s already registered.", agentName) + logger.InsertLog(logger.Info, fmt.Sprintf("Agent %s already registered.", agentName)) return nil } else { log.Printf("Failed to register agent, status: %v", resp.Status) + logger.InsertLog(logger.Error, fmt.Sprintf("Failed to register agent, status: %v", resp.Status)) return err } @@ -87,6 +88,7 @@ func getAgentDetails(agentId string) (*api.Agent, error) { resp, err := http.Get(agentURL) if err != nil { log.Printf("Failed to make GET request: %s", err) + logger.InsertLog(logger.Error, fmt.Sprintf("Failed to make GET request: %s", err)) return nil, err } defer resp.Body.Close() @@ -94,6 +96,7 @@ func getAgentDetails(agentId string) (*api.Agent, error) { doc, err := goquery.NewDocumentFromReader(resp.Body) if err != nil { log.Printf("Failed to parse HTML: %s", err) + logger.InsertLog(logger.Error, fmt.Sprintf("Failed to parse HTML: %s", err)) return nil, err } @@ -104,6 +107,7 @@ func getAgentDetails(agentId string) (*api.Agent, error) { agent.AgentID, err = strconv.Atoi(strings.TrimSpace(strings.TrimPrefix(text, "ID:"))) if err != nil { log.Printf("Converting string to integer failed in getAgentDetails(): %s", err) + logger.InsertLog(logger.Error, fmt.Sprintf("Converting string to integer failed in getAgentDetails(): %s", err)) } } else if strings.HasPrefix(text, "Name:") { agent.AgentName = strings.TrimSpace(strings.TrimPrefix(text, "Name:")) @@ -130,24 +134,28 @@ func getAgentIds() ([]string, error) { resp, err := http.Get(idURL) if err != nil { log.Printf("Failed to make GET request: %s", err) + logger.InsertLog(logger.Error, fmt.Sprintf("Failed to make GET request: %s", err)) return nil, err } defer resp.Body.Close() if resp.StatusCode != http.StatusOK { log.Printf("Unexpected status code: %d", resp.StatusCode) + logger.InsertLog(logger.Info, fmt.Sprintf("Unexpected status code: %d", resp.StatusCode)) return nil, nil } body, err := io.ReadAll(resp.Body) if err != nil { log.Printf("Failed to read response body: %s", err) + logger.InsertLog(logger.Error, fmt.Sprintf("Failed to read response body: %s", err)) return nil, err } var agentIds []string if err := json.Unmarshal(body, &agentIds); err != nil { log.Printf("Failed to parse JSON response: %s", err) + logger.InsertLog(logger.Error, fmt.Sprintf("Failed to parse JSON response: %s", err)) return nil, err } @@ -159,6 +167,7 @@ func (wsh webSocketHandler) ServeHTTP(w http.ResponseWriter, r *http.Request){ c, err := wsh.upgrader.Upgrade(w, r, nil) if err != nil { log.Printf("Error %s when upgrading connection to websocket", err) + logger.InsertLog(logger.Error, fmt.Sprintf("Error %s when upgrading connection to websocket", err)) return } @@ -183,13 +192,13 @@ func (wsh webSocketHandler) ServeHTTP(w http.ResponseWriter, r *http.Request){ if agentName == "" || agentIP == "" { log.Printf("Missing agentName or IPv4Address in query parameters") + logger.InsertLog(logger.Info, fmt.Sprintf("Missing agentName or IPv4Address in query parameters")) c.Close() return } - logLine := fmt.Sprintf("Agent %s connected: %s (%s)", agentId, agentName, agentIP) - log.Printf(logLine) - logger.InsertLog(logger.Info, logLine) + log.Printf("Agent %s connected: %s (%s)", agentId, agentName, agentIP) + logger.InsertLog(logger.Info, fmt.Sprintf("Agent %s connected: %s (%s)", agentId, agentName, agentIP)) agentSocketsMutex.Lock() agentSockets[agentName] = c @@ -200,22 +209,19 @@ func (wsh webSocketHandler) ServeHTTP(w http.ResponseWriter, r *http.Request){ delete(agentSockets, agentName) agentSocketsMutex.Unlock() c.Close() - logLine := fmt.Sprintf("Agent disconnected: %s (%s)", agentName, agentIP) - log.Printf(logLine) - logger.InsertLog(logger.Info, logLine) + log.Printf("Agent disconnected: %s (%s)", agentName, agentIP) + logger.InsertLog(logger.Info, fmt.Sprintf("Agent disconnected: %s (%s)", agentName, agentIP)) }() for { _, message, err := c.ReadMessage() if err != nil { - logLine := fmt.Sprintf("Error reading from agent %s: %v", agentName, err) - log.Printf(logLine) - logger.InsertLog(logger.Error, logLine) + log.Printf("Error reading from agent %s: %v", agentName, err) + logger.InsertLog(logger.Error, fmt.Sprintf("Error reading from agent %s: %v", agentName, err)) break } - logLine := fmt.Sprintf("Message from agent %s: %s", agentName, message) - log.Printf(logLine) - logger.InsertLog(logger.Debug, logLine) + log.Printf("Message from agent %s: %s", agentName, message) + logger.InsertLog(logger.Info, fmt.Sprintf("Message from agent %s: %s", agentName, message)) if ch, ok := responseChannels.Load(agentName); ok { responseChan := ch.(chan string) @@ -233,6 +239,7 @@ var executeCommand http.HandlerFunc = func(w http.ResponseWriter, r *http.Reques err := r.ParseForm() if err != nil { http.Error(w, "Invalid form data", http.StatusBadRequest) + logger.InsertLog(logger.Info, "Invalid form data") return } @@ -245,6 +252,7 @@ var executeCommand http.HandlerFunc = func(w http.ResponseWriter, r *http.Reques if !ok { http.Error(w, "Agent not connected", http.StatusNotFound) + logger.InsertLog(logger.Info, "Agent not connected") return } @@ -262,6 +270,7 @@ var executeCommand http.HandlerFunc = func(w http.ResponseWriter, r *http.Reques err = conn.WriteMessage(websocket.TextMessage, messageBytes) if err != nil { http.Error(w, "Failed to send command to the agent", http.StatusInternalServerError) + logger.InsertLog(logger.Error, "Failed to send command to the agent") return } @@ -275,13 +284,15 @@ var executeCommand http.HandlerFunc = func(w http.ResponseWriter, r *http.Reques payload, ok := parsedResponse["payload"] if !ok { http.Error(w, "Invalid response structure", http.StatusInternalServerError) + logger.InsertLog(logger.Error, "Invalid response structure") return } w.WriteHeader(http.StatusOK) w.Header().Set("Content-Type", "text/plain") w.Write([]byte(payload)) case <- time.After(10 * time.Second): - http.Error(w, "Agent repsonse timed out", http.StatusGatewayTimeout) + http.Error(w, "Agent response timed out", http.StatusGatewayTimeout) + logger.InsertLog(logger.Info, "Agent response timed out") } } diff --git a/templates/index.html b/templates/index.html index ba16f96..30c094c 100644 --- a/templates/index.html +++ b/templates/index.html @@ -11,45 +11,119 @@ g2: gommand & gontrol + +
@@ -58,15 +132,17 @@

Agents

- -
+
+

Command Execution

- +
@@ -80,33 +156,46 @@
- -
- -
- - -
-
- - -
-
- - -
-
- - -
-
- - -
- - -
-
+ + + + + + + + + + + + + + + + + + + + + + + + + + + + + +

Logs

+
+ + + + + + + + +
diff --git a/templates/logs.html b/templates/logs.html index bafd305..cc18902 100644 --- a/templates/logs.html +++ b/templates/logs.html @@ -5,6 +5,23 @@ Web Server Logs +

Web Server Logs

@@ -13,12 +30,51 @@
- - + + diff --git a/templates/partials/agent_list.html b/templates/partials/agent_list.html index b08f196..d78210e 100644 --- a/templates/partials/agent_list.html +++ b/templates/partials/agent_list.html @@ -22,7 +22,7 @@ - + {{.Status}} diff --git a/templates/partials/logs_partial.html b/templates/partials/logs_partial.html new file mode 100644 index 0000000..d2a771c --- /dev/null +++ b/templates/partials/logs_partial.html @@ -0,0 +1,5 @@ +{{range .}} +

+ ts={{.Timestamp}} level={{.Level}} msg="{{.Message}}" +

+{{end}}