gontrol/main.go

package main

import (
	"context"
	"encoding/json"
	"os"
	"strings"
	"time"

	// "errors"
	"fmt"
	// "io"
	// "io/ioutil"

	// "net"
	"database/sql"
	"html/template"
	"log"
	"net/http"
	"sync"
	"syscall"

	"os/signal"

	_ "github.com/go-sql-driver/mysql"
	"github.com/gorilla/websocket"
)

var tmpl *template.Template
var db 	 *sql.DB

const (
	keyServerAddr = "serverAddr"
	dbHost = "127.0.0.1"
	dbUser = "mysql"
	dbPort = 3306
	dbPassword = "mysql"
	dbName = "gomatic"
)

type Agent struct {
	AgentID        int    `json:"agentId"`
	AgentName      string `json:"agentName"`
	InitialContact string `json:"initialContact"`
	LastContact    string `json:"lastContact"`
	IPv4Address    string `json:"IPv4Address"`
	Status         string `json:"status"`
}

func init() {
	tmpl, _ = template.ParseGlob("templates/*.html")
}

func initDB () {
	var err  error
	dbOpen := fmt.Sprintf("%s:%s@(%s:%d)/%s?parseTime=true", dbUser,dbPassword, dbHost, dbPort, dbName)
	db, err = sql.Open("mysql", dbOpen)
	if err != nil {
		log.Fatal(err)
	}

	if err = db.Ping(); err != nil {
		log.Fatal(err)
	}
	log.Println("Connected to database")
}

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:
		deleteAgent(w, r, agentId)
	case http.MethodGet:
		if agentId == "" {
			listAgents(w, r)
		} else {
			getAgent(w, r, agentId)
		}
	case http.MethodPost:
		createAgent(w, r)
	case http.MethodPut:
		updateAgent(w, r, agentId)
	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) {
	agents, err := getAgents()
	if err != nil {
		http.Error(w, "Failed to fetch agents", http.StatusInternalServerError)
		return
	}

	renderTemplate(w, "templates/partials/agent_list.html", agents)
}

func deleteAgent(w http.ResponseWriter, r *http.Request, agentID string) {
	query := "DELETE FROM agents WHERE agentId = ?"
	_, err := db.Exec(query, agentID)
	if err != nil {
		http.Error(w, "Failed to delete agent", http.StatusInternalServerError)
		return
	}

	agents, err := getAgents()
	if err != nil {
		http.Error(w, "Failed to fetch agents", http.StatusInternalServerError)
		return
	}

	renderTemplate(w, "templates/partials/agent_list.html", agents)
}

func createAgent(w http.ResponseWriter, r * http.Request) {
	err := r.ParseForm()
	if err != nil {
		http.Error(w, "Invalid form data", http.StatusBadRequest)
	}

	agentName := r.FormValue("agentName")
	agentId := r.FormValue("agentId")
	IPv4Address := r.FormValue("IPv4Address")
	// initalContact := r.FormValue("initialContact")
	// lastContact := r.FormValue("lastContact")

	query := "INSERT INTO agents (agentId, agentName, IPv4Address, initialContact, lastContact) VALUES (?, ?, ?, NOW(), NOW())"
	_, err = db.Exec(query, agentId, agentName, IPv4Address)
	if err != nil {
		http.Error(w, "Failed to create agent", http.StatusInternalServerError)
		return
	}

	w.WriteHeader(http.StatusCreated)
	listAgents(w, r)
}

func updateAgent(w http.ResponseWriter, r *http.Request, agentId string) {
	err := r.ParseForm()
	if err != nil {
		http.Error(w, "Invalid form data", http.StatusBadRequest)
		return
	}

	agentName := r.FormValue("agentName")
	lastContact := r.FormValue("lastContact")
	IPv4Address := r.FormValue("IPv4Address")

	query := "UPDATE agents SET agentName = ?, IPv4Address = ?, lastContact = ? where agentId = ?"
	_, err = db.Exec(query, agentName, IPv4Address, lastContact, agentId)
	if err != nil {
		http.Error(w, "Failed to update agent", http.StatusInternalServerError)
		return
	}

	listAgents(w, r,)
}

func getAgents() ([]Agent, error) {
	query := "SELECT agentId, agentName, IPv4Address, initialContact, lastContact FROM agents"
	rows, err := db.Query(query)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var agents []Agent
	for rows.Next() {
		var agent Agent
		err := rows.Scan(&agent.AgentID, &agent.AgentName, &agent.IPv4Address, &agent.InitialContact, &agent.LastContact)
		if err != nil {
			return nil, err
		}
		agents = append(agents, agent)
	}
	return agents, rows.Err()
}

