bolt/bolt.go

package bolt

import (
	"errors"
	"fmt"
	"log"
	"os"
	"os/signal"
	"slices"
	"strings"
	"syscall"
	"time"

	dg "github.com/bwmarrin/discordgo"
)

const (
	TOKEN_ENV_VAR = "DISCORD_TOKEN" //label for token environment variable

	BOT_INTENTS = dg.IntentGuilds |
		dg.IntentGuildMembers |
		dg.IntentGuildPresences |
		dg.IntentMessageContent |
		dg.IntentsGuildMessages |
		dg.IntentGuildMessageReactions
)

// basic bot structure containing discordgo connection as well as the command map
type bolt struct {
	*dg.Session                    //holds discordgo internals
	commands    map[string]Command //maps trigger phrase to command struct for fast lookup
	indicator   string             //the indicator used to detect whether a message is a command
	logLvl      LogLevel           //determines how much the bot logs
}

type Bolt interface {
	Start() error
	AddCommands(cmd ...Command)
	//filtered methods
	stop() error
	messageHandler(s *dg.Session, msg *dg.MessageCreate)
	handleCommand(msgEvent *MessageCreateEvent, s *dg.Session, lg int) error
	createReply(content, message, channel, guild string) *dg.MessageSend
	getRemainingTimeout(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)
	joinUserVoice(guild, user string, s *dg.Session) (*dg.VoiceConnection, error)
}

// create a new bolt interface
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)
	}
	bot.Identify.Intents = BOT_INTENTS

	b := &bolt{
		Session:  bot,
		commands: make(map[string]Command, 0),
		logLvl:   LogLevelAll,
	}
	//set default command indicator
	b.indicator = "."

	//apply options
	for _, opt := range opts {
		opt(b)
	}

	return b, nil
}

// starts the bot, commands are added and the connection to Discord is opened, this is a BLOCKING call
// that handles safe shutdown of the bot
func (b *bolt) Start() error {
	//register commands and open connection
	b.AddHandler(b.messageHandler)

	err := b.Open()
	if err != nil {
		return fmt.Errorf("failed to open websocket connection with Discord: %e", err)
	}

	//safe shutdown handler
	sigChannel := make(chan os.Signal, 1)
	signal.Notify(sigChannel, syscall.SIGINT)
	<-sigChannel

	if err := b.stop(); err != nil {
		return err
	}

	return nil
}

// stops the bot
func (b *bolt) stop() error {
	return b.Close()
}

// adds commands to bot command map for use
func (b *bolt) AddCommands(cmd ...Command) {
	for _, c := range cmd {
		b.commands[c.Trigger] = c
	}
}

// handler function that parses message data, handles logging the message based on logLevel, and executes
// the payload function in a goroutine
func (b *bolt) messageHandler(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
	}

	//if there is no content it is likely an image, gif, or sticker, updating message content for
	//better logging and to avoid confusion
	if len(msg.Content) == 0 {
		msg.Content = "[Embedded Content]"
	}

	//the bot will ignore it's own messages to prevent command loops
	if msg.Author.ID == s.State.User.ID {
		if b.logLvl == LogLevelCmd || b.logLvl == LogLevelAll {
			//log command responses
			log.Printf("< %s | %s | %s > %s\n", server.Name, channel.Name, msg.Author.Username, msg.Content)
		}
		return
	}

	if b.logLvl == LogLevelAll {
		//log message
		log.Printf("< %s | %s | %s > %s\n", server.Name, channel.Name, msg.Author.Username, msg.Content)
	}

	//does the message have the command indicator
	lg := len(b.indicator)
	if msg.Content[:lg] == b.indicator {
		mCreate := &MessageCreateEvent{
			AuthorUsername: msg.Author.Username,
			AuthorID:       msg.Author.ID,
			AuthorRoles:    msg.Member.Roles,
			MsgID:          msg.ID,
			Msg:            msg.Content,
			MsgChanID:      msg.ChannelID,
			MsgGuildID:     msg.GuildID,
			MsgAttachments: msg.Attachments,
		}

		if b.logLvl == LogLevelCmd {
			//log commands
			log.Printf("< %s | %s | %s > %s\n", mCreate.MsgGuildName, mCreate.MsgChanName, mCreate.AuthorUsername, mCreate.Msg)
		}

		//handled in its own goroutine to allow for async commands
		go func() {
			err := b.handleCommand(mCreate, s, lg)
			if err != nil {
				log.Println(err)
			}
		}()
	}
}

