gontrol/main.go

package main

import (
	"context"
	"encoding/json"
	"embed"
	"os"
	"strings"
	"time"
	"slices"
	"database/sql"
	"fmt"
	"html/template"
	"log"
	"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



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)
	}

	// Agents database
	db = database.InitSQLiteDB("/tmp/gontrol_agents.db")
}

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
// 	}
// 	agentPort := r.URL.Query().Get("port")
// 	if agentIP == "" {
// 		http.Error(w, "Missing 'port' parameter", http.StatusBadRequest)
// 		return
// 	}

// 	// Construct the URL to proxy to
// 	agentURL := "http://" + agentIP + ":" + agentPort

// 	// 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 <link> 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 <head> section
// 	w.Header().Set("Content-Type", "text/html")
// 	w.WriteHeader(http.StatusOK)

// 	// Inject the stylesheets into the <head> section of your page
// 	fmt.Fprintf(w, "<html><head>")
// 	for _, stylesheet := range stylesheets {
// 		// Make sure the stylesheet is loaded properly (absolute URLs or proxy it)
// 		fmt.Fprintf(w, `<link rel="stylesheet" href="%s">`, stylesheet)
// 	}
// 	fmt.Fprintf(w, "</head><body>")

// 	// 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, `
// 		<iframe src="%s" width="100%" height="800px" style="border:none;">
// 			Your browser does not support iframes.
// 		</iframe>
// 	`, agentURL)

// 	fmt.Fprintf(w, "</body></html>")
// }

// proxyAgentHandler proxies requests to the agent server specified by IP and port query parameters.
// It strips the "/proxyAgent" prefix from the request path before forwarding.
func proxyAgentHandler(w http.ResponseWriter, r *http.Request) {
	agentIP := r.URL.Query().Get("ip")
	if agentIP == "" {
		http.Error(w, "Missing 'ip' parameter", http.StatusBadRequest)
		return
	}
	agentPort := r.URL.Query().Get("port")
	if agentPort == "" {
		http.Error(w, "Missing 'port' parameter", http.StatusBadRequest)
		return
	}

	agentBaseURL := "http://" + agentIP + ":" + agentPort

	targetURL, err := url.Parse(agentBaseURL)
	if err != nil {
		http.Error(w, "Invalid agent URL: "+err.Error(), http.StatusBadRequest)
		return
	}

	proxy := httputil.NewSingleHostReverseProxy(targetURL)

	// Modify the Director function to rewrite the request URL before sending to agent
	originalDirector := proxy.Director
	proxy.Director = func(req *http.Request) {
		originalDirector(req)

		// Strip "/proxyAgent" prefix from the path
		const prefix = "/proxyAgent"
		if strings.HasPrefix(req.URL.Path, prefix) {
			req.URL.Path = strings.TrimPrefix(req.URL.Path, prefix)
			if req.URL.Path == "" {
				req.URL.Path = "/"
			}
		}

		// Preserve original query parameters except ip and port (remove those for backend)
		query := req.URL.Query()
		query.Del("ip")
		query.Del("port")
		req.URL.RawQuery = query.Encode()

		// Optional: set the Host header to the target host
		req.Host = targetURL.Host
	}

	proxy.ServeHTTP(w, r)
}

// func proxyAgentHandler(w http.ResponseWriter, r *http.Request) {
//     agentIP := r.URL.Query().Get("ip")
//     if agentIP == "" {
//         http.Error(w, "Missing 'ip' parameter", http.StatusBadRequest)
//         return
//     }
//     agentPort := r.URL.Query().Get("port")
//     if agentPort == "" {
//         http.Error(w, "Missing 'port' parameter", http.StatusBadRequest)
//         return
//     }

//     agentURL := "http://" + agentIP + ":" + agentPort + "/"

//     resp, err := http.Get(agentURL)
//     if err != nil {
//         http.Error(w, "Failed to reach agent: "+err.Error(), http.StatusBadGateway)
//         return
//     }
//     defer resp.Body.Close()

//     // Copy headers from agent response (you might want to filter some)
//     for k, v := range resp.Header {
//         for _, vv := range v {
//             w.Header().Add(k, vv)
//         }
//     }
//     w.WriteHeader(resp.StatusCode)

//     // Stream the entire response body directly to the client
//     io.Copy(w, resp.Body)
// }


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)

	// 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")
}