vyvanse/bot/bot.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
}