163 lines
3.4 KiB
Go
163 lines
3.4 KiB
Go
package sshd
|
|
|
|
import (
|
|
"errors"
|
|
"flag"
|
|
"fmt"
|
|
"sort"
|
|
"strings"
|
|
|
|
"github.com/armon/go-radix"
|
|
)
|
|
|
|
// 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
|
|
}
|