mirror of
https://github.com/fankes/beszel.git
synced 2025-10-19 09:49:28 +08:00
rename /src to /internal (sorry i'll fix the prs)
This commit is contained in:
157
internal/cmd/agent/agent.go
Normal file
157
internal/cmd/agent/agent.go
Normal file
@@ -0,0 +1,157 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"log"
|
||||
"os"
|
||||
"strings"
|
||||
|
||||
"github.com/henrygd/beszel"
|
||||
"github.com/henrygd/beszel/agent"
|
||||
"github.com/henrygd/beszel/agent/health"
|
||||
"github.com/spf13/pflag"
|
||||
"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.
|
||||
// TODO: add hubURL and token
|
||||
// hubURL string // hubURL is the URL of the hub to use.
|
||||
// token string // token is the token to use for authentication.
|
||||
}
|
||||
|
||||
// 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 {
|
||||
subcommand := ""
|
||||
if len(os.Args) > 1 {
|
||||
subcommand = os.Args[1]
|
||||
}
|
||||
|
||||
// Subcommands that don't require any pflag parsing
|
||||
switch subcommand {
|
||||
case "-v", "version":
|
||||
fmt.Println(beszel.AppName+"-agent", beszel.Version)
|
||||
return true
|
||||
case "health":
|
||||
err := health.Check()
|
||||
if err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
fmt.Print("ok")
|
||||
return true
|
||||
}
|
||||
|
||||
// pflag.CommandLine.ParseErrorsWhitelist.UnknownFlags = true
|
||||
pflag.StringVarP(&opts.key, "key", "k", "", "Public key(s) for SSH authentication")
|
||||
pflag.StringVarP(&opts.listen, "listen", "l", "", "Address or port to listen on")
|
||||
// pflag.StringVarP(&opts.hubURL, "hub-url", "u", "", "URL of the hub to use")
|
||||
// pflag.StringVarP(&opts.token, "token", "t", "", "Token to use for authentication")
|
||||
chinaMirrors := pflag.BoolP("china-mirrors", "c", false, "Use mirror for update (gh.beszel.dev) instead of GitHub")
|
||||
help := pflag.BoolP("help", "h", false, "Show this help message")
|
||||
|
||||
// Convert old single-dash long flags to double-dash for backward compatibility
|
||||
flagsToConvert := []string{"key", "listen"}
|
||||
for i, arg := range os.Args {
|
||||
for _, flag := range flagsToConvert {
|
||||
singleDash := "-" + flag
|
||||
doubleDash := "--" + flag
|
||||
if arg == singleDash {
|
||||
os.Args[i] = doubleDash
|
||||
break
|
||||
} else if strings.HasPrefix(arg, singleDash+"=") {
|
||||
os.Args[i] = doubleDash + arg[len(singleDash):]
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pflag.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())
|
||||
pflag.PrintDefaults()
|
||||
}
|
||||
|
||||
// Parse all arguments with pflag
|
||||
pflag.Parse()
|
||||
|
||||
// Must run after pflag.Parse()
|
||||
switch {
|
||||
case *help || subcommand == "help":
|
||||
pflag.Usage()
|
||||
return true
|
||||
case subcommand == "update":
|
||||
agent.Update(*chinaMirrors)
|
||||
return true
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
}
|
336
internal/cmd/agent/agent_test.go
Normal file
336
internal/cmd/agent/agent_test.go
Normal file
@@ -0,0 +1,336 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"crypto/ed25519"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"github.com/henrygd/beszel/agent"
|
||||
|
||||
"github.com/spf13/pflag"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
"golang.org/x/crypto/ssh"
|
||||
)
|
||||
|
||||
func TestGetAddress(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
opts cmdOptions
|
||||
envVars map[string]string
|
||||
expected string
|
||||
}{
|
||||
{
|
||||
name: "default port when no config",
|
||||
opts: cmdOptions{},
|
||||
expected: ":45876",
|
||||
},
|
||||
{
|
||||
name: "use address from flag",
|
||||
opts: cmdOptions{
|
||||
listen: "8080",
|
||||
},
|
||||
expected: ":8080",
|
||||
},
|
||||
{
|
||||
name: "use unix socket from flag",
|
||||
opts: cmdOptions{
|
||||
listen: "/tmp/beszel.sock",
|
||||
},
|
||||
expected: "/tmp/beszel.sock",
|
||||
},
|
||||
{
|
||||
name: "use LISTEN env var",
|
||||
opts: cmdOptions{},
|
||||
envVars: map[string]string{
|
||||
"LISTEN": "1.2.3.4:9090",
|
||||
},
|
||||
expected: "1.2.3.4:9090",
|
||||
},
|
||||
{
|
||||
name: "use legacy PORT env var",
|
||||
opts: cmdOptions{},
|
||||
envVars: map[string]string{
|
||||
"PORT": "7070",
|
||||
},
|
||||
expected: ":7070",
|
||||
},
|
||||
{
|
||||
name: "use unix socket from env var",
|
||||
opts: cmdOptions{
|
||||
listen: "",
|
||||
},
|
||||
envVars: map[string]string{
|
||||
"LISTEN": "/tmp/beszel.sock",
|
||||
},
|
||||
expected: "/tmp/beszel.sock",
|
||||
},
|
||||
{
|
||||
name: "flag takes precedence over env vars",
|
||||
opts: cmdOptions{
|
||||
listen: ":8080",
|
||||
},
|
||||
envVars: map[string]string{
|
||||
"LISTEN": ":9090",
|
||||
"PORT": "7070",
|
||||
},
|
||||
expected: ":8080",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
// Setup environment
|
||||
for k, v := range tt.envVars {
|
||||
t.Setenv(k, v)
|
||||
}
|
||||
|
||||
addr := tt.opts.getAddress()
|
||||
assert.Equal(t, tt.expected, addr)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadPublicKeys(t *testing.T) {
|
||||
// Generate a test key
|
||||
_, priv, err := ed25519.GenerateKey(nil)
|
||||
require.NoError(t, err)
|
||||
signer, err := ssh.NewSignerFromKey(priv)
|
||||
require.NoError(t, err)
|
||||
pubKey := ssh.MarshalAuthorizedKey(signer.PublicKey())
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
opts cmdOptions
|
||||
envVars map[string]string
|
||||
setupFiles map[string][]byte
|
||||
wantErr bool
|
||||
errContains string
|
||||
}{
|
||||
{
|
||||
name: "load key from flag",
|
||||
opts: cmdOptions{
|
||||
key: string(pubKey),
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "load key from env var",
|
||||
envVars: map[string]string{
|
||||
"KEY": string(pubKey),
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "load key from file",
|
||||
envVars: map[string]string{
|
||||
"KEY_FILE": "testkey.pub",
|
||||
},
|
||||
setupFiles: map[string][]byte{
|
||||
"testkey.pub": pubKey,
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "error when no key provided",
|
||||
wantErr: true,
|
||||
errContains: "no key provided",
|
||||
},
|
||||
{
|
||||
name: "error on invalid key file",
|
||||
envVars: map[string]string{
|
||||
"KEY_FILE": "nonexistent.pub",
|
||||
},
|
||||
wantErr: true,
|
||||
errContains: "failed to read key file",
|
||||
},
|
||||
{
|
||||
name: "error on invalid key data",
|
||||
opts: cmdOptions{
|
||||
key: "invalid-key-data",
|
||||
},
|
||||
wantErr: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
// Create a temporary directory for test files
|
||||
if len(tt.setupFiles) > 0 {
|
||||
tmpDir := t.TempDir()
|
||||
for name, content := range tt.setupFiles {
|
||||
path := filepath.Join(tmpDir, name)
|
||||
err := os.WriteFile(path, content, 0600)
|
||||
require.NoError(t, err)
|
||||
if tt.envVars != nil {
|
||||
tt.envVars["KEY_FILE"] = path
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Set up environment
|
||||
for k, v := range tt.envVars {
|
||||
t.Setenv(k, v)
|
||||
}
|
||||
|
||||
keys, err := tt.opts.loadPublicKeys()
|
||||
if tt.wantErr {
|
||||
assert.Error(t, err)
|
||||
if tt.errContains != "" {
|
||||
assert.Contains(t, err.Error(), tt.errContains)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
require.NoError(t, err)
|
||||
assert.Len(t, keys, 1)
|
||||
assert.Equal(t, signer.PublicKey().Type(), keys[0].Type())
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetNetwork(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
opts cmdOptions
|
||||
envVars map[string]string
|
||||
expected string
|
||||
}{
|
||||
{
|
||||
name: "NETWORK env var",
|
||||
envVars: map[string]string{
|
||||
"NETWORK": "tcp4",
|
||||
},
|
||||
expected: "tcp4",
|
||||
},
|
||||
{
|
||||
name: "only port",
|
||||
opts: cmdOptions{listen: "8080"},
|
||||
expected: "tcp",
|
||||
},
|
||||
{
|
||||
name: "ipv4 address",
|
||||
opts: cmdOptions{listen: "1.2.3.4:8080"},
|
||||
expected: "tcp",
|
||||
},
|
||||
{
|
||||
name: "ipv6 address",
|
||||
opts: cmdOptions{listen: "[2001:db8::1]:8080"},
|
||||
expected: "tcp",
|
||||
},
|
||||
{
|
||||
name: "unix network",
|
||||
opts: cmdOptions{listen: "/tmp/beszel.sock"},
|
||||
expected: "unix",
|
||||
},
|
||||
{
|
||||
name: "env var network",
|
||||
opts: cmdOptions{listen: ":8080"},
|
||||
envVars: map[string]string{"NETWORK": "tcp4"},
|
||||
expected: "tcp4",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
// Setup environment
|
||||
for k, v := range tt.envVars {
|
||||
t.Setenv(k, v)
|
||||
}
|
||||
network := agent.GetNetwork(tt.opts.listen)
|
||||
assert.Equal(t, tt.expected, network)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseFlags(t *testing.T) {
|
||||
// Save original command line arguments and restore after test
|
||||
oldArgs := os.Args
|
||||
defer func() {
|
||||
os.Args = oldArgs
|
||||
pflag.CommandLine = pflag.NewFlagSet(os.Args[0], pflag.ExitOnError)
|
||||
}()
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
args []string
|
||||
expected cmdOptions
|
||||
}{
|
||||
{
|
||||
name: "no flags",
|
||||
args: []string{"cmd"},
|
||||
expected: cmdOptions{
|
||||
key: "",
|
||||
listen: "",
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "key flag only",
|
||||
args: []string{"cmd", "-key", "testkey"},
|
||||
expected: cmdOptions{
|
||||
key: "testkey",
|
||||
listen: "",
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "key flag double dash",
|
||||
args: []string{"cmd", "--key", "testkey"},
|
||||
expected: cmdOptions{
|
||||
key: "testkey",
|
||||
listen: "",
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "key flag short",
|
||||
args: []string{"cmd", "-k", "testkey"},
|
||||
expected: cmdOptions{
|
||||
key: "testkey",
|
||||
listen: "",
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "addr flag only",
|
||||
args: []string{"cmd", "-listen", ":8080"},
|
||||
expected: cmdOptions{
|
||||
key: "",
|
||||
listen: ":8080",
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "addr flag double dash",
|
||||
args: []string{"cmd", "--listen", ":8080"},
|
||||
expected: cmdOptions{
|
||||
key: "",
|
||||
listen: ":8080",
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "addr flag short",
|
||||
args: []string{"cmd", "-l", ":8080"},
|
||||
expected: cmdOptions{
|
||||
key: "",
|
||||
listen: ":8080",
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "both flags",
|
||||
args: []string{"cmd", "-key", "testkey", "-listen", ":8080"},
|
||||
expected: cmdOptions{
|
||||
key: "testkey",
|
||||
listen: ":8080",
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
// Reset flags for each test
|
||||
pflag.CommandLine = pflag.NewFlagSet(tt.args[0], pflag.ExitOnError)
|
||||
os.Args = tt.args
|
||||
|
||||
var opts cmdOptions
|
||||
opts.parse()
|
||||
pflag.Parse()
|
||||
|
||||
assert.Equal(t, tt.expected, opts)
|
||||
})
|
||||
}
|
||||
}
|
Reference in New Issue
Block a user