New Command pkg

internal/aes/ecb.go

@@ -3,38 +3,37 @@ package aes
 import (
 	"crypto/aes"
 	"errors"
-	"flag"
 	"fmt"
-	"os"
 
+	"bandr.me/p/pocryp/internal/cli/cmd"
 	"bandr.me/p/pocryp/internal/util"
 	"bandr.me/p/pocryp/internal/util/stdfile"
 )
 
-func EcbCmd(args ...string) error {
-	fset := flag.NewFlagSet("aes-ecb", flag.ContinueOnError)
-	fset.Usage = func() {
-		fmt.Fprint(os.Stderr, `Usage: pocryp aes-ecb [-bin] [-e/-d] -key|-key-file [-in INPUT] [-out OUTPUT]
+var EcbCmd = &cmd.Command{
+	Name:  "aes-ecb",
+	Run:   runEcb,
+	Brief: "Encrypt/Decrypt using AES-ECB",
+
+	Usage: `Usage: pocryp aes-ecb [-bin] [-e/-d] -key|-key-file [-in INPUT] [-out OUTPUT]
 
 Encrypt/Decrypt INPUT to OUTPUT using AES-ECB.
 
 If -in is not specified, stdin will be read.
 If -out is not specified, the output will be printed to stdout.
+`,
+}
 
-Options:
-`)
-		fset.PrintDefaults()
-	}
-
-	fEncrypt := fset.Bool("e", false, "Encrypt the input to the output. Default if omitted.")
-	fDecrypt := fset.Bool("d", false, "Decrypt the input to the output.")
-	fOutput := fset.String("out", "", "Write the result to the file at path OUTPUT.")
-	fInput := fset.String("in", "", "Read data from the file at path INPUT.")
-	fKey := fset.String("key", "", "Key as hex.")
-	fKeyFile := fset.String("key-file", "", "File which contains the key as binary/text.")
-	fBin := fset.Bool("bin", false, "Print output in binary form not hex.")
+func runEcb(cmd *cmd.Command) error {
+	fEncrypt := cmd.Flags.Bool("e", false, "Encrypt the input to the output. Default if omitted.")
+	fDecrypt := cmd.Flags.Bool("d", false, "Decrypt the input to the output.")
+	fOutput := cmd.Flags.String("out", "", "Write the result to the file at path OUTPUT.")
+	fInput := cmd.Flags.String("in", "", "Read data from the file at path INPUT.")
+	fKey := cmd.Flags.String("key", "", "Key as hex.")
+	fKeyFile := cmd.Flags.String("key-file", "", "File which contains the key as binary/text.")
+	fBin := cmd.Flags.Bool("bin", false, "Print output in binary form not hex.")
 
-	if err := fset.Parse(args); err != nil {
+	if err := cmd.Parse(); err != nil {
 		return err
 	}
 

internal/cli/app.go

@@ -0,0 +1,145 @@
+package cli
+
+import (
+	"flag"
+	"fmt"
+	"os"
+	"runtime/debug"
+	"strings"
+
+	"bandr.me/p/pocryp/internal/cli/cmd"
+)
+
+type App struct {
+	categories []category
+
+	printVersion bool
+}
+
+type category struct {
+	name     string
+	commands []*cmd.Command
+}
+
+func printAppVersion() {
+	bi, _ := debug.ReadBuildInfo()
+	valOf := func(k string) string {
+		for _, v := range bi.Settings {
+			if v.Key == k {
+				return v.Value
+			}
+		}
+		return ""
+	}
+	fmt.Println(
+		bi.Main.Version,
+		bi.GoVersion,
+		valOf("GOOS"),
+		valOf("GOARCH"),
+		valOf("vcs.revision"),
+		valOf("vcs.time"),
+	)
+}
+
+func (a *App) Run(args ...string) error {
+	fset := flag.NewFlagSet("pocryp", flag.ExitOnError)
+	fset.Usage = a.Usage
+	fset.BoolVar(&a.printVersion, "version", false, "")
+	if err := fset.Parse(args); err != nil {
+		return err
+	}
+
+	args = fset.Args()
+
+	if a.printVersion {
+		printAppVersion()
+		return nil
+	}
+
+	if len(args) == 0 {
+		a.Usage()
+		return nil
+	}
+
+	name := args[0]
+	args = args[1:]
+
+	for _, category := range a.categories {
+		for _, cmd := range category.commands {
+			if cmd.Name == name {
+				cmd.Args = args
+				return cmd.Run(cmd)
+			}
+		}
+	}
+
+	return fmt.Errorf("unknown command '%s'", name)
+}
+
+func (a *App) Add(categoryName string, cmds ...*cmd.Command) {
+	i := -1
+	for ii, v := range a.categories {
+		if v.name == categoryName {
+			i = ii
+		}
+	}
+	if i == -1 {
+		a.categories = append(a.categories, category{name: categoryName})
+		i = len(a.categories) - 1
+	}
+	cat := &a.categories[i]
+	for _, cmd := range cmds {
+		if err := cat.hasCmd(cmd.Name); err != nil {
+			panic(err.Error())
+		}
+		cmd.Init()
+	}
+	cat.commands = append(cat.commands, cmds...)
+}
+
+func (c category) hasCmd(name string) error {
+	for _, cmd := range c.commands {
+		if cmd.Name == name {
+			return fmt.Errorf("category '%s' already has a command name '%s'", c.name, name)
+		}
+	}
+	return nil
+}
+
+func (a App) maxCommandName(category string) int {
+	max := 0
+	for _, v := range a.categories {
+		if v.name != category {
+			continue
+		}
+		for _, cmd := range v.commands {
+			if len(cmd.Name) > max {
+				max = len(cmd.Name)
+			}
+		}
+	}
+	return max
+}
+
+func (a App) Usage() {
+	w := os.Stderr
+	fmt.Fprint(w, `Usage: pocryp command [ARGS]
+
+Flags:
+  -h, --help  Print this message
+  --version   Print version information
+
+Commands(by category):
+
+`)
+	for _, v := range a.categories {
+		fmt.Fprintf(w, "%s:\n", v.name)
+		mlen := a.maxCommandName(v.name)
+		for _, cmd := range v.commands {
+			padding := strings.Repeat(" ", mlen-len(cmd.Name))
+			fmt.Fprintf(w, "  %s%s  %s\n", cmd.Name, padding, cmd.Usage)
+		}
+		fmt.Fprint(w, "\n")
+	}
+	fmt.Fprint(w, "Run 'pocryp command -h' for more information about a command.\n")
+}

internal/cli/app_test.go

@@ -0,0 +1,48 @@
+package cli
+
+import (
+	"testing"
+
+	"bandr.me/p/pocryp/internal/cli/cmd"
+)
+
+func TestApp(t *testing.T) {
+	var app App
+
+	if err := app.Run(); err != nil {
+		t.Fatal(err)
+	}
+
+	if err := app.Run("nothing"); err == nil {
+		t.Fatal("expected an error")
+	}
+
+	app.Add("foo", &cmd.Command{
+		Name:  "bar",
+		Brief: "bar",
+		Usage: "bar",
+		Run:   func(*cmd.Command) error { return nil },
+	})
+	app.Add("foo", &cmd.Command{
+		Name:  "baz",
+		Brief: "baz",
+		Usage: "baz",
+		Run:   func(*cmd.Command) error { return nil },
+	})
+	app.Add("fizz", &cmd.Command{
+		Name:  "buzz",
+		Brief: "barr",
+		Usage: "barr",
+		Run:   func(*cmd.Command) error { return nil },
+	})
+
+	if err := app.Run("nothing"); err == nil {
+		t.Fatal("expected an error")
+	}
+
+	if err := app.Run("bar"); err != nil {
+		t.Fatal(err)
+	}
+
+	app.Usage()
+}

internal/cli/cmd/command.go

@@ -0,0 +1,43 @@
+package cmd
+
+import (
+	"flag"
+	"fmt"
+	"os"
+)
+
+type Command struct {
+	Run   func(*Command) error
+	Name  string
+	Brief string
+	Usage string
+
+	Args  []string
+	Flags *flag.FlagSet
+}
+
+func (c *Command) Parse() error {
+	return c.Flags.Parse(c.Args)
+}
+
+func (c *Command) Init() {
+	if c.Run == nil {
+		panic("cmd: missing Run")
+	}
+	if c.Name == "" {
+		panic("cmd: missing Name")
+	}
+	if c.Usage == "" {
+		panic("cmd: missing Usage")
+	}
+	if c.Brief == "" {
+		panic("cmd: missing Brief")
+	}
+	c.Flags = flag.NewFlagSet(c.Name, flag.ContinueOnError)
+	c.Flags.Usage = func() {
+		fmt.Fprint(os.Stderr, c.Usage)
+		fmt.Fprintln(os.Stderr, "Options:")
+		c.Flags.PrintDefaults()
+		fmt.Fprint(os.Stderr, "\n")
+	}
+}

internal/encoding/rsa/priv2pub.go

@@ -3,33 +3,31 @@ package rsa
 import (
 	"crypto/x509"
 	"encoding/pem"
-	"flag"
-	"fmt"
-	"os"
 
+	"bandr.me/p/pocryp/internal/cli/cmd"
 	"bandr.me/p/pocryp/internal/encoding/rsa/util"
 	"bandr.me/p/pocryp/internal/util/stdfile"
 )
 
-func Priv2PubCmd(args ...string) error {
-	fset := flag.NewFlagSet("rsa-priv2pub", flag.ContinueOnError)
-	fset.Usage = func() {
-		fmt.Fprint(os.Stderr, `Usage: pocryp rsa-priv2pub [-in INPUT] [-out OUTPUT]
+var Priv2PubCmd = &cmd.Command{
+	Name:  "rsa-priv2pub",
+	Run:   runPriv2Pub,
+	Brief: "Extract RSA public key from private key",
+
+	Usage: `Usage: pocryp rsa-priv2pub [-in INPUT] [-out OUTPUT]
 
 Extract RSA public key from private key, specified as PEM.
 
 If -in is not specified, stdin will be read.
 If -out is not specified, the output will be printed to stdout.
+`,
+}
 
-Options:
-`)
-		fset.PrintDefaults()
-	}
-
-	fOutput := fset.String("out", "", "Write the result to the file at path OUTPUT.")
-	fInput := fset.String("in", "", "Read data from the file at path INPUT.")
+func runPriv2Pub(cmd *cmd.Command) error {
+	fOutput := cmd.Flags.String("out", "", "Write the result to the file at path OUTPUT.")
+	fInput := cmd.Flags.String("in", "", "Read data from the file at path INPUT.")
 
-	if err := fset.Parse(args); err != nil {
+	if err := cmd.Parse(); err != nil {
 		return err
 	}