package main import ( "context" "encoding/json" "os" "strings" "time" "slices" "database/sql" "fmt" "html/template" "log" "net/http" "sync" "syscall" "gontrol/src/logger" "gontrol/src/randomname" api "gontrol/src/server/api" "gontrol/src/server/database" websocketserver "gontrol/src/server/websocket" "os/signal" _ "github.com/go-sql-driver/mysql" "github.com/kelseyhightower/envconfig" ) var ( tmpl *template.Template db *sql.DB ) type Config struct { Database struct { Username string `envconfig:"DB_USERNAME"` Password string `envconfig:"DB_PASSWORD"` Port int16 `envconfig:"DB_PORT"` Name string `envconfig:"DB_NAME"` Host string `envconfig:"DB_HOST"` } } func readEnv(cfg *Config) { err := envconfig.Process("", cfg) if err != nil { processError(err) } } func processError(err error) { fmt.Println(err) os.Exit(2) } func init() { tmpl, _ = template.ParseGlob("templates/*.html") // Sqlite3 err := logger.InitDB("/tmp/gontrol_logs.db") if err != nil { log.Fatal(err) } } func renderTemplate(w http.ResponseWriter, tmpl string, data interface{}) { t, err := template.ParseFiles(tmpl) if err != nil { log.Printf("Failed to load template %s: %v", tmpl, err) http.Error(w, "Internal Server Error", http.StatusInternalServerError) return } if err := t.Execute(w, data); err != nil { log.Printf("Failed to render template %s: %v", tmpl, err) http.Error(w, "Internal Server Error", http.StatusInternalServerError) return } } func agentsHandler(w http.ResponseWriter, r *http.Request) { parts := strings.Split(strings.TrimPrefix(r.URL.Path, "/agents/"), "/") agentId := "" if len(parts) > 0 && parts[0] != "" { agentId = parts[0] } switch r.Method { case http.MethodDelete: api.DeleteAgent(db, w, r, agentId) listAgents(w,r) case http.MethodGet: if agentId == "" { listAgents(w, r) } else { agent, _ := api.GetAgent(db, w, r, agentId) renderTemplate(w, "templates/partials/agent_detail.html", agent) } case http.MethodPost: api.CreateAgent(db, w, r) listAgents(w, r) case http.MethodPut: api.UpdateAgent(db, w, r, agentId) listAgents(w, r) default: http.Error(w, "Method not allowed", http.StatusMethodNotAllowed) } } func getHomepage(w http.ResponseWriter, r *http.Request) { tmpl.ExecuteTemplate(w, "index.html", nil) } func listAgents(w http.ResponseWriter, r *http.Request) { var agents []api.Agent agents, err := api.GetAgents(db) currentAgents := getAgentsStatus() for i := range agents { if slices.Contains(currentAgents, agents[i].AgentName) { agents[i].Status = "Connected" } else { agents[i].Status = "Disconnected" } } if err != nil { http.Error(w, "Failed to fetch agents", http.StatusInternalServerError) return } 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 } func getAgentIds(w http.ResponseWriter, r *http.Request) { api.GetAgentIds(db, w, r) return } 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 } 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) ctx, cancel := context.WithCancel(context.Background()) defer cancel() var wg sync.WaitGroup websocketServer := websocketserver.Server() webMux := http.NewServeMux() webMux.HandleFunc("/", getHomepage) webMux.HandleFunc("/agents", agentsHandler) webMux.HandleFunc("/agentNames", getAgentNames) webMux.HandleFunc("/agentIds", getAgentIds) webMux.HandleFunc("/agents/{agentId}", agentsHandler) webMux.HandleFunc("/logs", logsHandler) 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.Sprintf("Server instance: %s", name) webServer := &http.Server { Addr: ":3333", Handler: webMux, } wg.Add(1) go func() { defer wg.Done() logLine := "Websocket server is running on port 5555" log.Println(logLine) 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) } }() wg.Add(1) go func() { defer wg.Done() logLine := "Web server is running on port 3333" log.Println(logLine) if err := webServer.ListenAndServe(); err != http.ErrServerClosed { log.Fatalf("Web server failed: %s", err) } err := logger.InsertLog(logger.Info, logLine) if err != nil { log.Println("Error inserting log:", err) } }() shutdownCh := make(chan os.Signal, 1) signal.Notify(shutdownCh, os.Interrupt, syscall.SIGTERM) <-shutdownCh log.Println("Shutdown signal received") shutdownCtx, shutdownCancel := context.WithTimeout(ctx, 10*time.Second) defer shutdownCancel() if err := websocketServer.Shutdown(shutdownCtx); err != nil { log.Printf("error shutting down websocket server: %s", err) } if err := webServer.Shutdown(shutdownCtx); err != nil { log.Printf("Error shutting down web server: %s", err) } wg.Wait() log.Println("All servers stopped") }