rename /src to /internal (sorry i'll fix the prs)

This commit is contained in:
henrygd
2025-09-09 13:29:07 -04:00
parent 86ea23fe39
commit 8a13b05c20
177 changed files with 0 additions and 0 deletions

157
internal/cmd/agent/agent.go Normal file
View 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)
}
}

View 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)
})
}
}