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
This commit is contained in:
parent
47ba7aa229
commit
76d343e8ac
|
@ -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)
|
||||
}
|
||||
|
|
21
main.go
21
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)
|
||||
|
|
|
@ -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,45 +60,41 @@ func InitDB(dbPath string) error {
|
|||
level TEXT
|
||||
);`
|
||||
|
||||
_, err = Lite_db.Exec(CreateTableQuery)
|
||||
if _, err = db.Exec(CreateTableQuery); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return New(db), nil
|
||||
}
|
||||
|
||||
func (l *Logger) InsertLog(level LogLevel, message string) error {
|
||||
l.lock.Lock()
|
||||
defer l.lock.Unlock()
|
||||
|
||||
tx, err := l.db.Begin()
|
||||
if err != nil {
|
||||
return fmt.Errorf("Error creating table: %w", err)
|
||||
return fmt.Errorf("InsertLog: start tx: %w", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func InsertLog(level LogLevel, message string) error {
|
||||
lite_dbMutex.Lock()
|
||||
defer lite_dbMutex.Unlock()
|
||||
|
||||
// Future use may fulfill multiple transactions
|
||||
tx, err := Lite_db.Begin()
|
||||
if err != nil {
|
||||
return fmt.Errorf("Error starting transaction: %w", err)
|
||||
}
|
||||
|
||||
message = strings.ReplaceAll(message, `"`, `\"`)
|
||||
|
||||
insertQuery := `INSERT INTO logs (message, level) VALUES(?, ?)`
|
||||
_, err = tx.Exec(insertQuery, message, level)
|
||||
message = strings.ReplaceAll(message, `"`, `\"`)
|
||||
|
||||
_, 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{"%"}
|
||||
|
@ -102,12 +102,13 @@ func FetchLogs(limit int, levels []string) ([]LogEntry, error) {
|
|||
|
||||
var args[] interface {}
|
||||
placeholders := make([]string, len(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...))
|
||||
}
|
||||
|
|
|
@ -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()
|
||||
|
||||
|
|
|
@ -4,14 +4,13 @@ import (
|
|||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"testing"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/DATA-DOG/go-sqlmock"
|
||||
"gontrol/src/logger"
|
||||
)
|
||||
|
||||
func newTestApp(t *testing.T) (*App, sqlmock.Sqlmock) {
|
||||
func newTestAppWithSQLiteLogger(t *testing.T) *App {
|
||||
t.Helper()
|
||||
|
||||
tmpl, err := ParseTemplates()
|
||||
|
@ -19,43 +18,35 @@ func newTestApp(t *testing.T) (*App, sqlmock.Sqlmock) {
|
|||
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)
|
||||
}
|
||||
}
|
||||
|
||||
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())
|
||||
t.Cleanup(func() {
|
||||
os.Remove(tmp.Name())
|
||||
})
|
||||
|
||||
if err := logger.InitDB(tmp.Name()); err != nil {
|
||||
t.Fatalf("logger.InitDB: %v", err)
|
||||
logInstance, err := logger.Init(tmp.Name())
|
||||
if err != nil {
|
||||
t.Fatalf("logger.Init: %v", err)
|
||||
}
|
||||
defer logger.CloseDB()
|
||||
t.Cleanup(func() {
|
||||
logInstance.Close()
|
||||
})
|
||||
|
||||
logger.InsertLog(logger.Error, "fake-log-1")
|
||||
|
||||
app, mock := newTestApp(t)
|
||||
defer app.DB.Close()
|
||||
|
||||
if err := mock.ExpectationsWereMet(); err != nil {
|
||||
t.Fatalf("sqlmock expectations: %v", err)
|
||||
return &App{
|
||||
Tmpl: tmpl,
|
||||
Logger: logInstance,
|
||||
}
|
||||
}
|
||||
|
||||
// Test that the logsHandler returns HTML with logs correctly
|
||||
func TestLogsHandler_HTML(t *testing.T) {
|
||||
app := newTestAppWithSQLiteLogger(t)
|
||||
|
||||
// Insert a log entry to be fetched
|
||||
app.Logger.Error("fake-log-1")
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/logs/error?limit=1", nil)
|
||||
rec := httptest.NewRecorder()
|
||||
|
@ -73,11 +64,9 @@ func TestLogsHandler_HTML(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
// 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, _ := newTestApp(t) // helper from previous file
|
||||
app := newTestAppWithSQLiteLogger(t)
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/logs?limit=abc", nil)
|
||||
rec := httptest.NewRecorder()
|
||||
|
@ -92,11 +81,12 @@ func TestLogsHandler_InvalidLimit(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
// 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")
|
||||
|
@ -108,9 +98,9 @@ func TestLogsHandler_JSON(t *testing.T) {
|
|||
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)
|
||||
t.Errorf("Content-Type = %q; want application/json", ct)
|
||||
}
|
||||
if !strings.Contains(rec.Body.String(), "unit‑test") {
|
||||
if !strings.Contains(rec.Body.String(), "unit-test-json") {
|
||||
t.Errorf("JSON body missing expected log entry")
|
||||
}
|
||||
}
|
||||
|
|
|
@ -1,4 +1,3 @@
|
|||
// src/server/webapp/static_files_test.go
|
||||
package webapp
|
||||
|
||||
import (
|
||||
|
@ -12,7 +11,7 @@ import (
|
|||
// 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 {
|
||||
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
|
||||
|
@ -31,9 +30,10 @@ func TestStaticFileServer(t *testing.T) {
|
|||
}
|
||||
|
||||
//-----------------------------------------------------------------
|
||||
// build router with sqlmock DB (not used in this test)
|
||||
// build router with SQLite Logger (real DB) for test
|
||||
//-----------------------------------------------------------------
|
||||
app, _ := newTestApp(t)
|
||||
app := newTestAppWithSQLiteLogger(t)
|
||||
|
||||
ts := httptest.NewServer(BuildRouter(app))
|
||||
defer ts.Close()
|
||||
|
||||
|
|
|
@ -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 {
|
||||
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
|
||||
}
|
||||
|
||||
|
|
Loading…
Reference in New Issue