nebula/sshd/command.go

168 lines
3.6 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 {
// SetOutput() here in case fl.Parse dumps usage.
fl.SetOutput(w.GetWriter())
err := fl.Parse(args)
if err != nil {
// fl.Parse has dumped error information to the user via the w writer.
return err
}
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
}