discord-flauschbot/command.go

package main

import (
    "fmt"
    "github.com/bwmarrin/discordgo"
    "github.com/deckarep/golang-set"
    "strings"
    "time"
    "log"
    "regexp"
)

type CommandType int

// These are used to specify Command.CommandType when registering new commands
const (
    CommandTypePrefix      CommandType = 0
    CommandTypeFullMatch   CommandType = 1
    CommandTypeRegex       CommandType = 2
    CommandTypeContains    CommandType = 3
)

 /*
    This struct represents a command object.
    The options should be self-explanatory, but they are also explained in the readme.
    A struct can be initialized by passing any number of its attributes as parameters.
    Everything not set will be set to the go-usual defaults ("" for string, 0 for int, false for bool, nil for the rest)
    Any command that has a Trigger is valid (but useless if nothing else is specified)
 */
type Command struct {
    Trigger string                          // must be specified
    Output string                           // no output if unspecified
    OutputEmbed *discordgo.MessageEmbed     // no embed output if unspecified
    Type CommandType                        // defaults to Prefix
    Cooldown int                            // defaults to 0 (no cooldown)
    OutputIsReply bool
    RequiresMention bool
    DeleteInput bool
    DMOnly bool
    AdminOnly bool
    IgnoreCase bool
    // for custom commands that go beyond prints and deletions
    Function func(*discordgo.Session, *discordgo.MessageCreate)

    UsersOnCooldown mapset.Set            // don’t set this manually (it’s overwritten anyway)
}


// Performs basic input validation on a given command and adds it to the global command array
func registerCommand(command Command) {
    if command.Trigger == "" {
        fmt.Println("Cannot register a command with no trigger. Skipping.")
        return
    }
    if command.IgnoreCase {
        command.Trigger = strings.ToLower(command.Trigger)
    }
    command.UsersOnCooldown = mapset.NewSet()
    commands = append(commands, command)
}

 /*
    Any message that the bot can read is evaluated here.
    The message is matched against each of the command triggers depending on the respective match type.
    If one of the commands matches, execute that command and return.
    Only one command can be executed per message. Earlier defined commands take precedence.
    This is a deliberate choice (for now).
 */
func evaluateMessage(s *discordgo.Session, m *discordgo.MessageCreate) {
    if m.Author.ID == s.State.User.ID {
        log.Printf("<Self> %s", m.Content)
        return
    }
    for i, command := range commands {
        content := m.Content
        if command.IgnoreCase {
            content = strings.ToLower(content)
        }
        if command.RequiresMention {
            command.Trigger = fmt.Sprintf(command.Trigger, s.State.User.ID)
        }
        switch command.Type {
        case CommandTypePrefix:
            if strings.HasPrefix(content, command.Trigger) {
                executeCommand(s, m, command, i)
                return
            }
        case CommandTypeFullMatch:
            if content == command.Trigger {
                executeCommand(s, m, command, i)
                return
            }
        case CommandTypeRegex:
            match, _ := regexp.MatchString(command.Trigger, content)
            if match {
                executeCommand(s, m, command, i)
                return
            }
        case CommandTypeContains:
            if strings.Contains(content, command.Trigger) {
                executeCommand(s, m, command, i)
                return
            }
        }
    }
}

 /*
    Executes the given command on the given message and session.
    Sets command cooldowns if necessary and also clears them again.
 */
func executeCommand(session *discordgo.Session, message *discordgo.MessageCreate, command Command, commandIndex int) {
    if isAdmin(message.Author) || // no restrictions for admins
            (!command.AdminOnly && (isDM(session, message) || !commands[commandIndex].UsersOnCooldown.Contains(message.Author.ID)) &&
            (!command.DMOnly || isDM(session, message))) {

        log.Printf("Executed command %s triggered by user %s", command.Trigger, userToString(message.Author))
        if command.Cooldown > 0 && !isDM(session, message) && !isAdmin(message.Author) {
            commands[commandIndex].UsersOnCooldown.Add(message.Author.ID)
            go removeCooldown(commandIndex, message.Author.ID)
        }
        if command.Function == nil {
            // simple reply
            if command.OutputEmbed == nil {
                messageContent := generateReply(message, command)
                session.ChannelMessageSend(message.ChannelID, messageContent)
            } else {
                session.ChannelMessageSendEmbed(message.ChannelID, command.OutputEmbed)
            }
            if command.DeleteInput {
                session.ChannelMessageDelete(message.ChannelID, message.ID)
            }
        } else {
            // execute custom function
            command.Function(session, message)
        }
    } else {
        log.Printf("Denied command %s to user %s.", command.Trigger, userToString(message.Author))
    }
}

func removeCooldown(commandIndex int, uid string) {
    time.Sleep(time.Duration(commands[commandIndex].Cooldown) * time.Second)
    if commands[commandIndex].UsersOnCooldown.Contains(uid) {
        commands[commandIndex].UsersOnCooldown.Remove(uid)
    }
}

func generateReply(message *discordgo.MessageCreate, command Command) string {
    output := command.Output
    if command.OutputIsReply {
        output = fmt.Sprintf(output, message.Author.ID)
    }
    return output
}

 /*
    Any message passed to this method will be redirected to config.ModChannel.
    This is useful for anonymous complaints or similar messages.
 */
func redirectComplaint(s *discordgo.Session, m *discordgo.MessageCreate) {
    embed := &discordgo.MessageEmbed {
            Author:      &discordgo.MessageEmbedAuthor{},
            Color:       0xbb0000,
            Description: m.Content,
    }
    s.ChannelMessageSendEmbed(config.ModChannel, embed)
}

func echoMessage(s *discordgo.Session, m *discordgo.MessageCreate) {
    s.ChannelMessageSend(m.ChannelID, m.Content)
}

func giveAgeRole(s *discordgo.Session, m *discordgo.MessageCreate) {
    Member, _ := s.GuildMember(config.ServerID, m.Author.ID)
    dm, _ := s.UserChannelCreate(Member.User.ID)
    for command, role := range config.RoleCommands {
        if m.Content == command {
            // Found the command that was triggered
            // This is a restriction imposed by my own wrapper, 
            // but working around it is not actually necessary for performance and makes the code uglier in other places.
            for _, newRole := range config.RoleCommands {
                for _, curRole := range Member.Roles {
                    // If the user already has one of the available roles, tell him and exit
                    if newRole == curRole {
                        if curRole == role {
                            // User is trying to get the role they already have
                            s.ChannelMessageSend(dm.ID, "Baka, die Rolle hast du doch schon.")
                            log.Printf("Denied Role %s to %s. User already has %s", roleName(s.State, curRole), userToString(m.Author), roleName(s.State, curRole))
                        } else {
                            s.ChannelMessageSend(dm.ID, "Baka, du kannst nur eine der Rollen haben.")
                            log.Printf("Denied Role %s to %s. User already has %s", roleName(s.State, curRole), userToString(m.Author), roleName(s.State, curRole))
                        }
                        return
                    }
                }
            }
            log.Printf("Giving Role %s to %s", roleName(s.State, role), userToString(m.Author))
            s.ChannelMessageSend(dm.ID, "Haaai, Ryoukai desu~")
            s.GuildMemberRoleAdd(config.ServerID, m.Author.ID, role)
        }
    }
}