bolt/bolt.go

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)
	}

	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:
	}

	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 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
}