Basic CLI

Signed-off-by: Chris Koch <chrisko@google.com>

cli/cli.go

@@ -0,0 +1,93 @@
+// Copyright 2024 the u-root Authors. All rights reserved
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// Package cli provides a bare bones CLI with commands.
+//
+// cli supports standard Go flags.
+package cli
+
+import (
+	"flag"
+	"fmt"
+	"io"
+	"os"
+	"slices"
+	"strings"
+	"text/tabwriter"
+)
+
+// Command is a CLI command.
+type Command struct {
+	// Name or Aliases are used to match argv[1] and select this command.
+	Name    string
+	Aliases []string
+
+	// Short is a one-line description for the command.
+	Short string
+
+	// Run is called if argv[1] matches the command. Flags are parsed before Run is called.
+	Run func(args []string)
+
+	flags *flag.FlagSet
+}
+
+func (c *Command) Flags() *flag.FlagSet {
+	if c.flags == nil {
+		c.flags = flag.NewFlagSet(c.Name, flag.ContinueOnError)
+	}
+	return c.flags
+}
+
+// An App is composed of many commands.
+type App []Command
+
+// Help returns the app's help string.
+func (a App) Help() string {
+	var s strings.Builder
+	w := tabwriter.NewWriter(&s, 1, 2, 4, ' ', 0)
+	fmt.Fprintf(w, "Commands:\n\n")
+	for _, cmd := range a {
+		fmt.Fprintf(w, "\t%s\t%s\n", cmd.Name, cmd.Short)
+	}
+	w.Flush()
+	return s.String()
+}
+
+func (a App) commandFor(args []string) *Command {
+	if len(args) == 0 {
+		return nil
+	}
+	for _, cmd := range a {
+		if args[0] == cmd.Name || slices.Contains(cmd.Aliases, args[0]) {
+			return &cmd
+		}
+	}
+	return nil
+}
+
+func (a App) run(errW io.Writer, args []string) int {
+	if len(args) == 0 {
+		fmt.Fprintf(errW, "No program name provided\n")
+		return 1
+	}
+	cmd := a.commandFor(args[1:])
+	if cmd == nil {
+		fmt.Fprintf(errW, "%s", a.Help())
+		return 1
+	}
+
+	cmd.Flags().SetOutput(errW)
+	if err := cmd.Flags().Parse(args[2:]); err != nil {
+		cmd.Flags().Output()
+		return 1
+	}
+
+	cmd.Run(cmd.Flags().Args())
+	return 0
+}
+
+// Run runs the app. Expects program name as the first arg, and an optional command name next.
+func (a App) Run(args []string) {
+	os.Exit(a.run(os.Stderr, args))
+}

cli/cli_test.go

@@ -0,0 +1,145 @@
+// Copyright 2024 the u-root Authors. All rights reserved
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package cli
+
+import (
+	"reflect"
+	"strings"
+	"testing"
+)
+
+func TestCLI(t *testing.T) {
+	type flags struct {
+		output string
+		input  string
+	}
+	var f flags
+	var cmd string
+	var cmdArgs []string
+
+	makeCmd := Command{
+		Name:  "make",
+		Short: "create uimage",
+		Run: func(args []string) {
+			cmdArgs = args
+			cmd = "make"
+		},
+	}
+	makeCmd.Flags().StringVar(&f.output, "o", "", "Output")
+	makeCmd.Flags().StringVar(&f.input, "i", "", "Input")
+
+	listCmd := Command{
+		Name:    "list",
+		Short:   "list uimage",
+		Aliases: []string{"ls", "l"},
+		Run: func(args []string) {
+			cmdArgs = args
+			cmd = "list"
+		},
+	}
+	app := App{makeCmd, listCmd}
+
+	for _, tt := range []struct {
+		name        string
+		args        []string
+		wantCmd     string
+		wantCmdArgs []string
+		wantFlags   flags
+		wantExit    int
+		wantPrint   string
+	}{
+		{
+			name:        "cmd with flag",
+			args:        []string{"uimage", "make", "-o", "high", "foobar", "bla"},
+			wantCmdArgs: []string{"foobar", "bla"},
+			wantCmd:     "make",
+			wantFlags: flags{
+				output: "high",
+			},
+			wantExit:  0,
+			wantPrint: "",
+		},
+		{
+			name:      "not exist",
+			args:      []string{"uimage", "notmake", "-o", "low"},
+			wantExit:  1,
+			wantPrint: "Commands:\n\n    make    create uimage\n    list    list uimage\n",
+		},
+		{
+			name:      "cmd exist but flag doesn't",
+			args:      []string{"uimage", "list", "-o", "low"},
+			wantExit:  1,
+			wantPrint: "flag provided but not defined: -o\nUsage of list:\n",
+		},
+		{
+			name:        "cmd with no flags",
+			args:        []string{"uimage", "list", "anything"},
+			wantExit:    0,
+			wantCmd:     "list",
+			wantCmdArgs: []string{"anything"},
+		},
+		{
+			name:        "alias",
+			args:        []string{"uimage", "ls", "anything"},
+			wantExit:    0,
+			wantCmd:     "list",
+			wantCmdArgs: []string{"anything"},
+		},
+		{
+			name:      "no program name",
+			wantExit:  1,
+			wantPrint: "No program name provided\n",
+		},
+		{
+			name:      "no command name",
+			args:      []string{"uimage"},
+			wantExit:  1,
+			wantPrint: "Commands:\n\n    make    create uimage\n    list    list uimage\n",
+		},
+		{
+			name:      "cmd help",
+			args:      []string{"uimage", "make", "-h"},
+			wantExit:  1,
+			wantPrint: "Usage of make:\n  -i string\n    \tInput\n  -o string\n    \tOutput\n",
+		},
+		{
+			name:      "app help",
+			args:      []string{"uimage", "-h"},
+			wantExit:  1,
+			wantPrint: "Commands:\n\n    make    create uimage\n    list    list uimage\n",
+		},
+		{
+			name:      "app help 2",
+			args:      []string{"uimage", "help"},
+			wantExit:  1,
+			wantPrint: "Commands:\n\n    make    create uimage\n    list    list uimage\n",
+		},
+	} {
+		t.Run(tt.name, func(t *testing.T) {
+			f = flags{}
+			cmd = ""
+			cmdArgs = nil
+
+			var s strings.Builder
+			exitCode := app.run(&s, tt.args)
+			t.Logf("App:\n%s", s.String())
+			if exitCode != tt.wantExit {
+				t.Errorf("run = %d, want %d", exitCode, tt.wantExit)
+			}
+			if cmd != tt.wantCmd {
+				t.Errorf("run = cmd %s, want cmd %s", cmd, tt.wantCmd)
+			}
+			if !reflect.DeepEqual(cmdArgs, tt.wantCmdArgs) {
+				t.Errorf("run = args %+v, want %+v", cmdArgs, tt.wantCmdArgs)
+			}
+			if !reflect.DeepEqual(f, tt.wantFlags) {
+				t.Errorf("run = flags %+v, want %+v", f, tt.wantFlags)
+			}
+			if got := s.String(); got != tt.wantPrint {
+				t.Errorf("run = %#v, want %#v", got, tt.wantPrint)
+			}
+		})
+	}
+}