package bot import ( "context" "errors" "log" "strings" "sync" "time" "github.com/bwmarrin/discordgo" opentracing "github.com/opentracing/opentracing-go" otlog "github.com/opentracing/opentracing-go/log" "golang.org/x/time/rate" ) var ( ErrRateLimitExceeded = errors.New("bot: per-command rate limit exceeded") ) type command struct { aliases []string verb string helptext string } func (c *command) Verb() string { return c.verb } func (c *command) Helptext() string { return c.helptext } // Handler is the type that bot command functions need to implement. Errors // should be returned. type Handler func(context.Context, *discordgo.Session, *discordgo.Message, []string) error // CommandHandler is a generic interface for types that implement a bot // command. It is akin to http.Handler, but more comprehensive. type CommandHandler interface { Verb() string Helptext() string Handler(context.Context, *discordgo.Session, *discordgo.Message, []string) error Permissions(context.Context, *discordgo.Session, *discordgo.Message, []string) error } type basicCommand struct { *command handler Handler permissions Handler limiter *rate.Limiter } func (bc *basicCommand) Handler(ctx context.Context, s *discordgo.Session, m *discordgo.Message, parv []string) error { return bc.handler(ctx, s, m, parv) } func (bc *basicCommand) Permissions(ctx context.Context, s *discordgo.Session, m *discordgo.Message, parv []string) error { if !bc.limiter.Allow() { return ErrRateLimitExceeded } return bc.permissions(ctx, s, m, parv) } // The "default" command set, useful for simple bot projects. var ( DefaultCommandSet = NewCommandSet() ) // Command handling errors. var ( ErrAlreadyExists = errors.New("bot: command already exists") ErrNoSuchCommand = errors.New("bot: no such command exists") ErrNoPermissions = errors.New("bot: you do not have permissions for this command") ErrParvCountMismatch = errors.New("bot: parameter count mismatch") ) // The default command prefix. Command `foo` becomes `.foo` in chat, etc. const ( DefaultPrefix = "." ) // NewCommand creates an anonymous command and adds it to the default CommandSet. func NewCommand(verb, helptext string, handler, permissions Handler) error { return DefaultCommandSet.Add(NewBasicCommand(verb, helptext, handler, permissions)) } // NewBasicCommand creates a CommandHandler instance using the implementation // functions supplied as arguments. func NewBasicCommand(verb, helptext string, permissions, handler Handler) CommandHandler { return &basicCommand{ command: &command{ verb: verb, helptext: helptext, }, handler: handler, permissions: permissions, limiter: rate.NewLimiter(rate.Every(5*time.Second), 1), } } // CommandSet is a group of bot commands similar to an http.ServeMux. type CommandSet struct { sync.Mutex cmds map[string]CommandHandler Prefix string } // NewCommandSet creates a new command set with the `help` command pre-loaded. func NewCommandSet() *CommandSet { cs := &CommandSet{ cmds: map[string]CommandHandler{}, Prefix: DefaultPrefix, } cs.AddCmd("help", "Shows help for the bot", NoPermissions, cs.help) return cs } // NoPermissions is a simple middelware function that allows all command invocations // to pass the permissions check. func NoPermissions(ctx context.Context, s *discordgo.Session, m *discordgo.Message, parv []string) error { return nil } // AddCmd is syntactic sugar for cs.Add(NewBasicCommand(args...)) func (cs *CommandSet) AddCmd(verb, helptext string, permissions, handler Handler) error { return cs.Add(NewBasicCommand(verb, helptext, permissions, handler)) } // Add adds a single command handler to the CommandSet. This can be done at runtime // but it is suggested to only add commands on application boot. func (cs *CommandSet) Add(h CommandHandler) error { cs.Lock() defer cs.Unlock() v := strings.ToLower(h.Verb()) if _, ok := cs.cmds[v]; ok { return ErrAlreadyExists } cs.cmds[v] = h return nil } // Run executes a single command handler if applicable func (cs *CommandSet) Run(ctx context.Context, s *discordgo.Session, msg *discordgo.Message) error { sp, ctx := opentracing.StartSpanFromContext(ctx, "bot.commandset.run") defer sp.Finish() cs.Lock() defer cs.Unlock() ctx, cancel := context.WithCancel(ctx) defer cancel() if strings.HasPrefix(msg.Content, cs.Prefix) { params := strings.Fields(msg.Content) verb := strings.ToLower(params[0][1:]) sp.LogFields( otlog.String("message-id", msg.ID), otlog.String("author", msg.Author.ID), otlog.String("channel-id", msg.ChannelID), otlog.String("verb", verb), otlog.Int("parc", len(params)), ) cmd, ok := cs.cmds[verb] if !ok { sp.LogFields(otlog.Error(ErrNoSuchCommand)) return ErrNoSuchCommand } err := cmd.Permissions(ctx, s, msg, params) if err != nil { sp.LogFields(otlog.Error(err)) log.Printf("Permissions error: %s: %v", msg.Author.Username, err) s.ChannelMessageSend(msg.ChannelID, "You don't have permissions for that, sorry.") return ErrNoPermissions } err = cmd.Handler(ctx, s, msg, params) if err != nil { log.Printf("command handler error: %v", err) s.ChannelMessageSend(msg.ChannelID, "error when running that command: "+err.Error()) return err } } return nil }