diff --git a/beszel/cmd/agent/agent.go b/beszel/cmd/agent/agent.go index 86b00e2..c2bc049 100644 --- a/beszel/cmd/agent/agent.go +++ b/beszel/cmd/agent/agent.go @@ -4,12 +4,12 @@ import ( "beszel" "beszel/internal/agent" "beszel/internal/agent/health" - "flag" "fmt" "log" "os" "strings" + "github.com/spf13/pflag" "golang.org/x/crypto/ssh" ) @@ -17,43 +17,24 @@ import ( 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 { - 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] } + // Subcommands that don't require any pflag parsing 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 { @@ -63,7 +44,57 @@ func (opts *cmdOptions) parse() bool { return true } - flag.Parse() + // 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 } diff --git a/beszel/cmd/agent/agent_test.go b/beszel/cmd/agent/agent_test.go index 1983ab4..9f41563 100644 --- a/beszel/cmd/agent/agent_test.go +++ b/beszel/cmd/agent/agent_test.go @@ -3,11 +3,11 @@ package main import ( "beszel/internal/agent" "crypto/ed25519" - "flag" "os" "path/filepath" "testing" + "github.com/spf13/pflag" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "golang.org/x/crypto/ssh" @@ -245,7 +245,7 @@ func TestParseFlags(t *testing.T) { oldArgs := os.Args defer func() { os.Args = oldArgs - flag.CommandLine = flag.NewFlagSet(os.Args[0], flag.ExitOnError) + pflag.CommandLine = pflag.NewFlagSet(os.Args[0], pflag.ExitOnError) }() tests := []struct { @@ -269,6 +269,22 @@ func TestParseFlags(t *testing.T) { 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"}, @@ -277,6 +293,22 @@ func TestParseFlags(t *testing.T) { 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"}, @@ -290,12 +322,12 @@ func TestParseFlags(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { // Reset flags for each test - flag.CommandLine = flag.NewFlagSet(tt.args[0], flag.ExitOnError) + pflag.CommandLine = pflag.NewFlagSet(tt.args[0], pflag.ExitOnError) os.Args = tt.args var opts cmdOptions opts.parse() - flag.Parse() + pflag.Parse() assert.Equal(t, tt.expected, opts) }) diff --git a/beszel/cmd/hub/hub.go b/beszel/cmd/hub/hub.go index f0bc96c..048b8a9 100644 --- a/beszel/cmd/hub/hub.go +++ b/beszel/cmd/hub/hub.go @@ -45,11 +45,13 @@ func getBaseApp() *pocketbase.PocketBase { baseApp.RootCmd.Use = beszel.AppName baseApp.RootCmd.Short = "" // add update command - baseApp.RootCmd.AddCommand(&cobra.Command{ + updateCmd := &cobra.Command{ Use: "update", Short: "Update " + beszel.AppName + " to the latest version", Run: hub.Update, - }) + } + updateCmd.Flags().Bool("china-mirrors", false, "Use mirror (gh.beszel.dev) instead of GitHub") + baseApp.RootCmd.AddCommand(updateCmd) // add health command baseApp.RootCmd.AddCommand(newHealthCmd()) diff --git a/beszel/internal/agent/update.go b/beszel/internal/agent/update.go index 0c39926..e3e9b15 100644 --- a/beszel/internal/agent/update.go +++ b/beszel/internal/agent/update.go @@ -60,7 +60,7 @@ func detectRestarter() restarter { // Update checks GitHub for a newer release of beszel-agent, applies it, // fixes SELinux context if needed, and restarts the service. -func Update() error { +func Update(useMirror bool) error { exePath, _ := os.Executable() dataDir, err := getDataDir() @@ -70,6 +70,7 @@ func Update() error { updated, err := ghupdate.Update(ghupdate.Config{ ArchiveExecutable: "beszel-agent", DataDir: dataDir, + UseMirror: useMirror, }) if err != nil { log.Fatal(err) @@ -99,6 +100,8 @@ func Update() error { if err := r.Restart(); err != nil { ghupdate.ColorPrintf(ghupdate.ColorYellow, "Warning: failed to restart service: %v", err) ghupdate.ColorPrint(ghupdate.ColorYellow, "Please restart the service manually.") + } else { + ghupdate.ColorPrint(ghupdate.ColorGreen, "Service restarted successfully") } } else { ghupdate.ColorPrint(ghupdate.ColorYellow, "No supported init system detected; please restart manually if needed.") diff --git a/beszel/internal/ghupdate/ghupdate.go b/beszel/internal/ghupdate/ghupdate.go index 711c699..575fe9a 100644 --- a/beszel/internal/ghupdate/ghupdate.go +++ b/beszel/internal/ghupdate/ghupdate.go @@ -23,7 +23,7 @@ import ( const ( colorReset = "\033[0m" ColorYellow = "\033[33m" - colorGreen = "\033[32m" + ColorGreen = "\033[32m" colorCyan = "\033[36m" colorGray = "\033[90m" ) @@ -64,6 +64,10 @@ type Config struct { // The data directory to use when fetching and downloading the latest release. DataDir string + + // UseMirror specifies whether to use the beszel.dev mirror instead of GitHub API. + // When false (default), always uses api.github.com. When true, uses gh.beszel.dev. + UseMirror bool } type updater struct { @@ -106,21 +110,19 @@ func (p *updater) update() (updated bool, err error) { var latest *release var useMirror bool + // Determine the API endpoint based on UseMirror flag + apiURL := fmt.Sprintf("https://api.github.com/repos/%s/%s/releases/latest", p.config.Owner, p.config.Repo) + if p.config.UseMirror { + useMirror = true + apiURL = fmt.Sprintf("https://gh.beszel.dev/repos/%s/%s/releases/latest?api=true", p.config.Owner, p.config.Repo) + ColorPrint(ColorYellow, "Using mirror for update.") + } + latest, err = fetchLatestRelease( p.config.Context, p.config.HttpClient, - fmt.Sprintf("https://api.github.com/repos/%s/%s/releases/latest", p.config.Owner, p.config.Repo), + apiURL, ) - // if the first fetch fails, try the beszel.dev API (fallback for China) - if err != nil { - ColorPrint(ColorYellow, "Failed to fetch release. Trying beszel.dev mirror...") - useMirror = true - latest, err = fetchLatestRelease( - p.config.Context, - p.config.HttpClient, - fmt.Sprintf("https://gh.beszel.dev/repos/%s/%s/releases/latest?api=true", p.config.Owner, p.config.Repo), - ) - } if err != nil { return false, err } @@ -129,7 +131,7 @@ func (p *updater) update() (updated bool, err error) { newVersion := semver.MustParse(strings.TrimPrefix(latest.Tag, "v")) if newVersion.LTE(currentVersion) { - ColorPrintf(colorGreen, "You already have the latest version %s.", p.currentVersion) + ColorPrintf(ColorGreen, "You already have the latest version %s.", p.currentVersion) return false, nil } @@ -209,14 +211,11 @@ func (p *updater) update() (updated bool, err error) { } ColorPrint(colorGray, "---") - ColorPrint(colorGreen, "Update completed successfully! You can start the executable as usual.") + ColorPrint(ColorGreen, "Update completed successfully!") // print the release notes if latest.Body != "" { fmt.Print("\n") - ColorPrintf(colorCyan, "Here is a list with some of the %s changes:", latest.Tag) - // remove the update command note to avoid "stuttering" - // (@todo consider moving to a config option) releaseNotes := strings.TrimSpace(strings.Replace(latest.Body, "> _To update the prebuilt executable you can run `./"+p.config.ArchiveExecutable+" update`._", "", 1)) ColorPrint(colorCyan, releaseNotes) fmt.Print("\n") diff --git a/beszel/internal/hub/update.go b/beszel/internal/hub/update.go index 246f579..b6aebee 100644 --- a/beszel/internal/hub/update.go +++ b/beszel/internal/hub/update.go @@ -11,7 +11,7 @@ import ( ) // Update updates beszel to the latest version -func Update(_ *cobra.Command, _ []string) { +func Update(cmd *cobra.Command, _ []string) { dataDir := os.TempDir() // set dataDir to ./beszel_data if it exists @@ -19,9 +19,13 @@ func Update(_ *cobra.Command, _ []string) { dataDir = "./beszel_data" } + // Check if china-mirrors flag is set + useMirror, _ := cmd.Flags().GetBool("china-mirrors") + updated, err := ghupdate.Update(ghupdate.Config{ ArchiveExecutable: "beszel", DataDir: dataDir, + UseMirror: useMirror, }) if err != nil { log.Fatal(err) @@ -49,13 +53,13 @@ func restartService() { // Check if beszel service exists and is active cmd := exec.Command("systemctl", "is-active", "beszel.service") if err := cmd.Run(); err == nil { - fmt.Println("Restarting beszel service...") + ghupdate.ColorPrint(ghupdate.ColorYellow, "Restarting beszel service...") restartCmd := exec.Command("systemctl", "restart", "beszel.service") if err := restartCmd.Run(); err != nil { - fmt.Printf("Warning: Failed to restart service: %v\n", err) - fmt.Println("Please restart the service manually: sudo systemctl restart beszel") + ghupdate.ColorPrintf(ghupdate.ColorYellow, "Warning: Failed to restart service: %v\n", err) + ghupdate.ColorPrint(ghupdate.ColorYellow, "Please restart the service manually: sudo systemctl restart beszel") } else { - fmt.Println("Service restarted successfully") + ghupdate.ColorPrint(ghupdate.ColorGreen, "Service restarted successfully") } return } @@ -65,17 +69,17 @@ func restartService() { if _, err := exec.LookPath("rc-service"); err == nil { cmd := exec.Command("rc-service", "beszel", "status") if err := cmd.Run(); err == nil { - fmt.Println("Restarting beszel service...") + ghupdate.ColorPrint(ghupdate.ColorYellow, "Restarting beszel service...") restartCmd := exec.Command("rc-service", "beszel", "restart") if err := restartCmd.Run(); err != nil { - fmt.Printf("Warning: Failed to restart service: %v\n", err) - fmt.Println("Please restart the service manually: sudo rc-service beszel restart") + ghupdate.ColorPrintf(ghupdate.ColorYellow, "Warning: Failed to restart service: %v\n", err) + ghupdate.ColorPrint(ghupdate.ColorYellow, "Please restart the service manually: sudo rc-service beszel restart") } else { - fmt.Println("Service restarted successfully") + ghupdate.ColorPrint(ghupdate.ColorGreen, "Service restarted successfully") } return } } - fmt.Println("Note: Service restart not attempted. If running as a service, restart manually.") + ghupdate.ColorPrint(ghupdate.ColorYellow, "Service restart not attempted. If running as a service, restart manually.") }