canopeas/discord/discord.go

package discord

import (
	"strings"
	"time"

	"go.etztech.xyz/sedbot/config"
	"go.etztech.xyz/sedbot/database"

	"github.com/bwmarrin/discordgo"
	"go.jolheiser.com/beaver"
)

// Register commands to this map
var (
	commands       = make(map[string]command)
	messageRoleMap = make(map[string]map[string]string)

	memeRateLimit *rateLimit
)

type commandInit struct {
	session  *discordgo.Session
	message  *discordgo.Message
	config   *config.Config
	database *database.Database
}

type command struct {
	validate func(cmd commandInit) bool
	run      func(cmd commandInit) (string, error)
	help     string
}

func Bot(cfg *config.Config, db *database.Database) (*discordgo.Session, error) {
	bot, err := discordgo.New("Bot " + cfg.Token)
	if err != nil {
		return nil, err
	}

	for _, messageRole := range cfg.MessageRoles {
		if messageRoleMap[messageRole.MessageID] == nil {
			messageRoleMap[messageRole.MessageID] = make(map[string]string)
		}
		_ = bot.MessageReactionAdd(messageRole.ChannelID, messageRole.MessageID, messageRole.Emoji)
		messageRoleMap[messageRole.MessageID][messageRole.Emoji] = messageRole.RoleID
	}

	Echo(cfg)
	bot.AddHandler(readyHandler())
	bot.AddHandler(commandHandler(cfg, db))
	bot.AddHandler(messageHandler(cfg, db))
	bot.AddHandler(reactionAddHandler())
	bot.AddHandler(reactionRemoveHandler())

	// Rate limits
	d, err := time.ParseDuration(cfg.MemeRate)
	if err != nil {
		return nil, err
	}
	memeRateLimit = NewRateLimit(d)

	return bot, nil
}

func sendTyping(s *discordgo.Session, channelID string) {
	if err := s.ChannelTyping(channelID); err != nil {
		beaver.Errorf("could not send typing status: %v", err)
	}
}

func sendMessage(s *discordgo.Session, channelID, content string, scrub bool) *discordgo.Message {
	var msg *discordgo.Message
	var err error
	if scrub {
		msg, err = s.ChannelMessageSendComplex(channelID, &discordgo.MessageSend{
			Content:         content,
			AllowedMentions: &discordgo.MessageAllowedMentions{},
		})
	} else {
		msg, err = s.ChannelMessageSend(channelID, content)
	}
	if err != nil {
		beaver.Errorf("could not send message: %v", err)
		return nil
	}
	return msg
}

func sendEmbed(s *discordgo.Session, channelID string, embed *discordgo.MessageEmbed) *discordgo.Message {
	msg, err := s.ChannelMessageSendEmbed(channelID, embed)
	if err != nil {
		beaver.Errorf("could not send embed: %v", err)
		return nil
	}
	return msg
}

func isStaff(authorRoleIDs, staffRoleIDs []string) bool {
	for _, aRole := range authorRoleIDs {
		for _, sRole := range staffRoleIDs {
			if aRole == sRole {
				return true
			}
		}
	}
	return false
}

func readyHandler() func(s *discordgo.Session, m *discordgo.Ready) {
	return func(s *discordgo.Session, r *discordgo.Ready) {
		beaver.Infof("https://discord.com/api/oauth2/authorize?client_id=%s&permissions=0&redirect_uri=https://birbmc.com&scope=bot", r.User.ID)
	}
}

func commandHandler(cfg *config.Config, db *database.Database) func(s *discordgo.Session, m *discordgo.MessageCreate) {
	return func(s *discordgo.Session, m *discordgo.MessageCreate) {
		// Ignore bots
		if m.Author.Bot {
			return
		}

		// Check prefix
		if !strings.HasPrefix(m.Content, cfg.Prefix) {
			return
		}

		content := m.Content[1:]
		args := strings.Fields(content)
		if len(args) == 0 {
			return
		}

		cmdArg := strings.ToLower(args[0])

		cmd, ok := commands[cmdArg]
		if !ok {
			return
		}

		cmdInit := commandInit{
			session:  s,
			message:  m.Message,
			config:   cfg,
			database: db,
		}
		if !cmd.validate(cmdInit) {
			sendMessage(s, m.ChannelID, "You cannot run this command.", false)
			return
		}
		feedback, err := cmd.run(cmdInit)
		if err != nil {
			feedback = "Internal error"
			beaver.Errorf("error while running %s: %v", cmdArg, err)
		}
		if len(feedback) > 0 {
			sendMessage(s, m.ChannelID, feedback, false)
		}
	}
}

func messageHandler(cfg *config.Config, db *database.Database) func(s *discordgo.Session, m *discordgo.MessageCreate) {
	return func(s *discordgo.Session, m *discordgo.MessageCreate) {
		// Ignore bots
		if m.Author.Bot {
			return
		}

		// [FIRED] increment
		for _, role := range m.MentionRoles {
			if cfg.FiredRole == role {
				if err := db.IncrementPing(cfg.FiredRole); err != nil {
					beaver.Errorf("could not increment ping for %s: %v", cfg.FiredRole, err)
				}
			}
		}
	}
}

func reactionAddHandler() func(s *discordgo.Session, m *discordgo.MessageReactionAdd) {
	return func(s *discordgo.Session, m *discordgo.MessageReactionAdd) {
		reactionHandler(true, s, m.MessageReaction)
	}
}

func reactionRemoveHandler() func(s *discordgo.Session, m *discordgo.MessageReactionRemove) {
	return func(s *discordgo.Session, m *discordgo.MessageReactionRemove) {
		reactionHandler(false, s, m.MessageReaction)
	}
}

func reactionHandler(add bool, s *discordgo.Session, m *discordgo.MessageReaction) {
	if _, ok := messageRoleMap[m.MessageID]; ok {
		if r, ok := messageRoleMap[m.MessageID][m.Emoji.APIName()]; ok {
			var roleCmd func(guildID, userID, roleID string) (err error)
			if add {
				roleCmd = s.GuildMemberRoleAdd
			} else {
				roleCmd = s.GuildMemberRoleRemove
			}
			if err := roleCmd(m.GuildID, m.UserID, r); err != nil {
				beaver.Errorf("could not modify role %s for user %s", r, m.UserID)
			}
		}
	}
}