chore: better io handling and minor style enhancements

main.go

@@ -1,13 +1,22 @@
 package main
 
 import (
-	"bufio"
 	"fmt"
+	"io"
 	"net"
 	"os"
 	"os/exec"
+	"strings"
 )
 
+var commands = map[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)
@@ -18,34 +27,33 @@ func ExecuteCommand(command string) string {
 	return string(output)
 }
 
+// handleConnection handles one command per connection
 func handleConnection(conn net.Conn) {
 	defer conn.Close()
 
-	scanner := bufio.NewScanner(conn)
-	for scanner.Scan() {
-		command := scanner.Text()
-
-		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, err := io.ReadAll(conn)
+	if err != nil {
+		fmt.Println(err)
+		return
 	}
+
+	cmd, valid := commands[strings.TrimSpace(string(command))]
+	if !valid {
+		fmt.Fprintln(conn, "Invalid command")
+		return
+	}
+
+	stdout := ExecuteCommand(cmd)
+	fmt.Fprintln(conn, stdout)
 }
 
 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)