package main import ( "context" "encoding/json" "embed" "os" "strings" "time" "slices" "database/sql" "fmt" "html/template" "log" "net/http" "strconv" "sync" "syscall" // "io" "golang.org/x/net/html" "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 ) //go:embed static/* var staticFiles embed.FS 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 strings.Contains(r.Header.Get("Accept"), "json") { w.Header().Set("Content-Type", "application/json") jsonData, err := json.Marshal(agents) if err != nil { http.Error(w, "Failed to encode agents to JSON", http.StatusInternalServerError) return } w.Write(jsonData) return } 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) { // Warning this bit me in the nose: var countStr is []string, but // variable = countStr[0] is casted to int automatically // when the string is a number. Jesus Christ, this is odd behavior! levels := r.URL.Query()["level"] countStr := r.URL.Query()["limit"] var limit int = 128 if len(countStr) > 0 { parsedCount, err := strconv.Atoi(countStr[0]) if err == nil { limit = parsedCount } else { http.Error(w, "Invalid count value", http.StatusBadRequest) return } } // This enables not only `level` GET parameters but also selecting by paths // For example /logs/error is now identical to /logs?level=error parts := strings.Split(strings.TrimPrefix(r.URL.Path, "/logs/"), "/") if (len(levels) == 0) && len(parts) > 0 && parts[0] != "" { levels = []string{parts[0]} } // Call the police... I mean logger logs, err := logger.FetchLogs(limit, levels) if err != nil { http.Error(w, "Error fetching logs", http.StatusInternalServerError) return } renderTemplate(w, "templates/partials/logs_partial.html", logs) } func proxyAgentHandler(w http.ResponseWriter, r *http.Request) { agentIP := r.URL.Query().Get("ip") // e.g., 10.0.0.42 if agentIP == "" { http.Error(w, "Missing 'ip' parameter", http.StatusBadRequest) return } // Construct the URL to proxy to agentURL := "http://" + agentIP + ":8080" // Send request to agent server resp, err := http.Get(agentURL) if err != nil { http.Error(w, "Failed to reach agent: "+err.Error(), http.StatusBadGateway) return } defer resp.Body.Close() // Parse the HTML from the agent's response doc, err := html.Parse(resp.Body) if err != nil { http.Error(w, "Failed to parse agent response: "+err.Error(), http.StatusInternalServerError) return } // Extract all elements for stylesheets var stylesheets []string var extractStylesheets func(*html.Node) extractStylesheets = func(n *html.Node) { if n.Type == html.ElementNode && n.Data == "link" { for _, attr := range n.Attr { if attr.Key == "rel" && attr.Val == "stylesheet" { for _, attr := range n.Attr { if attr.Key == "href" { stylesheets = append(stylesheets, attr.Val) } } } } } for c := n.FirstChild; c != nil; c = c.NextSibling { extractStylesheets(c) } } extractStylesheets(doc) // Return the HTML and inject the stylesheets in the
section w.Header().Set("Content-Type", "text/html") w.WriteHeader(http.StatusOK) // Inject the stylesheets into the section of your page fmt.Fprintf(w, "") for _, stylesheet := range stylesheets { // Make sure the stylesheet is loaded properly (absolute URLs or proxy it) fmt.Fprintf(w, ``, stylesheet) } fmt.Fprintf(w, "") // Now, serve the HTML content of the agent web app (or an iframe) // Output the rest of the HTML (including the agent's content inside iframe) fmt.Fprintf(w, ` `, agentURL) fmt.Fprintf(w, "") } 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) webMux.HandleFunc("/logs/{level}", logsHandler) webMux.Handle("/static/", http.FileServer(http.FS(staticFiles))) webMux.HandleFunc("/proxyAgent", proxyAgentHandler) 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") }