rcon-cli/main.go

package main

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"net"
	"os"
	"runtime"
	"strings"

	"github.com/hamburghammer/rcon"
	flags "github.com/spf13/pflag"
)

const (
	defaultHost     = "localhost"
	defaultPort     = "25575"
	defaultPassword = ""

	envVarPrefix = "RCON_CLI_"
)

// resetColor is the ASCII code for resetting the color.
const resetColor = "\u001B[0m"

// colors a map with the ASCII color codes for Unix terms.
var colors = map[string]string{
	"0": "\u001B[30m", // black
	"1": "\u001B[34m", // dark blue
	"2": "\u001B[32m", // dark green
	"3": "\u001B[36m", // dark aqua
	"4": "\u001B[31m", // dark red
	"5": "\u001B[35m", // dark purple
	"6": "\u001B[33m", // gold
	"7": "\u001B[37m", // gray
	"8": "\u001B[30m", // dark gray
	"9": "\u001B[34m", // blue
	"a": "\u001B[32m", // green
	"b": "\u001B[32m", // aqua
	"c": "\u001B[31m", // red
	"d": "\u001B[35m", // light purple
	"e": "\u001B[33m", // yellow
	"f": "\u001B[37m", // white
	"k": "",           // random
	"m": "\u001B[9m",  // strikethrough
	"o": "\u001B[3m",  // italic
	"l": "\u001B[1m",  // bold
	"n": "\u001B[4m",  // underline
	"r": resetColor,   // reset
}

// vars holding the parsed configuration
var (
	host     string
	port     string
	password string
	help     bool
)

func main() {
	commands, _ := parseArgs(os.Args[1:])
	loadEnvIfNotSet()
	if help {
		printHelp()
		os.Exit(0)
		return
	}

	err := run(strings.Join(commands, " "))
	if err != nil {
		_, err = fmt.Fprintln(os.Stderr, err.Error())
		if err != nil {
			panic(err)
		}
		os.Exit(1)
	}
}

func parseArgs(args []string) ([]string, error) {
	flags.StringVar(&host, "host", defaultHost, "RCON server's hostname.")
	flags.StringVar(&port, "port", defaultPort, "RCON server's port.")
	flags.StringVar(&password, "password", defaultPassword, "RCON server's password.")
	flags.BoolVarP(&help, "help", "h", false, "Prints this help message and exits.")

	err := flags.CommandLine.Parse(args)
	if err != nil {
		return []string{}, err
	}

	command := flags.Args()

	return command, nil
}

func loadEnvIfNotSet() {
	envHost := os.Getenv(fmt.Sprintf("%sHOST", envVarPrefix))
	if envHost != "" && host == defaultHost {
		host = envHost
	}

	envPort := os.Getenv(fmt.Sprintf("%sPORT", envVarPrefix))
	if envPort != "" && port == defaultPort {
		port = envPort
	}

	envPassword := os.Getenv(fmt.Sprintf("%sPASSWORD", envVarPrefix))
	if envPassword != "" && password == defaultPassword {
		password = envPassword
	}
}

func run(cmd string) error {
	hostPort := net.JoinHostPort(host, port)
	remoteConsole, err := rcon.Dial(hostPort, password)
	if err != nil {
		return fmt.Errorf("failed to connect to RCON server: %w", err)
	}
	defer remoteConsole.Close()

	if len(cmd) == 0 {
		err = executeInteractive(remoteConsole)
	} else {
		resp, err := execute(cmd, remoteConsole)
		if err != nil {
			return err
		}
		_, err = fmt.Println(resp)
	}

	return err
}

func printHelp() {
	_, _ = fmt.Println(`rcon-cli is a CLI to interact with a RCON server.
It can be run in an interactive mode or to execute a single command.

USAGE:
	rcon-cli [FLAGS] [RCON command ...]
	
FLAGS:`)
	flags.PrintDefaults()
	fmt.Printf(`
ENVIRONMENT VARIABLE:
	All flags can be set through the flag name in capslock with the %s prefix (see examples).
	Flags have allways priority over env vars!

EXAMPLES:
	rcon-cli --host 127.0.0.1 --port 25575
	rcon-cli --password admin123 stop
	%sPORT=25575 rcon-cli stop
`, envVarPrefix, envVarPrefix)
}

func executeInteractive(remoteConsole *rcon.RemoteConsole) error {
	scanner := bufio.NewScanner(os.Stdin)
	_, _ = fmt.Println("To quit the session type 'exit'.")
	_, _ = fmt.Print("> ")
	for scanner.Scan() {
		cmd := scanner.Text()
		if cmd == "exit" {
			return nil
		}
		resp, err := execute(cmd, remoteConsole)
		if err != nil {
			return err
		}

		_, _ = fmt.Println(resp)
		_, _ = fmt.Print("> ")
	}

	if err := scanner.Err(); err != nil {
		return fmt.Errorf("reading standard input: %w", err)
	}

	return nil
}

func execute(cmd string, remoteConsole *rcon.RemoteConsole) (string, error) {
	resp := ""
	reqID, err := remoteConsole.Write(cmd)
	if err != nil {
		return resp, fmt.Errorf("failed to send command: %w", err)
	}

	resp, respID, err := remoteConsole.Read()
	if err != nil {
		if err == io.EOF {
			return resp, nil
		}
		return resp, fmt.Errorf("failed to read command: %w", err)
	}

	if reqID != respID {
		return resp, errors.New("the response id didn't match the request id")
	}

	resp = colorize(resp)

	return resp, nil
}

// colorize tries to add the color codes for the terminal.
// works on none windows machines.
func colorize(str string) string {
	const sectionSign = "§"

	for code := range colors {
		if runtime.GOOS == "windows" {
			str = strings.ReplaceAll(str, sectionSign+code, "")
		} else {
			str = strings.ReplaceAll(str, sectionSign+code, colors[code])
		}
	}
	// reset color after each new line
	if runtime.GOOS != "windows" {
		return strings.ReplaceAll(str, "\n", "\n"+resetColor)
	}

	return str
}