Public Release

sshd/command.go

@@ -0,0 +1,161 @@
+package sshd
+
+import (
+	"errors"
+	"flag"
+	"fmt"
+	"github.com/armon/go-radix"
+	"sort"
+	"strings"
+)
+
+// CommandFlags is a function called before help or command execution to parse command line flags
+// It should return a flag.FlagSet instance and a pointer to the struct that will contain parsed flags
+type CommandFlags func() (*flag.FlagSet, interface{})
+
+// CommandCallback is the function called when your command should execute.
+// fs will be a a pointer to the struct provided by Command.Flags callback, if there was one. -h and -help are reserved
+// and handled automatically for you.
+// a will be any unconsumed arguments, if no Command.Flags was available this will be all the flags passed in.
+// w is the writer to use when sending messages back to the client.
+// If an error is returned by the callback it is logged locally, the callback should handle messaging errors to the user
+// where appropriate
+type CommandCallback func(fs interface{}, a []string, w StringWriter) error
+
+type Command struct {
+	Name             string
+	ShortDescription string
+	Help             string
+	Flags            CommandFlags
+	Callback         CommandCallback
+}
+
+func execCommand(c *Command, args []string, w StringWriter) error {
+	var (
+		fl *flag.FlagSet
+		fs interface{}
+	)
+
+	if c.Flags != nil {
+		fl, fs = c.Flags()
+		if fl != nil {
+			//TODO: handle the error
+			fl.Parse(args)
+			args = fl.Args()
+		}
+	}
+
+	return c.Callback(fs, args, w)
+}
+
+func dumpCommands(c *radix.Tree, w StringWriter) {
+	err := w.WriteLine("Available commands:")
+	if err != nil {
+		//TODO: log
+		return
+	}
+
+	cmds := make([]string, 0)
+	for _, l := range allCommands(c) {
+		cmds = append(cmds, fmt.Sprintf("%s - %s", l.Name, l.ShortDescription))
+	}
+
+	sort.Strings(cmds)
+	err = w.Write(strings.Join(cmds, "\n") + "\n\n")
+	if err != nil {
+		//TODO: log
+	}
+}
+
+func lookupCommand(c *radix.Tree, sCmd string) (*Command, error) {
+	cmd, ok := c.Get(sCmd)
+	if !ok {
+		return nil, nil
+	}
+
+	command, ok := cmd.(*Command)
+	if !ok {
+		return nil, errors.New("failed to cast command")
+	}
+
+	return command, nil
+}
+
+func matchCommand(c *radix.Tree, cmd string) []string {
+	cmds := make([]string, 0)
+	c.WalkPrefix(cmd, func(found string, v interface{}) bool {
+		cmds = append(cmds, found)
+		return false
+	})
+	sort.Strings(cmds)
+	return cmds
+}
+
+func allCommands(c *radix.Tree) []*Command {
+	cmds := make([]*Command, 0)
+	c.WalkPrefix("", func(found string, v interface{}) bool {
+		cmd, ok := v.(*Command)
+		if ok {
+			cmds = append(cmds, cmd)
+		}
+		return false
+	})
+	return cmds
+}
+
+func helpCallback(commands *radix.Tree, a []string, w StringWriter) (err error) {
+	// Just typed help
+	if len(a) == 0 {
+		dumpCommands(commands, w)
+		return nil
+	}
+
+	// We are printing a specific commands help text
+	cmd, err := lookupCommand(commands, a[0])
+	if err != nil {
+		//TODO: handle error
+		//TODO: message the user
+		return
+	}
+
+	if cmd != nil {
+		err = w.WriteLine(fmt.Sprintf("%s - %s", cmd.Name, cmd.ShortDescription))
+		if err != nil {
+			return err
+		}
+
+		if cmd.Help != "" {
+			err = w.WriteLine(fmt.Sprintf("  %s", cmd.Help))
+			if err != nil {
+				return err
+			}
+		}
+
+		if cmd.Flags != nil {
+			fs, _ := cmd.Flags()
+			if fs != nil {
+				fs.SetOutput(w.GetWriter())
+				fs.PrintDefaults()
+			}
+		}
+
+		return nil
+	}
+
+	err = w.WriteLine("Command not available " + a[0])
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func checkHelpArgs(args []string) bool {
+	for _, a := range args {
+		if a == "-h" || a == "-help" {
+			return true
+		}
+	}
+
+	return false
+}