From 1a6a2a64f2123964feb6d4e77255bc2e98e040e6 Mon Sep 17 00:00:00 2001
From: Alexander Mnich <56564725+a-mnich@users.noreply.github.com>
Date: Wed, 19 Feb 2025 06:19:49 +0100
Subject: [PATCH] implement agent startup arguments and help message (#581)

---
 beszel/cmd/agent/agent.go | 50 +++++++++++++++++++++++++++++++++------
 1 file changed, 43 insertions(+), 7 deletions(-)

diff --git a/beszel/cmd/agent/agent.go b/beszel/cmd/agent/agent.go
index 914378a..9ad4834 100644
--- a/beszel/cmd/agent/agent.go
+++ b/beszel/cmd/agent/agent.go
@@ -3,6 +3,7 @@ package main
 import (
 	"beszel"
 	"beszel/internal/agent"
+	"flag"
 	"fmt"
 	"log"
 	"os"
@@ -10,26 +11,53 @@ import (
 )
 
 func main() {
+	// Define flags for key and port
+	keyFlag := flag.String("key", "", "Public key")
+	portFlag := flag.String("port", "45876", "Port number")
+
+	flag.Usage = func() {
+		fmt.Printf("Usage: %s [options] [subcommand]\n", os.Args[0])
+		fmt.Println("\nOptions:")
+		flag.PrintDefaults()
+		fmt.Println("\nSubcommands:")
+		fmt.Println("  version      Display the version")
+		fmt.Println("  help         Display this help message")
+		fmt.Println("  update       Update the agent to the latest version")
+	}
+
+	// Parse the flags
+	flag.Parse()
+
 	// handle flags / subcommands
 	if len(os.Args) > 1 {
 		switch os.Args[1] {
-		case "-v":
+		case "version":
 			fmt.Println(beszel.AppName+"-agent", beszel.Version)
+			os.Exit(0)
+		case "help":
+			flag.Usage()
+			os.Exit(0)
 		case "update":
 			agent.Update()
+			os.Exit(0)
 		}
-		os.Exit(0)
 	}
 
-	// Try to get the key from the KEY environment variable.
-	key, _ := agent.GetEnv("KEY")
-	pubKey := []byte(key)
+	var pubKey []byte
+	// Override the key if the -key flag is provided
+	if *keyFlag != "" {
+		pubKey = []byte(*keyFlag)
+	} else {
+		// Try to get the key from the KEY environment variable.
+		key, _ := agent.GetEnv("KEY")
+		pubKey = []byte(key)
+	}
 
 	// If KEY is not set, try to read the key from the file specified by KEY_FILE.
 	if len(pubKey) == 0 {
 		keyFile, exists := agent.GetEnv("KEY_FILE")
 		if !exists {
-			log.Fatal("Must set KEY or KEY_FILE environment variable")
+			log.Fatal("Must set KEY or KEY_FILE environment variable or supply as input argument. Use 'beszel-agent help' for more information.")
 		}
 		var err error
 		pubKey, err = os.ReadFile(keyFile)
@@ -38,7 +66,10 @@ func main() {
 		}
 	}
 
-	addr := ":45876"
+	// Init with default port
+	addr := ":" + *portFlag
+	
+	//Use port from ENV if it exists
 	// TODO: change env var to ADDR
 	if portEnvVar, exists := agent.GetEnv("PORT"); exists {
 		// allow passing an address in the form of "127.0.0.1:45876"
@@ -48,5 +79,10 @@ func main() {
 		addr = portEnvVar
 	}
 
+	// Override the default and ENV port if the -port flag is provided and is non default
+	if *portFlag != "45876" {
+		addr = ":" + *portFlag
+	}
+
 	agent.NewAgent().Run(pubKey, addr)
 }