From 76d343e8ac290edac666c29b59aef19640906f4b Mon Sep 17 00:00:00 2001 From: Stefan Etringer Date: Tue, 15 Jul 2025 13:16:48 +0000 Subject: [PATCH] added the logger interface, json logs, and logs input from agents to the server. This means there is an new type which is named logs. It can be sent from the agents to the server --- agents/agent.go | 29 +++++ main.go | 21 +-- src/logger/logger.go | 132 +++++++++++-------- src/server/webapp/handlers.go | 6 +- src/server/webapp/handlers_test.go | 164 +++++++++++------------- src/server/webapp/static_files_test.go | 76 +++++------ src/server/websocket/websocketServer.go | 118 +++++++++++++---- 7 files changed, 333 insertions(+), 213 deletions(-) diff --git a/agents/agent.go b/agents/agent.go index 4a2b8cb..4e5cbf0 100644 --- a/agents/agent.go +++ b/agents/agent.go @@ -35,6 +35,7 @@ type Agent struct { type Message struct { Type string `json:"type"` + Level string `json:"level"` Payload string `json:"payload"` } @@ -84,6 +85,12 @@ func reconnectToWebSocket(agentName, agentId, agentIp, agentType, hostName strin backoff := 2 * time.Second maxBackoff := 1 * time.Minute + err := sendLog(conn) + + if err != nil { + log.Println("sendLog() error") + } + for { log.Println("Attempting to reconnect to WebSocket...") err := connectToWebSocket(agentName, agentId, agentIp, agentType, hostName) @@ -102,11 +109,28 @@ func reconnectToWebSocket(agentName, agentId, agentIp, agentType, hostName strin } } +func sendLog(conn *websocket.Conn) error { + + response := Message{ + Type: "response", + Level: "Fatal", + Payload: "Remote logging is working!", + } + responseBytes, _ := json.Marshal(response) + if err := conn.WriteMessage(websocket.TextMessage, responseBytes); err != nil { + log.Printf("Error sending output: %v", err) + return err + } + + return nil +} + 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, hostName); reconnectErr != nil { @@ -203,5 +227,10 @@ func main() { log.Fatalf("Websocket connection failed: %v", err) } + err := sendLog(conn) + if err != nil { + log.Println("sendLog() error") + } + listenForCommands(agentName, agentId, agentIp, agentType, hostName) } diff --git a/main.go b/main.go index 486bb26..ef7ac67 100644 --- a/main.go +++ b/main.go @@ -12,7 +12,9 @@ import ( "gontrol/src/logger" "gontrol/src/server/database" "gontrol/src/server/webapp" - "gontrol/src/server/websocket" + websocketserver "gontrol/src/server/websocket" + + // _ "net/http/pprof" ) func main() { @@ -24,12 +26,13 @@ func main() { db := database.InitSQLiteDB("/tmp/gontrol_agents.db") defer db.Close() - if err := logger.InitDB("/tmp/gontrol_logs.db"); err != nil { + logService, err := logger.Init("/tmp/gontrol_logs.db") + if err != nil { log.Fatalf("Init log db: %v", err) } - defer logger.CloseDB() + defer logService.Close() - app := &webapp.App{Tmpl: tmpl, DB: db} + app := &webapp.App{Tmpl: tmpl, DB: db, Logger: logService} srv := &http.Server { Addr: ":3333", @@ -38,6 +41,10 @@ func main() { WriteTimeout: 10 * time.Second, } + // go func () { + // log.Println(http.ListenAndServe("localhost:6060", nil)) + // }() + go func() { log.Println("Web server is running on port :3333") if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed { @@ -48,14 +55,12 @@ func main() { go func() { logLine := "Websocket server is running on port :5555" log.Println(logLine) + websocketserver.SetLogger(logService) websocketServer := websocketserver.Server() if err := websocketServer.ListenAndServe(); err != http.ErrServerClosed { log.Fatalf("Websocket server failed: %s", err) } - err := logger.InsertLog(logger.Info, logLine) - if err != nil { - log.Println("Error inserting log:", err) - } + app.Logger.Log(logger.Info, logLine) }() stop := make(chan os.Signal, 1) diff --git a/src/logger/logger.go b/src/logger/logger.go index 3988602..abf8a79 100644 --- a/src/logger/logger.go +++ b/src/logger/logger.go @@ -3,8 +3,6 @@ package logger import ( "database/sql" "fmt" - "log" - // "net/http" "strings" "sync" "time" @@ -12,14 +10,16 @@ import ( _ "github.com/mattn/go-sqlite3" ) - -var ( - Lite_db *sql.DB - lite_dbMutex sync.Mutex - - logMutex sync.Mutex - logLimit = 100 -) +type LoggerInterface interface { + Log(level LogLevel, message string) + FetchLogs(limit int, levels []string) ([]LogEntry, error) + Error(format string, args ...any) + Info(format string, args ...any) + Warning(format string, args ...any) + Debug(format string, args ...any) + Fatal(format string, args ...any) + Close() error +} const ( Debug LogLevel = "debug" @@ -37,16 +37,20 @@ type LogEntry struct { Level LogLevel } -func ToLog(logLine string) string { - log := fmt.Sprintf("%s",time.Now().Format(time.RFC3339) + " " + logLine) - return log +type Logger struct { + db *sql.DB + lock sync.Mutex } -func InitDB(dbPath string) error { - var err error - Lite_db, err = sql.Open("sqlite3", dbPath) +func New(db *sql.DB) *Logger { + return &Logger{db: db} +} + + +func Init(dbPath string) (*Logger, error) { + db, err := sql.Open("sqlite3", dbPath) if err != nil { - return fmt.Errorf("Error opening DB: %w", err) + return nil, fmt.Errorf("Error opening DB: %w", err) } CreateTableQuery := `CREATE TABLE IF NOT EXISTS logs ( @@ -56,58 +60,55 @@ func InitDB(dbPath string) error { level TEXT );` - _, err = Lite_db.Exec(CreateTableQuery) - if err != nil { - return fmt.Errorf("Error creating table: %w", err) + if _, err = db.Exec(CreateTableQuery); err != nil { + return nil, err } - return nil + return New(db), nil } -func InsertLog(level LogLevel, message string) error { - lite_dbMutex.Lock() - defer lite_dbMutex.Unlock() +func (l *Logger) InsertLog(level LogLevel, message string) error { + l.lock.Lock() + defer l.lock.Unlock() - // Future use may fulfill multiple transactions - tx, err := Lite_db.Begin() + tx, err := l.db.Begin() if err != nil { - return fmt.Errorf("Error starting transaction: %w", err) + return fmt.Errorf("InsertLog: start tx: %w", err) } + insertQuery := `INSERT INTO logs (message, level) VALUES(?, ?)` message = strings.ReplaceAll(message, `"`, `\"`) - insertQuery := `INSERT INTO logs (message, level) VALUES (?, ?)` _, err = tx.Exec(insertQuery, message, level) - if err != nil { tx.Rollback() - return fmt.Errorf("Error inserting log: %v", err) + return fmt.Errorf("InsertLog: exec: %w", err) } - err = tx.Commit() - if err != nil { - return fmt.Errorf("Error committing transaction: %w", err) + if err := tx.Commit(); err != nil { + return fmt.Errorf("InsertLog: commit: %w", err) } return nil } -func FetchLogs(limit int, levels []string) ([]LogEntry, error) { - lite_dbMutex.Lock() - defer lite_dbMutex.Unlock() +func (l *Logger) FetchLogs(limit int, levels []string) ([]LogEntry, error) { + l.lock.Lock() + defer l.lock.Unlock() if len(levels) == 0 { levels = []string{"%"} } - var args []interface{} + var args[] interface {} placeholders := make([]string, len(levels)) - - for i, level := range levels { + for i, level := range levels{ placeholders[i] = "level LIKE ?" args = append(args, level) } + args = append(args, limit) + query := fmt.Sprintf(` SELECT timestamp, level, message FROM logs @@ -115,31 +116,54 @@ func FetchLogs(limit int, levels []string) ([]LogEntry, error) { ORDER BY timestamp DESC LIMIT ?`, strings.Join(placeholders, " OR ")) - args = append(args, limit) - - rows, err := Lite_db.Query(query, args...) + rows, err := l.db.Query(query, args...) if err != nil { - return nil, fmt.Errorf("Error fetching logs: %w", err) + return nil, fmt.Errorf("FetchLogs: query: %w", err) } defer rows.Close() var logs []LogEntry for rows.Next() { - var logEntry LogEntry - if err := rows.Scan( &logEntry.Timestamp, &logEntry.Level, &logEntry.Message); err != nil { - return nil, fmt.Errorf("Error scanning row: %w", err) + var entry LogEntry + if err := rows.Scan(&entry.Timestamp, &entry.Level, &entry.Message); err != nil { + return nil, fmt.Errorf("FetchLogs: scan: %w", err) } - logs = append(logs, logEntry) + logs = append(logs, entry) } return logs, nil } -func CloseDB() { - if Lite_db != nil { - err := Lite_db.Close() - if err != nil { - log.Printf("Error closing database: %v", err) - } - } + +func (l * Logger) Close() error { + return l.db.Close() +} + +func ToLog(logLine string) string { + log := fmt.Sprintf("%s",time.Now().Format(time.RFC3339) + " " + logLine) + return log +} + +func (l *Logger) Log(level LogLevel, message string) { + _ = l.InsertLog(level, message) +} + +func (l *Logger) Error(format string, args ...any) { + l.Log(Error, fmt.Sprintf(format, args...)) +} + +func (l *Logger) Info(format string, args ...any) { + l.Log(Info, fmt.Sprintf(format, args...)) +} + +func (l *Logger) Warning(format string, args ...any) { + l.Log(Warning, fmt.Sprintf(format, args...)) +} + +func (l *Logger) Debug(format string, args ...any) { + l.Log(Debug, fmt.Sprintf(format, args...)) +} + +func (l *Logger) Fatal(format string, args ...any) { + l.Log(Fatal, fmt.Sprintf(format, args...)) } diff --git a/src/server/webapp/handlers.go b/src/server/webapp/handlers.go index 3d6dd7e..935407b 100644 --- a/src/server/webapp/handlers.go +++ b/src/server/webapp/handlers.go @@ -25,6 +25,8 @@ import ( type App struct { Tmpl *template.Template DB *sql.DB + // Logger *logger.Logger + Logger logger.LoggerInterface } func (a *App) renderTemplate(w http.ResponseWriter, name string, data any) { @@ -139,7 +141,7 @@ func (a *App) logsHandler(w http.ResponseWriter, r *http.Request) { } // Call the police... I mean logger - logs, err := logger.FetchLogs(limit, levels) + logs, err := a.Logger.FetchLogs(limit, levels) if err != nil { http.Error(w, "Error fetching logs", http.StatusInternalServerError) @@ -165,7 +167,7 @@ func (a *App) 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") + a.Logger.Log(logger.Error, "Error fetching agent names from websocketServer") } defer resp.Body.Close() diff --git a/src/server/webapp/handlers_test.go b/src/server/webapp/handlers_test.go index f8036bc..100acc8 100644 --- a/src/server/webapp/handlers_test.go +++ b/src/server/webapp/handlers_test.go @@ -1,116 +1,106 @@ package webapp import ( - "net/http" - "net/http/httptest" - "os" - "testing" - "strings" + "net/http" + "net/http/httptest" + "os" + "strings" + "testing" - "github.com/DATA-DOG/go-sqlmock" - "gontrol/src/logger" + "gontrol/src/logger" ) -func newTestApp(t *testing.T) (*App, sqlmock.Sqlmock) { - t.Helper() +func newTestAppWithSQLiteLogger(t *testing.T) *App { + t.Helper() - tmpl, err := ParseTemplates() - if err != nil { - t.Fatalf("ParseTemplates: %v", err) - } + tmpl, err := ParseTemplates() + if err != nil { + t.Fatalf("ParseTemplates: %v", err) + } - db, mock, err := sqlmock.New() - if err != nil { - t.Fatalf("sqlmock: %v", err) - } - return &App{Tmpl: tmpl, DB: db}, mock -} - -// initFakeLogs puts one record in the in‑memory log DB used by logger.FetchLogs. -func initFakeLogs(t *testing.T) { - t.Helper() - if err := logger.InsertLog(logger.Info, "unit‑test"); err != nil { - t.Fatalf("InsertLog: %v", err) - } + tmp, err := os.CreateTemp("", "logs_*.db") + if err != nil { + t.Fatalf("Temp DB: %v", err) + } + tmp.Close() + t.Cleanup(func() { + os.Remove(tmp.Name()) + }) + + logInstance, err := logger.Init(tmp.Name()) + if err != nil { + t.Fatalf("logger.Init: %v", err) + } + t.Cleanup(func() { + logInstance.Close() + }) + + return &App{ + Tmpl: tmpl, + Logger: logInstance, + } } +// Test that the logsHandler returns HTML with logs correctly func TestLogsHandler_HTML(t *testing.T) { - tmp, err := os.CreateTemp("", "logs_*.db") - if err != nil { - t.Fatalf("Temp DB: %v", err) - } - tmp.Close() - defer os.Remove(tmp.Name()) + app := newTestAppWithSQLiteLogger(t) - if err := logger.InitDB(tmp.Name()); err != nil { - t.Fatalf("logger.InitDB: %v", err) - } - defer logger.CloseDB() + // Insert a log entry to be fetched + app.Logger.Error("fake-log-1") - logger.InsertLog(logger.Error, "fake-log-1") + req := httptest.NewRequest(http.MethodGet, "/logs/error?limit=1", nil) + rec := httptest.NewRecorder() - app, mock := newTestApp(t) - defer app.DB.Close() + app.logsHandler(rec, req) - if err := mock.ExpectationsWereMet(); err != nil { - t.Fatalf("sqlmock expectations: %v", err) - } - - - req := httptest.NewRequest(http.MethodGet, "/logs/error?limit=1", nil) - rec := httptest.NewRecorder() - - app.logsHandler(rec, req) - - if rec.Code != http.StatusOK { - t.Fatalf("status = %d; want 200", rec.Code) - } - if ct := rec.Header().Get("Content-Type"); ct != "text/html; charset=utf-8" { - t.Errorf("Content-Type = %s; want text/html", ct) - } - if !strings.Contains(rec.Body.String(), "fake-log-1") { - t.Errorf("rendered HTML missing our fake log line") - } + if rec.Code != http.StatusOK { + t.Fatalf("status = %d; want 200", rec.Code) + } + if ct := rec.Header().Get("Content-Type"); ct != "text/html; charset=utf-8" { + t.Errorf("Content-Type = %s; want text/html", ct) + } + if !strings.Contains(rec.Body.String(), "fake-log-1") { + t.Errorf("rendered HTML missing our fake log line") + } } -// Test that `limit=abc` is rejected with 400 and FetchLogs is **not** hit. +// Test that invalid 'limit' query param returns 400 Bad Request func TestLogsHandler_InvalidLimit(t *testing.T) { - initFakeLogs(t) + app := newTestAppWithSQLiteLogger(t) - app, _ := newTestApp(t) // helper from previous file + req := httptest.NewRequest(http.MethodGet, "/logs?limit=abc", nil) + rec := httptest.NewRecorder() - req := httptest.NewRequest(http.MethodGet, "/logs?limit=abc", nil) - rec := httptest.NewRecorder() + app.logsHandler(rec, req) - app.logsHandler(rec, req) - - if rec.Code != http.StatusBadRequest { - t.Fatalf("status = %d; want 400", rec.Code) - } - if !strings.Contains(rec.Body.String(), "Invalid count value") { - t.Errorf("missing error message") - } + if rec.Code != http.StatusBadRequest { + t.Fatalf("status = %d; want 400", rec.Code) + } + if !strings.Contains(rec.Body.String(), "Invalid count value") { + t.Errorf("missing error message") + } } -// Test JSON response when client sends Accept: application/json. +// Test JSON output when client sends Accept: application/json func TestLogsHandler_JSON(t *testing.T) { - initFakeLogs(t) + app := newTestAppWithSQLiteLogger(t) - app, _ := newTestApp(t) + // Insert a log entry to be fetched + app.Logger.Info("unit-test-json") - req := httptest.NewRequest(http.MethodGet, "/logs?limit=1", nil) - req.Header.Set("Accept", "application/json") - rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/logs?limit=1", nil) + req.Header.Set("Accept", "application/json") + rec := httptest.NewRecorder() - app.logsHandler(rec, req) + app.logsHandler(rec, req) - if rec.Code != http.StatusOK { - t.Fatalf("status = %d; want 200", rec.Code) - } - if ct := rec.Header().Get("Content-Type"); ct != "application/json" { - t.Errorf("Content‑Type = %q; want application/json", ct) - } - if !strings.Contains(rec.Body.String(), "unit‑test") { - t.Errorf("JSON body missing expected log entry") - } + if rec.Code != http.StatusOK { + t.Fatalf("status = %d; want 200", rec.Code) + } + if ct := rec.Header().Get("Content-Type"); ct != "application/json" { + t.Errorf("Content-Type = %q; want application/json", ct) + } + if !strings.Contains(rec.Body.String(), "unit-test-json") { + t.Errorf("JSON body missing expected log entry") + } } diff --git a/src/server/webapp/static_files_test.go b/src/server/webapp/static_files_test.go index 0389f6c..1d923ab 100644 --- a/src/server/webapp/static_files_test.go +++ b/src/server/webapp/static_files_test.go @@ -1,54 +1,54 @@ -// src/server/webapp/static_files_test.go package webapp import ( - "io" - "io/fs" - "net/http" - "net/http/httptest" - "testing" + "io" + "io/fs" + "net/http" + "net/http/httptest" + "testing" ) // findFirstStaticFile walks the embedded FS and returns the first file path. func findFirstStaticFile() (webPath string, content []byte, ok bool) { - _ = fs.WalkDir(assets, "static", func(path string, d fs.DirEntry, err error) error { - if !d.IsDir() && ok == false { - data, _ := assets.ReadFile(path) // ignore err; test will skip if nil - webPath = "/" + path // e.g. /static/css/main.css - content = data - ok = true - } - return nil - }) - return + _ = fs.WalkDir(assets, "static", func(path string, d fs.DirEntry, err error) error { + if !d.IsDir() && !ok { + data, _ := assets.ReadFile(path) // ignore err; test will skip if nil + webPath = "/" + path // e.g. /static/css/main.css + content = data + ok = true + } + return nil + }) + return } // Requires at least one file under static/. Skips if none embedded. func TestStaticFileServer(t *testing.T) { - webPath, wantBytes, ok := findFirstStaticFile() - if !ok { - t.Skip("no embedded static files to test") - } + webPath, wantBytes, ok := findFirstStaticFile() + if !ok { + t.Skip("no embedded static files to test") + } - //----------------------------------------------------------------- - // build router with sqlmock DB (not used in this test) - //----------------------------------------------------------------- - app, _ := newTestApp(t) - ts := httptest.NewServer(BuildRouter(app)) - defer ts.Close() + //----------------------------------------------------------------- + // build router with SQLite Logger (real DB) for test + //----------------------------------------------------------------- + app := newTestAppWithSQLiteLogger(t) - res, err := http.Get(ts.URL + webPath) - if err != nil { - t.Fatalf("GET %s: %v", webPath, err) - } - defer res.Body.Close() + ts := httptest.NewServer(BuildRouter(app)) + defer ts.Close() - if res.StatusCode != http.StatusOK { - t.Fatalf("status = %d; want 200", res.StatusCode) - } + res, err := http.Get(ts.URL + webPath) + if err != nil { + t.Fatalf("GET %s: %v", webPath, err) + } + defer res.Body.Close() - gotBytes, _ := io.ReadAll(res.Body) - if len(gotBytes) == 0 || string(gotBytes) != string(wantBytes) { - t.Errorf("served file differs from embedded asset") - } + if res.StatusCode != http.StatusOK { + t.Fatalf("status = %d; want 200", res.StatusCode) + } + + gotBytes, _ := io.ReadAll(res.Body) + if len(gotBytes) == 0 || string(gotBytes) != string(wantBytes) { + t.Errorf("served file differs from embedded asset") + } } diff --git a/src/server/websocket/websocketServer.go b/src/server/websocket/websocketServer.go index f5d6e57..5a75b49 100644 --- a/src/server/websocket/websocketServer.go +++ b/src/server/websocket/websocketServer.go @@ -1,7 +1,6 @@ package websocketserver import ( - "database/sql" "encoding/json" "fmt" "gontrol/src/logger" @@ -22,12 +21,16 @@ import ( ) var responseChannels sync.Map // Key: agentName, Value: chan string -var db *sql.DB +var logService logger.LoggerInterface type webSocketHandler struct { upgrader websocket.Upgrader } +func SetLogger(logger logger.LoggerInterface) { + logService = logger +} + var agentSockets = make(map[string]*websocket.Conn) var agentSocketsMutex sync.Mutex @@ -62,22 +65,26 @@ func registerAgent(agentName, agentId, agentIp, agentType, addPort, hostname str 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)) + // logger.InsertLog(logger.Error, fmt.Sprintf("Error registering agent: %v", err)) + logService.Info(fmt.Sprintf("Error registering agent: %v", err)) return err } defer resp.Body.Close() if resp.StatusCode == http.StatusCreated { log.Printf("Agent %s successfully registered.", agentName) - logger.InsertLog(logger.Info, fmt.Sprintf("Agent %s successfully registered.", agentName)) + // logger.InsertLog(logger.Info, fmt.Sprintf("Agent %s successfully registered.", agentName)) + logService.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)) + // logger.InsertLog(logger.Info, fmt.Sprintf("Agent %s already registered.", agentName)) + logService.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)) + // logger.InsertLog(logger.Error, fmt.Sprintf("Failed to register agent, status: %v", resp.Status)) + logService.Error(fmt.Sprintf("Failer to register agent, status: %v", resp.Status)) return err } @@ -89,7 +96,8 @@ 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)) + // logger.InsertLog(logger.Error, fmt.Sprintf("Failed to make GET request: %s", err)) + logService.Error(fmt.Sprintf("Failed to make GET request: %s", err)) return nil, err } defer resp.Body.Close() @@ -97,7 +105,8 @@ 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)) + // logger.InsertLog(logger.Error, fmt.Sprintf("Failed to parse HTML: %s", err)) + logService.Error(fmt.Sprintf("Failed to parse HTML: %s", err)) return nil, err } @@ -108,7 +117,8 @@ 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)) + // logger.InsertLog(logger.Error, fmt.Sprintf("Converting string to integer failed in getAgentDetails(): %s", err)) + logService.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:")) @@ -134,28 +144,32 @@ 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)) + // logger.InsertLog(logger.Error, fmt.Sprintf("Failed to make GET request: %s", err)) + logService.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)) + // logger.InsertLog(logger.Info, fmt.Sprintf("Unexpected status code: %d", resp.StatusCode)) + logService.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)) + // logger.InsertLog(logger.Error, fmt.Sprintf("Failed to read response body: %s", err)) + logService.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)) + // logger.InsertLog(logger.Error, fmt.Sprintf("Failed to parse JSON response: %s", err)) + logService.Error(fmt.Sprintf("Failed to parse JSON response: %s", err)) return nil, err } @@ -167,7 +181,8 @@ 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)) + // logger.InsertLog(logger.Error, fmt.Sprintf("Error %s when upgrading connection to websocket", err)) + logService.Error(fmt.Sprintf("Error %s when upgrading connection to websocket", err)) return } @@ -200,13 +215,15 @@ 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")) + // logger.InsertLog(logger.Info, fmt.Sprintf("Missing agentName or IPv4Address in query parameters")) + logService.Info(fmt.Sprintf("Missing agentName or IPv4Address in query parameters")) c.Close() return } log.Printf("Agent %s connected: %s (%s)", agentId, agentName, agentIP) - logger.InsertLog(logger.Info, fmt.Sprintf("Agent %s connected: %s (%s)", agentId, agentName, agentIP)) + // logger.InsertLog(logger.Info, fmt.Sprintf("Agent %s connected: %s (%s)", agentId, agentName, agentIP)) + logService.Info(fmt.Sprintf("Agent %s connected: %s (%s)", agentId, agentName, agentIP)) agentSocketsMutex.Lock() agentSockets[agentName] = c @@ -218,19 +235,69 @@ func (wsh webSocketHandler) ServeHTTP(w http.ResponseWriter, r *http.Request){ agentSocketsMutex.Unlock() c.Close() log.Printf("Agent disconnected: %s (%s)", agentName, agentIP) - logger.InsertLog(logger.Info, fmt.Sprintf("Agent disconnected: %s (%s)", agentName, agentIP)) + // logger.InsertLog(logger.Info, fmt.Sprintf("Agent disconnected: %s (%s)", agentName, agentIP)) + logService.Info(fmt.Sprintf("Agent disconnected: %s (%s)", agentName, agentIP)) }() + // for { + // _, message, err := c.ReadMessage() + // if err != nil { + // log.Printf("Error reading from agent %s: %v", agentName, err) + // // logger.InsertLog(logger.Error, fmt.Sprintf("Error reading from agent %s: %v", agentName, err)) + // logService.Error(fmt.Sprintf("Error reading from agent %s: %v", agentName, err)) + // break + // } + // // log.Printf("Message from agent %s: %s", agentName, message) + // log.Printf("Message from agent %s received", agentName) + // // logger.InsertLog(logger.Debug, fmt.Sprintf("Message from agent %s: %s", agentName, message)) + // logService.Debug(fmt.Sprintf("Message from agent %s: %s", agentName, message)) + + // if ch, ok := responseChannels.Load(agentName); ok { + // responseChan := ch.(chan string) + // responseChan <- string(message) + // } + // } + for { _, message, err := c.ReadMessage() - if err != nil { + if err != nil { log.Printf("Error reading from agent %s: %v", agentName, err) - logger.InsertLog(logger.Error, fmt.Sprintf("Error reading from agent %s: %v", agentName, err)) + logService.Error(fmt.Sprintf("Agent disconnected: %s (%s)", agentName, agentIP)) break } - // log.Printf("Message from agent %s: %s", agentName, message) - log.Printf("Message from agent %s received", agentName) - logger.InsertLog(logger.Debug, fmt.Sprintf("Message from agent %s: %s", agentName, message)) + + var generic map [string]interface{} + if err := json.Unmarshal(message, &generic); err != nil { + logService.Error(fmt.Sprintf("Invalid JSON from %s: %v", agentName, err)) + continue + } + + // In case the message coming in is of type log, it gets switched into the corresponding level + if msgType, ok := generic["type"].(string); ok && msgType == "log" { + level := strings.ToLower(fmt.Sprintf("%v", generic["level"])) + content := fmt.Sprintf("%v", generic["payload"]) + + formatted := fmt.Sprintf("Log from %s: %s", agentName, content) + + switch level { + case "info": + logService.Info(formatted) + case "error": + logService.Error(formatted) + case "debug": + logService.Debug(formatted) + case "warn", "warning": + logService.Fatal(formatted) + case "fatal": + logService.Fatal(formatted) + default: + logService.Info(formatted) + } + continue + } + + // From here things get prepared for /executeCommand + logService.Debug(fmt.Sprintf("Message from agent %s: %s", agentName, message)) if ch, ok := responseChannels.Load(agentName); ok { responseChan := ch.(chan string) @@ -241,6 +308,7 @@ func (wsh webSocketHandler) ServeHTTP(w http.ResponseWriter, r *http.Request){ type Message struct { Type string `json:"type"` + Level string `json:"level"` Payload string `json:"payload"` } @@ -248,7 +316,8 @@ 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") + // logger.InsertLog(logger.Info, "Invalid form data") + logService.Info("Invalid form data") return } @@ -268,7 +337,8 @@ var executeCommand http.HandlerFunc = func(w http.ResponseWriter, r *http.Reques if len(agentNames) == 0 || command == "" { http.Error(w, "Missing agent or command", http.StatusBadRequest) - logger.InsertLog(logger.Error, "Missing agent or command") + // logger.InsertLog(logger.Error, "Missing agent or command") + logService.Error("Missing agent or command") return }