1
0
Fork 0
mirror of https://github.com/diamondburned/arikawa.git synced 2024-12-11 16:05:00 +00:00
arikawa/utils/bot/ctx.go
Cléo Rebert 8f548d2607 go fmt + gci
Signed-off-by: Cléo Rebert <cleo.rebert@gmail.com>
2023-09-19 11:37:24 -07:00

541 lines
16 KiB
Go

package bot
import (
"context"
"errors"
"fmt"
"log"
"os"
"os/signal"
"strings"
"sync"
"syscall"
"time"
"github.com/diamondburned/arikawa/v3/api"
"github.com/diamondburned/arikawa/v3/gateway"
"github.com/diamondburned/arikawa/v3/session"
"github.com/diamondburned/arikawa/v3/session/shard"
"github.com/diamondburned/arikawa/v3/state"
"github.com/diamondburned/arikawa/v3/state/store/defaultstore"
"github.com/diamondburned/arikawa/v3/utils/bot/extras/shellwords"
"github.com/diamondburned/arikawa/v3/utils/handler"
)
// Prefixer checks a message if it starts with the desired prefix. By default,
// NewPrefix() is used.
type Prefixer func(*gateway.MessageCreateEvent) (prefix string, ok bool)
// NewPrefix creates a simple prefix checker using strings. As the default
// prefix is "!", the function is called as NewPrefix("!").
func NewPrefix(prefixes ...string) Prefixer {
return func(msg *gateway.MessageCreateEvent) (string, bool) {
for _, prefix := range prefixes {
if strings.HasPrefix(msg.Content, prefix) {
return prefix, true
}
}
return "", false
}
}
// ArgsParser is the function type for parsing message content into fields,
// usually delimited by spaces.
type ArgsParser func(content string) ([]string, error)
// DefaultArgsParser implements a parser similar to that of shell's,
// implementing quotes as well as escapes.
var DefaultArgsParser = shellwords.Parse
// NewShardFunc creates a shard constructor that shares the same internal store.
// If opts sets its own cabinet, then a new store isn't created.
func NewShardFunc(fn func(*state.State) (*Context, error)) shard.NewShardFunc {
if fn == nil {
panic("bot.NewShardFunc missing fn")
}
return func(m *shard.Manager, id *gateway.Identifier) (shard.Shard, error) {
sessn := session.NewCustom(*id, api.NewClient(id.Token), handler.New())
state := state.NewFromSession(sessn, defaultstore.New())
bot, err := fn(state)
if err != nil {
return nil, fmt.Errorf("failed to create bot instance: %w", err)
}
return bot, nil
}
}
// Context is the bot state for commands and subcommands.
//
// # Commands
//
// A command can be created by making it a method of Commands, or whatever
// struct was given to the constructor. This following example creates a command
// with a single integer argument (which can be ran with "~example 123"):
//
// func (c *Commands) Example(
// m *gateway.MessageCreateEvent, i int) (string, error) {
//
// return fmt.Sprintf("You sent: %d", i)
// }
//
// Commands' exported methods will all be used as commands. Messages are parsed
// with its first argument (the command) mapped accordingly to c.MapName, which
// capitalizes the first letter automatically to reflect the exported method
// name.
//
// A command can either return either an error, or data and error. The only data
// types allowed are string, *discord.Embed, and *api.SendMessageData. Any other
// return types will invalidate the method.
//
// # Events
//
// An event can only have one argument, which is the pointer to the event
// struct. It can also only return error.
//
// func (c *Commands) Example(o *gateway.TypingStartEvent) error {
// log.Println("Someone's typing!")
// return nil
// }
type Context struct {
*Subcommand
*state.State
// Descriptive (but optional) bot name
Name string
// Descriptive help body
Description string
// Called to parse message content, default to DefaultArgsParser().
ParseArgs ArgsParser
// Called to check a message's prefix. The default prefix is "!". Refer to
// NewPrefix().
HasPrefix Prefixer
// AllowBot makes the router also process MessageCreate events from bots.
// This is false by default and only applies to MessageCreate.
AllowBot bool
// QuietUnknownCommand, if true, will not make the bot reply with an unknown
// command error into the chat. This will apply to all other subcommands.
// SilentUnknown controls whether or not an UnknownCommandError should be
// returned (instead of a silent error).
SilentUnknown struct {
// Command when true will silent only unknown commands. Known
// subcommands with unknown commands will still error out.
Command bool
// Subcommand when true will suppress unknown subcommands.
Subcommand bool
}
// FormatError formats any errors returned by anything, including the method
// commands or the reflect functions. This also includes invalid usage
// errors or unknown command errors. Returning an empty string means
// ignoring the error.
//
// By default, this field replaces all @ with @\u200b, which prevents an
// @everyone mention.
FormatError func(error) string
// ErrorLogger logs any error that anything makes and the library can't
// reply to the client. This includes any event callback errors that aren't
// Message Create.
ErrorLogger func(error)
// ReplyError when true replies to the user the error. This only applies to
// MessageCreate events.
ReplyError bool
// ErrorReplier is an optional function that allows changing how the error
// is replied. It overrides ReplyError and is only used for MessageCreate
// events.
//
// Note that errors that are passed in here will bypas FormatError; in other
// words, the implementation might only care about ErrorReplier and leave
// FormatError as it is.
ErrorReplier func(err error, src *gateway.MessageCreateEvent) api.SendMessageData
// EditableCommands when true will also listen for MessageUpdateEvent and
// treat them as newly created messages. This is convenient if you want
// to quickly edit a message and re-execute the command.
EditableCommands bool
// Subcommands contains all the registered subcommands. This is not
// exported, as it shouldn't be used directly.
subcommands []*Subcommand
// Quick access map from event types to pointers. This map will never have
// MessageCreateEvent's type.
typeCache sync.Map // map[reflect.Type][]*CommandContext
stopFunc func() // unbind function, see Start()
}
// Start quickly starts a bot with the given command. It will prepend "Bot"
// into the token automatically. Refer to example/ for usage.
func Start(
token string, cmd interface{},
opts func(*Context) error) (wait func() error, err error) {
if token == "" {
return nil, errors.New("token is not given")
}
if !strings.HasPrefix(token, "Bot ") {
token = "Bot " + token
}
newShard := NewShardFunc(func(s *state.State) (*Context, error) {
ctx, err := New(s, cmd)
if err != nil {
return nil, err
}
// fail api request if they (will) take up more than 5 minutes
ctx.Client.Client.Timeout = 5 * time.Minute
if opts != nil {
if err := opts(ctx); err != nil {
return nil, err
}
}
ctx.AddIntents(ctx.DeriveIntents())
ctx.AddIntents(gateway.IntentGuilds) // for channel event caching
return ctx, nil
})
m, err := shard.NewManager(token, newShard)
if err != nil {
return nil, fmt.Errorf("failed to create shard manager: %w", err)
}
if err := m.Open(context.Background()); err != nil {
return nil, fmt.Errorf("failed to open: %w", err)
}
return func() error {
WaitForInterrupt()
// Close the shards first.
closeErr := m.Close()
// Remove all handlers to clean up.
m.ForEach(func(s shard.Shard) {
ctx := s.(*Context)
stop := ctx.Start()
stop()
})
return closeErr
}, nil
}
// Run starts the bot, prints a message into the console, and blocks until
// SIGINT. "Bot" is prepended into the token automatically, similar to Start.
// The function will call os.Exit(1) on an initialization or cleanup error.
func Run(token string, cmd interface{}, opts func(*Context) error) {
wait, err := Start(token, cmd, opts)
if err != nil {
log.Fatalln("failed to start:", err)
}
log.Println("Bot is running.")
if err := wait(); err != nil {
log.Fatalln("cleanup error:", err)
}
}
// WaitForInterrupt blocks until SIGINT.
func WaitForInterrupt() {
sigs := make(chan os.Signal, 1)
signal.Notify(sigs, os.Interrupt, syscall.SIGTERM)
<-sigs
}
// New makes a new context with a "~" as the prefix. cmds must be a pointer to a
// struct with a *Context field. Example:
//
// type Commands struct {
// Ctx *Context
// }
//
// cmds := &Commands{}
// c, err := bot.New(session, cmds)
//
// The default prefix is "~", which means commands must start with "~" followed
// by the command name in the first argument, else it will be ignored.
//
// c.Start() should be called afterwards to actually handle incoming events.
func New(s *state.State, cmd interface{}) (*Context, error) {
c, err := NewSubcommand(cmd)
if err != nil {
return nil, err
}
ctx := &Context{
Subcommand: c,
State: s,
ParseArgs: DefaultArgsParser,
HasPrefix: NewPrefix("~"),
FormatError: func(err error) string {
// Escape all pings, including @everyone.
return strings.Replace(err.Error(), "@", "@\u200b", -1)
},
ErrorLogger: func(err error) {
log.Println("Bot error:", err)
},
ReplyError: true,
}
if err := ctx.InitCommands(ctx); err != nil {
return nil, fmt.Errorf("failed to initialize with given cmds: %w", err)
}
return ctx, nil
}
// AddIntents adds the given Gateway Intent into the Gateway. This is a
// convenient function that calls Gateway's AddIntent.
func (ctx *Context) AddIntents(i gateway.Intents) {
ctx.Session.AddIntents(i)
}
// Subcommands returns the slice of subcommands. To add subcommands, use
// RegisterSubcommand().
func (ctx *Context) Subcommands() []*Subcommand {
// Getter is not useless, refer to the struct doc for reason.
return ctx.subcommands
}
// FindMethod finds a method based on the struct and method name. The queried
// names will have their flags stripped.
//
// // Find a command from the main context:
// cmd := ctx.FindMethod("", "Method")
// // Find a command from a subcommand:
// cmd = ctx.FindMethod("Starboard", "Reset")
func (ctx *Context) FindCommand(structName, methodName string) *MethodContext {
if structName == "" {
return ctx.Subcommand.FindCommand(methodName)
}
for _, sub := range ctx.subcommands {
if sub.StructName == structName {
return sub.FindCommand(methodName)
}
}
return nil
}
// MustRegisterSubcommand tries to register a subcommand, and will panic if it
// fails. This is recommended, as subcommands won't change after initializing
// once in runtime, thus fairly harmless after development.
//
// If no names are given or if the first name is empty, then the subcommand name
// will be derived from the struct name. If one name is given, then that name
// will override the struct name. Any other name values will be aliases.
//
// It is recommended to use this method to add subcommand aliases over manually
// altering the Aliases slice of each Subcommand, as it does collision checks
// against other subcommands as well.
func (ctx *Context) MustRegisterSubcommand(cmd interface{}, names ...string) *Subcommand {
s, err := ctx.RegisterSubcommand(cmd, names...)
if err != nil {
panic(err)
}
return s
}
// RegisterSubcommand registers and adds cmd to the list of subcommands. It will
// also return the resulting Subcommand. Refer to MustRegisterSubcommand for the
// names argument.
func (ctx *Context) RegisterSubcommand(cmd interface{}, names ...string) (*Subcommand, error) {
s, err := NewSubcommand(cmd)
if err != nil {
return nil, fmt.Errorf("failed to add subcommand: %w", err)
}
// Register the subcommand's name.
s.NeedsName()
if len(names) > 0 && names[0] != "" {
s.Command = names[0]
}
if len(names) > 1 {
// Copy the slice for expected behaviors.
s.Aliases = append([]string(nil), names[1:]...)
}
if err := s.InitCommands(ctx); err != nil {
return nil, fmt.Errorf("failed to initialize subcommand: %w", err)
}
// Check if the existing command name already exists. This could really be
// optimized, but since it's in a cold path, who cares.
var subcommandNames = append([]string{s.Command}, s.Aliases...)
for _, name := range subcommandNames {
for _, sub := range ctx.subcommands {
// Check each alias against the subcommand name.
if sub.Command == name {
return nil, fmt.Errorf("new subcommand has duplicate name: %q", name)
}
// Also check each alias against other subcommands' aliases.
for _, subalias := range sub.Aliases {
if subalias == name {
return nil, fmt.Errorf("new subcommand has duplicate alias: %q", name)
}
}
}
}
ctx.subcommands = append(ctx.subcommands, s)
return s, nil
}
// emptyMentionTypes is used by Start() to not parse any mentions.
var emptyMentionTypes = []api.AllowedMentionType{}
// Start adds itself into the session handlers. If Start is called more than
// once, then it does nothing. The caller doesn't have to call Start if they
// call Open.
//
// The returned function is a delete function, which removes itself from the
// Session handlers. The delete function is not safe to use concurrently.
func (ctx *Context) Start() func() {
if ctx.stopFunc == nil {
cancel := ctx.State.AddHandler(func(v interface{}) {
if err := ctx.callCmd(v); err != nil {
ctx.ErrorLogger(fmt.Errorf("command error: %w", err))
}
})
ctx.stopFunc = func() {
cancel()
ctx.stopFunc = nil
}
}
return ctx.stopFunc
}
// Open starts the bot context and the gateway connection. It automatically
// binds the needed handlers.
func (ctx *Context) Open(cancelCtx context.Context) error {
ctx.Start()
return ctx.State.Open(cancelCtx)
}
// Call should only be used if you know what you're doing.
func (ctx *Context) Call(event interface{}) error {
return ctx.callCmd(event)
}
// Help generates a full Help message. It serves mainly as a reference for
// people to reimplement and change. It doesn't show hidden commands.
func (ctx *Context) Help() string {
return ctx.HelpGenerate(false)
}
// HelpGenerate generates a full Help message. It serves mainly as a reference
// for people to reimplement and change. If showHidden is true, then hidden
// subcommands and commands will be shown.
func (ctx *Context) HelpGenerate(showHidden bool) string {
// Generate the header.
buf := strings.Builder{}
buf.WriteString("__Help__")
// Name an
if ctx.Name != "" {
buf.WriteString(": " + ctx.Name)
}
if ctx.Description != "" {
buf.WriteString("\n" + IndentLines(ctx.Description))
}
// Separators
buf.WriteString("\n---\n")
// Generate all commands
if help := ctx.Subcommand.Help(); help != "" {
buf.WriteString("__Commands__\n")
buf.WriteString(IndentLines(help))
buf.WriteByte('\n')
}
var subcommands = ctx.Subcommands()
var subhelps = make([]string, 0, len(subcommands))
for _, sub := range subcommands {
if sub.Hidden && !showHidden {
continue
}
help := sub.HelpShowHidden(showHidden)
if help == "" {
continue
}
help = IndentLines(help)
builder := strings.Builder{}
builder.WriteString("**")
builder.WriteString(sub.Command)
builder.WriteString("**")
for _, alias := range sub.Aliases {
builder.WriteString("|")
builder.WriteString("**")
builder.WriteString(alias)
builder.WriteString("**")
}
if sub.Description != "" {
builder.WriteString(": ")
builder.WriteString(sub.Description)
}
builder.WriteByte('\n')
builder.WriteString(help)
subhelps = append(subhelps, builder.String())
}
if len(subhelps) > 0 {
buf.WriteString("---\n")
buf.WriteString("__Subcommands__\n")
buf.WriteString(IndentLines(strings.Join(subhelps, "\n")))
}
return buf.String()
}
// IndentLine prefixes every line from input with a single-level indentation.
func IndentLines(input string) string {
const indent = " "
var lines = strings.Split(input, "\n")
for i := range lines {
lines[i] = indent + lines[i]
}
return strings.Join(lines, "\n")
}
// DeriveIntents derives all possible gateway intents from this context and all
// its subcommands' method handlers and middlewares.
func (ctx *Context) DeriveIntents() gateway.Intents {
var intents = ctx.Subcommand.DeriveIntents()
for _, subcmd := range ctx.subcommands {
intents |= subcmd.DeriveIntents()
}
return intents
}