func getAgent(w http.ResponseWriter, r *http.Request, agentId string) {
	query := "Select agentId, agentName, initialContact, lastContact from agents where agentId = ?"
	var agent Agent
	err := db.QueryRow(query, agentId).Scan(&agent.AgentID, &agent.AgentName, &agent.InitialContact, &agent.LastContact)
	if err == sql.ErrNoRows {
		http.Error(w, "Agent not found", http.StatusNotFound)
		return
	} else if err != nil {
		http.Error(w, "Failed to fetch agent", http.StatusInternalServerError)
		return
	}

	renderTemplate(w, "templates/partials/agent_detail.html", agent)

}

func getAgentNames(w http.ResponseWriter, r *http.Request) {
	query := "SELECT agentName from agents"
	rows, err := db.Query(query)
	if err != nil {
		http.Error(w, "Failed to fetch agent names", http.StatusInternalServerError)
		return
	}
	defer rows.Close()

	var agentNames []string
	for rows.Next() {
		var name string
		if err := rows.Scan(&name); err != nil {
			http.Error(w, "Error reading agent names", http.StatusInternalServerError)
			return
		}
		agentNames = append(agentNames, name)
	}

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(agentNames)
}

type webSocketHandler struct {
	upgrader websocket.Upgrader
}

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)
		return
	}
	defer c.Close()

	for {
		mt, message, err := c.ReadMessage()
		if err != nil {
			log.Printf("Error reading: message: %s", err)
		}

		log.Printf("Received message: %s", message)
		if err = c.WriteMessage(mt, message); err !=nil {
			log.Printf("Error writing the message: %s", err)
			break
		}
	}
}

var agentSockets = make(map[string]*websocket.Conn)
var agentSocketsMutex sync.Mutex

var executeCommand http.HandlerFunc = func(w http.ResponseWriter, r *http.Request){
	err := r.ParseForm()
	if err != nil {
		http.Error(w, "Invalid form data", http.StatusBadRequest)
	}

	agentName := r.FormValue("agentName")
	command := r.FormValue("command")

	agentSocketsMutex.Lock()
	conn, ok := agentSockets[agentName]
	agentSocketsMutex.Unlock()

	if !ok {
		http.Error(w, "Agent not connected", http.StatusNotFound)
		return
	}

	err = conn.WriteMessage(websocket.TextMessage, []byte(command))
	if err != nil {
		http.Error(w, "Failed to send command to the agent", http.StatusInternalServerError)
		return
	}

	w.WriteHeader(http.StatusOK)
	w.Write([]byte("Command sent successfully"))
}

func main() {

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	var wg sync.WaitGroup

	webSocketHandler := webSocketHandler {
		upgrader: websocket.Upgrader{},
	}

	webSocketMux := http.NewServeMux()
	webSocketMux.Handle("/data", webSocketHandler)
	webSocketMux.Handle("/executeCommand", executeCommand)
	websocketServer := &http.Server{
		Addr: ":5555",
		Handler: webSocketMux,
	}

	webMux := http.NewServeMux()
	webMux.HandleFunc("/", getHomepage)
	// webMux.HandleFunc("/index", getRoot)
	// webMux.HandleFunc("/hello", getHello)
	// webMux.HandleFunc("/agents", fetchAgents)
	webMux.HandleFunc("/agents", agentsHandler)
	webMux.HandleFunc("/agentNames", getAgentNames)
	webMux.HandleFunc("/agents/{agentId}", agentsHandler)

	initDB()
	defer db.Close()

	webServer := &http.Server {
		Addr: ":3333",
		Handler: webMux,
		// BaseContext: func(l net.Listener) context.Context {
		// 	ctx = context.WithValue(ctx, keyServerAddr, l.Addr().String())
		// 	return ctx
		// },
	}

	wg.Add(1)
	go func() {
		defer wg.Done()
		log.Println("Websocket server is running on port 5555")
		if err := websocketServer.ListenAndServe(); err != http.ErrServerClosed {
			log.Fatalf("Websocket server failed: %s", err)
		}
	}()

	wg.Add(1)
	go func() {
		defer wg.Done()
		log.Println("Web server is running on port 3333")

		if err := webServer.ListenAndServe(); err != http.ErrServerClosed {
			log.Fatalf("Web server failed: %s", 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")
}