Files
beszel-ipv6/beszel/cmd/agent/agent.go
2025-08-02 17:04:38 -04:00

127 lines
2.9 KiB
Go

package main
import (
"beszel"
"beszel/internal/agent"
"beszel/internal/agent/health"
"flag"
"fmt"
"log"
"os"
"strings"
"golang.org/x/crypto/ssh"
)
// cli options
type cmdOptions struct {
key string // key is the public key(s) for SSH authentication.
listen string // listen is the address or port to listen on.
}
// parse parses the command line flags and populates the config struct.
// It returns true if a subcommand was handled and the program should exit.
func (opts *cmdOptions) parse() bool {
flag.StringVar(&opts.key, "key", "", "Public key(s) for SSH authentication")
flag.StringVar(&opts.listen, "listen", "", "Address or port to listen on")
flag.Usage = func() {
builder := strings.Builder{}
builder.WriteString("Usage: ")
builder.WriteString(os.Args[0])
builder.WriteString(" [command] [flags]\n")
builder.WriteString("\nCommands:\n")
builder.WriteString(" health Check if the agent is running\n")
builder.WriteString(" help Display this help message\n")
builder.WriteString(" update Update to the latest version\n")
builder.WriteString("\nFlags:\n")
fmt.Print(builder.String())
flag.PrintDefaults()
}
subcommand := ""
if len(os.Args) > 1 {
subcommand = os.Args[1]
}
switch subcommand {
case "-v", "version":
fmt.Println(beszel.AppName+"-agent", beszel.Version)
return true
case "help":
flag.Usage()
return true
case "update":
agent.Update()
return true
case "health":
err := health.Check()
if err != nil {
log.Fatal(err)
}
fmt.Print("ok")
return true
}
flag.Parse()
return false
}
// loadPublicKeys loads the public keys from the command line flag, environment variable, or key file.
func (opts *cmdOptions) loadPublicKeys() ([]ssh.PublicKey, error) {
// Try command line flag first
if opts.key != "" {
return agent.ParseKeys(opts.key)
}
// Try environment variable
if key, ok := agent.GetEnv("KEY"); ok && key != "" {
return agent.ParseKeys(key)
}
// Try key file
keyFile, ok := agent.GetEnv("KEY_FILE")
if !ok {
return nil, fmt.Errorf("no key provided: must set -key flag, KEY env var, or KEY_FILE env var. Use 'beszel-agent help' for usage")
}
pubKey, err := os.ReadFile(keyFile)
if err != nil {
return nil, fmt.Errorf("failed to read key file: %w", err)
}
return agent.ParseKeys(string(pubKey))
}
func (opts *cmdOptions) getAddress() string {
return agent.GetAddress(opts.listen)
}
func main() {
var opts cmdOptions
subcommandHandled := opts.parse()
if subcommandHandled {
return
}
var serverConfig agent.ServerOptions
var err error
serverConfig.Keys, err = opts.loadPublicKeys()
if err != nil {
log.Fatal("Failed to load public keys:", err)
}
addr := opts.getAddress()
serverConfig.Addr = addr
serverConfig.Network = agent.GetNetwork(addr)
a, err := agent.NewAgent()
if err != nil {
log.Fatal("Failed to create agent: ", err)
}
if err := a.Start(serverConfig); err != nil {
log.Fatal("Failed to start server: ", err)
}
}