go-stats-server/main.go

package main

import (
	"bufio"
	"fmt"
	"log"
	"net"
	"os"
	"os/exec"
	"strings"
)

var systemdServiceUnit = strings.TrimSpace(`
[Unit]
Description=Stats Server
After=network.target

[Service]
Type=simple
Restart=always
RestartSec=15
ExecStart=/usr/local/bin/stats-server serve

[Install]
WantedBy=default.target
`)

// cmdStrHelper removes indentation and newlines from a shell command string for
// better readability.
//
// Precisely, it first trims the string, then splits it by newlines, then trims
// each line, then joins them back with a space.
func cmdStrHelper(command string) string {
	lines := strings.Split(strings.TrimSpace(command), "\n")
	for i, line := range lines {
		lines[i] = strings.TrimSpace(line)
	}

	return strings.Join(lines, " ")
}

// commands is a map of commands that can be run on the server, originally made
// by @BachoSeven
var commands = map[string]string{
	"cpu": cmdStrHelper(`
		top -bn1
		| grep "Cpu(s)"
		| sed "s/.*, *\([0-9.]*\)%* id.*/\1/"
		| sed "s/^/100 - /"
		| bc
	`),
	"memory": cmdStrHelper(`
		free -m
		| awk '/Mem/{print $3 " " $2}'
	`),
	"network": cmdStrHelper(`
		cat /sys/class/net/[e]*/statistics/{r,t}x_bytes
	`),
	"storage": cmdStrHelper(`
		df -Ph
		| grep mmcblk0p5
		| awk '{print $2 " " $3}'
		| sed 's/G//g'
	`),
	"uptime": cmdStrHelper(`
		cut -f1 -d. /proc/uptime
	`),
	"temperature": cmdStrHelper(`
		paste
			<(cat /sys/class/thermal/thermal_zone*/type)
			<(awk
				'{printf "%.3f\n", $1/1000}'
				/sys/class/thermal/thermal_zone*/temp
			)
	`),
}

func init() {
	log.SetFlags(0)
}

// runShellCommand runs a system command and returns its output
func runShellCommand(command string) (string, error) {
	cmd := exec.Command("bash", "-c", command)
	output, err := cmd.CombinedOutput()
	if err != nil {
		return string(output), err
	}

	return string(output), nil
}

// handleConnection handles one command per connection
func handleConnection(conn net.Conn) {
	defer conn.Close()

	scanner := bufio.NewScanner(conn)
	if !scanner.Scan() {
		log.Printf("error reading from %s, %s", conn.RemoteAddr(), scanner.Err())
		return
	}

	command := scanner.Text()
	log.Printf("received command %s from %s", command, conn.RemoteAddr())

	shellCmd, valid := commands[strings.TrimSpace(string(command))]
	if !valid {
		fmt.Fprintln(conn, "invalid command")
		return
	}

	output, err := runShellCommand(shellCmd)
	if err != nil {
		log.Fatal(err)
		return
	}

	fmt.Fprintln(conn, output)
}

func runCommand(args ...string) error {
	cmd := exec.Command(args[0], args[1:]...)

	if output, err := cmd.CombinedOutput(); err != nil {
		return fmt.Errorf("%w: %s", err, output)
	}

	return nil
}

func setupSystemdService() error {
	if _, err := os.Stat("/usr/local/bin/stats-server"); err == nil {
		log.Println("binary already exists in /usr/local/bin, updating...")
	}

	log.Println("copying itself to /usr/local/bin/stats-server")
	if err := runCommand("cp", "-f", os.Args[0], "/usr/local/bin/stats-server"); err != nil {
		return err
	}

	log.Println("writing unit to /etc/systemd/system/stats-server.service")
	if err := os.WriteFile(
		"/etc/systemd/system/stats-server.service",
		[]byte(systemdServiceUnit),
		os.ModePerm,
	); err != nil {
		return err
	}

	log.Println("running systemctl daemon-reload")
	if err := runCommand("systemctl", "daemon-reload"); err != nil {
		return err
	}

	log.Println("running systemctl enable stats-server.service")
	if err := runCommand("systemctl", "enable", "stats-server.service"); err != nil {
		return err
	}

	log.Println("running systemctl (re)start stats-server.service")
	if err := runCommand("systemctl", "restart", "stats-server.service"); err != nil {
		return err
	}

	return nil
}

func startTCPServer() error {
	host, ok := os.LookupEnv("HOST")
	if !ok {
		host = ":12345"
	}

	ln, err := net.Listen("tcp", host)
	if err != nil {
		return err
	}
	defer ln.Close()

	log.Printf("listening on %s...", host)
	for {
		conn, err := ln.Accept()
		if err != nil {
			log.Printf("error accepting connection: %s", err)
			continue
		}

		log.Printf("connection from %s", conn.RemoteAddr())
		go handleConnection(conn)
	}
}

var helpText = strings.TrimSpace(`
usage: stats-server [setup|serve]

This is a simple tcp server that returns system stats, it can be run as a
systemd service or as a standalone server.

protocol commands:
    cpu      returns cpu usage
    memory   returns memory usage
    network  returns network usage
    storage  returns storage usage
    uptime   returns uptime

subcommands:
    setup    auto-install and setup systemd service
    serve    start tcp server

config, environment variables:
    HOST     tcp host to bind to (default: :12345)
`)

func showHelp() {
	fmt.Println(helpText)
	os.Exit(1)
}

func main() {
	if len(os.Args) != 2 {
		showHelp()
	}

	switch os.Args[1] {
	case "setup":
		log.Println("starting setup")
		if err := setupSystemdService(); err != nil {
			log.Fatal(err)
		}

		log.Println("setup completed")

	case "serve":
		log.Println("starting server")
		if err := startTCPServer(); err != nil {
			log.Fatal(err)
		}

		log.Println("server exited")

	default:
		showHelp()
	}
}