199 lines
5.2 KiB
Go
199 lines
5.2 KiB
Go
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 makes a CommandSet compatible with discordgo event dispatching.
|
|
func (cs *CommandSet) Run(s *discordgo.Session, msg *discordgo.Message) error {
|
|
cs.Lock()
|
|
defer cs.Unlock()
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
|
|
sp, ctx := opentracing.StartSpanFromContext(ctx, "CommandSet.Run")
|
|
defer sp.Finish()
|
|
|
|
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 {
|
|
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
|
|
}
|