package bolt import ( "context" "fmt" "log" "os" "os/signal" "slices" "strings" "sync" "time" dg "github.com/bwmarrin/discordgo" ) const ( //the name of the environment variable that should contain the token for the bot, it is //required for bolt to run TOKEN_ENV_VAR = "DISCORD_TOKEN" //bot default command indicator, if messages begin with this substring they are processed //through the command handler instead of the generic message handler DEFAULT_INDICATOR = "." //max amount of concurrent goroutines that bolt can use for events. A lower amount //may lower the resource usage of bolt but may cause a delay in event handling DEFAULT_MAX_GOROUTINES = 250 //minimum intents for bots to function, intents can be changed with options DEFAULT_INTENTS = dg.IntentGuilds | dg.IntentGuildMembers | dg.IntentGuildPresences | dg.IntentMessageContent | dg.IntentsGuildMessages | dg.IntentGuildMessageReactions ) type bolt struct { //discordgo internals *dg.Session //maps trigger phrase to command struct for instant lookup commands map[string]Command //used to detect whether a message is a command indicator string //verbosity of logs bolt outputs logLvl LogLevel //waitgroup for event routines wg sync.WaitGroup //pool is a buffered channel used as a semaphore for event handler routines, it is limited to //only spawn maxRoutines to handle events pool chan struct{} maxRoutines int //generic message handler func msgHandlerf Payload } type Bolt interface { Start() error AddCommands(cmd ...Command) AddMessageHandler(p Payload) //filtered methods stop() error msgEventHandler(s *dg.Session, msg *dg.MessageCreate) mapEventToMsg(msg *dg.MessageCreate, channel *dg.Channel, server *dg.Guild) Message handleCommand(msgEvent *Message, lg int) error handleMessage(event *Message) error createReply(content, message, channel, guild string) *dg.MessageSend remainingTimeout(timeout time.Time) string roleCheck(guild string, roles []string, s *dg.Session, run Command) (bool, error) timeoutCheck(msgID, channelID, guildID string, s *dg.Session, run Command) (bool, error) } // New creates a new bolt instance and applies any supplied options func New(opts ...Option) (Bolt, error) { _, check := os.LookupEnv(TOKEN_ENV_VAR) if !check { return nil, fmt.Errorf("environment variable %s must be set", TOKEN_ENV_VAR) } bot, err := dg.New(fmt.Sprintf("Bot %s", os.Getenv(TOKEN_ENV_VAR))) if err != nil { return nil, fmt.Errorf("failed to create Discord session: %e", err) } b := &bolt{ Session: bot, commands: make(map[string]Command, 0), logLvl: LogLevelAll, indicator: DEFAULT_INDICATOR, wg: sync.WaitGroup{}, maxRoutines: DEFAULT_MAX_GOROUTINES, } b.Identify.Intents = DEFAULT_INTENTS for _, opt := range opts { opt(b) } //options can change max routine number, so create after b.pool = make(chan struct{}, b.maxRoutines) return b, nil } // Start applies the message event handler function to the bot and opens the initial websocket connection // with Discord. Start is a blocking call that also handles safe shutdown, on Interrupt, bolt will give // command routines a window to finish before closing the connection. func (b *bolt) Start() error { b.AddHandler(b.msgEventHandler) err := b.Open() if err != nil { return fmt.Errorf("failed to open websocket connection with Discord: %e", err) } log.Println("bot started") sigChannel := make(chan os.Signal, 1) signal.Notify(sigChannel, os.Interrupt) <-sigChannel //give handler routines a 5 second window to finish processes before closing connection ctx, cancel := context.WithTimeout(context.Background(), time.Second*5) defer cancel() closeChan := make(chan struct{}, 0) go func() { b.wg.Wait() close(closeChan) }() select { case <-ctx.Done(): log.Println("shutdown timed out waiting for handlers to finish, some may have been incomplete") case <-closeChan: log.Println("handler routines cleaned") } log.Println("exiting") return b.stop() } // AddCommands registers command handlers. Any messages that begin with the command indicator will be forwarded // to the handler func (b *bolt) AddCommands(cmd ...Command) { for _, c := range cmd { b.commands[c.Trigger] = c } } // AddMessageHandler registers the generic message handler, any messages that are not commands will be forwarded // to the handler func (b *bolt) AddMessageHandler(p Payload) { b.msgHandlerf = p } // stop closes the websocket connection to Discord func (b *bolt) stop() error { return b.Close() } // msgEventHandler is a beefy boy that handles message logging, command parsing, and executing payload functions. It needs cleanup then // i'll worry about this comment // msgEventHandler handles the routing of messages to either the command or message handlers. If LogLvl is set, logging is handled here before // the event is mapped to the Message struct and forwarded to the handlers. Each message is handled in its own goroutine, the max allowed goroutines // is set as a default and can be altered using options for better performance. func (b *bolt) msgEventHandler(s *dg.Session, msg *dg.MessageCreate) { //get server information server, err := s.Guild(msg.GuildID) if err != nil { log.Printf("failed to get guild: %e\n", err) return } channel, err := s.Channel(msg.ChannelID) if err != nil { log.Printf("failed to get channel from guild: %e\n", err) return } //the bot will ignore it's own messages to prevent command loops if msg.Author.ID == s.State.User.ID { if b.logLvl != LogLevelErr { log.Printf("< %s | %s | %s > %s\n", server.Name, channel.Name, msg.Author.Username, msg.Content) } return } if b.logLvl == LogLevelAll { log.Printf("< %s | %s | %s > %s\n", server.Name, channel.Name, msg.Author.Username, msg.Content) } m := b.mapEventToMsg(msg, channel, server) lg := len(b.indicator) if msg.Content[:lg] == b.indicator { if b.logLvl == LogLevelCmd { log.Printf("< %s | %s | %s > %s\n", m.Server, m.Channel, m.Author.Name, m.Content) } b.pool <- struct{}{} //'aquire' a routine b.wg.Go(func() { err := b.handleCommand(&m, lg) if err != nil { log.Println(err) } <-b.pool //release routine }) } else { b.pool <- struct{}{} //'aquire' a routine b.wg.Go(func() { err := b.handleMessage(&m) if err != nil { log.Println(err) } <-b.pool //release routine }) } } // mapEventToMsg maps the Discord event message struct into bolt's user-friendly Message struct func (b *bolt) mapEventToMsg(msg *dg.MessageCreate, channel *dg.Channel, server *dg.Guild) Message { m := Message{ Author: Author{ Name: msg.Author.Username, ID: msg.Author.ID, Roles: msg.Member.Roles, }, ID: msg.ID, Content: msg.Content, Channel: channel.Name, ChannelID: channel.ID, Server: server.Name, ServerID: server.ID, } w := strings.Fields(msg.Content) if len(w) > 0 { m.Words = w } if len(msg.Mentions) > 0 { m.Mentions = msg.Mentions } if len(msg.Attachments) > 0 { var att []MessageAttachment for _, a := range msg.Attachments { att = append(att, MessageAttachment{ ID: a.ID, URL: a.URL, ProxyURL: a.ProxyURL, Filename: a.Filename, ContentType: a.ContentType, Width: a.Width, Height: a.Height, Size: a.Size, DurationSecs: a.DurationSecs, }) } m.Attachments = att } return m } // handleMessage forwards the message data to the handler function, if one is set func (b *bolt) handleMessage(event *Message) error { if b.msgHandlerf != nil { return b.msgHandlerf(&Context{ Message: event, bolt: b, }) } return nil } // handleCommand maps the first word of the message to the command payload, if it exists. It then forwards the message // data to the handler after checking the timeout and role restrictions. If restrictions have not been met a generic // response is sent to the message // TODO: accept a string for timeout/role rejection messages to allow customization func (b *bolt) handleCommand(msg *Message, lg int) error { run, ok := b.commands[msg.Words[0][lg:]] if !ok { return nil //command doesn't exist } //has command met its timeout requirements, if timeout has not expired a response is sent to the message //from the timeoutCheck method tc, err := b.timeoutCheck(msg.ID, msg.ChannelID, msg.ServerID, b.Session, run) if err != nil { return fmt.Errorf("failed to calculate timeout for %s\n%e", run.Trigger, err) } //method handles sending a response, so we do not need to send another here, simply return if !tc { return nil } //does user have correct permissions to run this command if run.Roles != nil { check, err := b.roleCheck(msg.ServerID, msg.Author.Roles, b.Session, run) if err != nil { return fmt.Errorf("failed to perform permission checks for %s\n%e", run.Trigger, err) } if !check { //this alert should probably be moved into the roleCheck function to match the pattern of the timeoutCheck method reply := b.createReply("you do not have permissions to run that command", msg.ID, msg.ChannelID, msg.ServerID) _, err := b.Session.ChannelMessageSendComplex(msg.ChannelID, reply) if err != nil { return err } return nil } } //execute handler func with message context err = run.Payload(&Context{ Message: msg, bolt: b, }) if err != nil { return fmt.Errorf("encountered an error while handling command (%s): %e", msg.Words[0], err) } //update run time run.lastRun = time.Now() b.commands[run.Trigger] = run return nil } // createReploy is a basic wrapper function to map message data to the actual MessageSend struct func (b *bolt) createReply(content, message, channel, guild string) *dg.MessageSend { details := &dg.MessageReference{ MessageID: message, ChannelID: channel, GuildID: guild, } return &dg.MessageSend{ Content: content, Reference: details, } } // remainingTimeout calculates the amount of time left before a command can be run again. Returning a string // representing a readable version of the time left, example: 1h (1 hour) func (b *bolt) remainingTimeout(timeout time.Time) string { r := time.Until(timeout) var ( timeLeft int metric string ) timeLeft = int(r.Hours()) metric = "h" if timeLeft < 1 { timeLeft = int(r.Minutes()) metric = "m" if timeLeft < 1 { timeLeft = int(r.Seconds()) metric = "s" } } //provide a user-friendly time string to send back to the user return fmt.Sprintf("%d%s", timeLeft, metric) } // checks if the author of msg has the correct role to run the requested command // roleCheck loops through the provided user role ID's grabs the role "name" and ensures the role is present // in the commands whitelist. If the role exists in the Command whitelist a true response is returned func (b *bolt) roleCheck(guild string, roles []string, s *dg.Session, run Command) (bool, error) { var found bool //this needs a bit of love, looping through roles could get messy as server and role lists grow on big servers //especially with the Role() call inside of the loop for _, r := range roles { //get role name from ID n, err := s.State.Role(guild, r) if err != nil { return false, fmt.Errorf("failed to get role from ID %s\n%e", guild, err) } //does this role exist in command roles check := slices.Contains(run.Roles, n.Name) if check { found = true break } } return found, nil } // timeoutCheck determines if a Command timeout has expired, if it hasn't expired a response is sent back to the message alerting the user // of the remaining time func (b *bolt) timeoutCheck(msgID, channelID, guildID string, s *dg.Session, run Command) (bool, error) { wait := run.lastRun.Add(run.Timeout) now := time.Now() if !now.After(wait) && !now.Equal(wait) { reply := b.createReply(fmt.Sprintf("that command cannot be run for another %s", b.remainingTimeout(wait)), msgID, channelID, guildID) _, err := s.ChannelMessageSendComplex(channelID, reply) if err != nil { return false, fmt.Errorf("failed to send timeout response: %e", err) } return false, nil } return true, nil }