fix: random changes

main.go

@@ -3,55 +3,68 @@ package main
 import (
 	"bufio"
 	"fmt"
+	"log"
 	"net"
 	"os"
 	"os/exec"
+	"strings"
 )
 
-// ExecuteCommand runs a system command and returns its output
+var commands = map[string]string{
-func ExecuteCommand(command string) string {
+	"cpu":     `top -bn1 | grep "Cpu(s)" | sed "s/.*, *\([0-9.]*\)%* id.*/\1/" | sed "s/^/100 - /" | bc`,
+	"memory":  `free -m | awk '/Mem/{print $3 " " $2}'`,
+	"network": `cat /sys/class/net/[e]*/statistics/{r,t}x_bytes`,
+	"storage": `df -Ph | grep mmcblk0p5 | awk '{print $2 " " $3}' | sed 's/G//g'`,
+	"uptime":  `cut -f1 -d. /proc/uptime`,
+}
+
+// executeCommand runs a system command and returns its output
+func executeCommand(command string) string {
 	cmd := exec.Command("bash", "-c", command)
 	output, err := cmd.CombinedOutput()
 	if err != nil {
-		return fmt.Sprintf("Error: %s", err)
+		return err.Error()
 	}
+
 	return string(output)
 }
 
+// handleConnection handles one command per connection
 func handleConnection(conn net.Conn) {
 	defer conn.Close()
 
 	scanner := bufio.NewScanner(conn)
-	for scanner.Scan() {
+	if !scanner.Scan() {
-		command := scanner.Text()
+		log.Printf("Error reading from %s, %s", conn.RemoteAddr(), scanner.Err())
-
+		return
-		switch command {
-		case "cpu":
-			fmt.Fprintln(conn, ExecuteCommand("top -bn1 | grep \"Cpu(s)\" | sed \"s/.*, *\\([0-9.]*\\)%* id.*/\\1/\" | sed \"s/^/100 - /\" | bc"))
-		case "memory":
-			fmt.Fprintln(conn, ExecuteCommand("free -m | awk '/Mem/{print $3\" \"$2}'"))
-		case "network":
-			fmt.Fprintln(conn, ExecuteCommand("cat /sys/class/net/[e]*/statistics/{r,t}x_bytes"))
-		case "storage":
-			fmt.Fprintln(conn, ExecuteCommand("df -Ph | grep mmcblk0p5 | awk '{print $2\" \"$3}' | sed 's/G//g'"))
-		case "uptime":
-			fmt.Fprintln(conn, ExecuteCommand("cut -f1 -d. /proc/uptime"))
-		case "exit":
-			return
-		default:
-			fmt.Fprintln(conn, "Invalid command")
-		}
 	}
+
+	command := scanner.Text()
+
+	cmd, valid := commands[strings.TrimSpace(string(command))]
+	if !valid {
+		fmt.Fprintln(conn, "invalid command")
+		return
+	}
+
+	output := executeCommand(cmd)
+	fmt.Fprintln(conn, output)
 }
 
 func main() {
-	ln, err := net.Listen("tcp", ":12345")
+	host, ok := os.LookupEnv("HOST")
+	if !ok {
+		host = ":12345"
+	}
+
+	ln, err := net.Listen("tcp", host)
 	if err != nil {
 		fmt.Println("Error:", err)
 		os.Exit(1)
 	}
 	defer ln.Close()
 
+	log.Printf("Listening on %s...", host)
 	for {
 		conn, err := ln.Accept()
 		if err != nil {
@@ -59,6 +72,7 @@ func main() {
 			continue
 		}
 
+		log.Printf("Connection from %s", conn.RemoteAddr())
 		go handleConnection(conn)
 	}
 }