400 lines
11 KiB
Go
400 lines
11 KiB
Go
package main
|
||
|
||
import (
|
||
"context"
|
||
"encoding/json"
|
||
"embed"
|
||
"os"
|
||
"strings"
|
||
"time"
|
||
"slices"
|
||
"database/sql"
|
||
"fmt"
|
||
"html/template"
|
||
"log"
|
||
"net"
|
||
"net/http"
|
||
"strconv"
|
||
"sync"
|
||
"syscall"
|
||
|
||
"net/http/httputil"
|
||
"net/url"
|
||
// "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
|
||
|
||
// //go:embed templates
|
||
// var templateFiles 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")
|
||
|
||
// var err error
|
||
// tmpl, err = template.ParseFS(
|
||
// templateFiles,
|
||
// "templates/*.html",
|
||
// "templates/partials/*.html",
|
||
// )
|
||
|
||
// if err != nil {
|
||
// log.Fatalf("Failed to parse embedded templates: %v", err)
|
||
// }
|
||
|
||
// Sqlite3
|
||
err := logger.InitDB("/tmp/gontrol_logs.db")
|
||
|
||
if err != nil {
|
||
log.Fatal(err)
|
||
}
|
||
|
||
// Agents database
|
||
db = database.InitSQLiteDB("/tmp/gontrol_agents.db")
|
||
}
|
||
|
||
// func renderTemplate(w http.ResponseWriter, tmplPath string, data interface{}) {
|
||
// t := tmpl.Lookup(strings.TrimPrefix(tmplPath, "templates/"))
|
||
// if t == nil {
|
||
// log.Printf("Template %s not found", tmplPath)
|
||
// 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", tmplPath, err)
|
||
// http.Error(w, "Internal Server Error", http.StatusInternalServerError)
|
||
// }
|
||
|
||
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)
|
||
// renderTemplate(w, "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)
|
||
// renderTemplate(w, "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)
|
||
// renderTemplate(w, "logs_partial.html", logs)
|
||
}
|
||
|
||
// proxyAgentHandler tunnels HTTP and WebSocket traffic to an agent
|
||
// selected by ?ip=…&port=… . It keeps the path *after* /proxyAgent,
|
||
// removes the two query parameters, disables HTTP/2 (mandatory for WS),
|
||
// and streams with no buffering.
|
||
func proxyAgentHandler(w http.ResponseWriter, r *http.Request) {
|
||
ip := r.URL.Query().Get("ip")
|
||
port := r.URL.Query().Get("port")
|
||
if ip == "" || port == "" {
|
||
http.Error(w, "ip and port query parameters are required", http.StatusBadRequest)
|
||
return
|
||
}
|
||
|
||
// We leave the scheme "http" even for WebSockets – the Upgrade
|
||
// header does the rest. (Only cosmetic to change it to "ws”.)
|
||
target, err := url.Parse("http://" + ip + ":" + port)
|
||
if err != nil {
|
||
http.Error(w, "invalid ip/port: "+err.Error(), http.StatusBadRequest)
|
||
return
|
||
}
|
||
|
||
proxy := &httputil.ReverseProxy{
|
||
Director: func(req *http.Request) {
|
||
// Point to the agent
|
||
req.URL.Scheme = target.Scheme
|
||
req.URL.Host = target.Host
|
||
|
||
// Trim the first "/proxyAgent" prefix
|
||
if strings.HasPrefix(req.URL.Path, "/proxyAgent") {
|
||
req.URL.Path = strings.TrimPrefix(req.URL.Path, "/proxyAgent")
|
||
if req.URL.Path == "" {
|
||
req.URL.Path = "/"
|
||
}
|
||
}
|
||
|
||
// Scrub ip/port from downstream query
|
||
q := req.URL.Query()
|
||
q.Del("ip")
|
||
q.Del("port")
|
||
req.URL.RawQuery = q.Encode()
|
||
|
||
// Preserve Host (many CLI-style servers care)
|
||
req.Host = target.Host
|
||
},
|
||
|
||
// Critical tweaks for WebSockets
|
||
Transport: &http.Transport{
|
||
Proxy: http.ProxyFromEnvironment,
|
||
DialContext: (&net.Dialer{Timeout: 30 * time.Second, KeepAlive: 30 * time.Second}).DialContext,
|
||
ForceAttemptHTTP2: false, // MUST be HTTP/1.1 for WS
|
||
ResponseHeaderTimeout: 0,
|
||
},
|
||
|
||
FlushInterval: -1, // stream bytes immediately
|
||
}
|
||
|
||
proxy.ServeHTTP(w, r)
|
||
}
|
||
|
||
func main() {
|
||
|
||
// sqlite3 has been initialized in init()
|
||
defer logger.CloseDB()
|
||
|
||
defer db.Close()
|
||
|
||
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)
|
||
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()
|
||
log.Println("Server instance: ", 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")
|
||
}
|