// parses command from message and handles timeout checks, role checks, and command execution. All command responses are sent back to Discord
func (b *bolt) handleCommand(msgEvent *MessageCreateEvent, s *dg.Session, lg int) error {
	words := strings.Split(msgEvent.Msg, " ")
	run, ok := b.commands[words[0][lg:]]
	if !ok {
		return nil //command doesn't exist, maybe log or respond to author
	}

	//has command met its timeout requirements
	tc, err := b.timeoutCheck(msgEvent.MsgID, msgEvent.MsgChanID, msgEvent.MsgGuildID, s, run)
	if err != nil {
		return fmt.Errorf("failed to calculate timeout for %s\n%e", run.Trigger, err)
	}
	if !tc {
		return nil
	}

	//does user have correct permissions
	if run.Roles != nil {
		check, err := b.roleCheck(msgEvent.MsgGuildID, msgEvent.AuthorRoles, s, run)
		if err != nil {
			return fmt.Errorf("failed to perform permission checks for %s\n%e", run.Trigger, err)
		}
		if !check {
			reply := b.createReply("you do not have permissions to run that command", msgEvent.MsgID, msgEvent.MsgChanID, msgEvent.MsgGuildID)
			_, err := s.ChannelMessageSendComplex(msgEvent.MsgChanID, reply)
			if err != nil {
				return err
			}
			return nil
		}
	}

	//populate message struct exposed to client
	plMsg := Message{
		Author:    msgEvent.AuthorUsername,
		ID:        msgEvent.AuthorID,
		msgID:     msgEvent.MsgID,
		Words:     words,
		Content:   msgEvent.Msg,
		Channel:   msgEvent.MsgChanName,
		channelID: msgEvent.MsgChanID,
		Server:    msgEvent.MsgGuildName,
		serverID:  msgEvent.MsgGuildID,
		sesh:      b,
	}

	//check for file attachments
	if len(msgEvent.MsgAttachments) > 0 {
		var att []MessageAttachment
		for _, a := range msgEvent.MsgAttachments {
			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,
			})
		}

		plMsg.Attachments = att
	}

	//run command payload
	err = run.Payload(plMsg)
	if err != nil {
		return fmt.Errorf("failed to execute payload function: %e", err)
	}

	//update run time
	run.lastRun = time.Now()
	b.commands[run.Trigger] = run
	return nil
}

// basic wrapper function to create easy Discord responses
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,
	}
}

// used to calculate the remaining time left in a timeout and returning it in a human-readable format
func (b *bolt) getRemainingTimeout(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"
		}
	}

	return fmt.Sprintf("%d%s", timeLeft, metric)
}

// checks if the author of msg has the correct role to run the requested command
func (b *bolt) roleCheck(guild string, roles []string, s *dg.Session, run Command) (bool, error) {
	var found bool
	//loop thru author roles, there may be a better way to check for this UNION
	//TODO: improve role search performance to support bigger lists
	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
		}
	}

	//can't find role, don't run command
	if !found {
		return false, nil
	}

	return true, nil
}

// check if the command timeout has been met, responding with remaining time if timeout has not been met yet.
func (b *bolt) timeoutCheck(msgID, channelID, guildID string, s *dg.Session, run Command) (bool, error) {
	wait := run.lastRun.Add(run.Timeout)
	if !time.Now().After(wait) {
		reply := b.createReply(fmt.Sprintf("that command cannot be run for another %s", b.getRemainingTimeout(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
}

func (b *bolt) joinUserVoice(guild, user string, s *dg.Session) (*dg.VoiceConnection, error) {
	g, err := s.State.Guild(guild)
	if err != nil {
		return nil, err
	}

	var chanID string
	for _, x := range g.VoiceStates {
		if user == x.UserID {
			chanID = x.ChannelID
		}
	}

	if chanID == "" {
		return nil, errors.New("user is not in a voice channel")
	}

	//joining voice channel: false = not muted, true = deafened
	dgv, err := s.ChannelVoiceJoin(guild, chanID, false, true)
	if err != nil {
		return nil, err
	}

	return dgv, nil
}