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 _, 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)
                return
            }
        case CommandTypeFullMatch:
            if content == command.Trigger {
                executeCommand(s, m, command)
                return
            }
        case CommandTypeRegex:
            match, _ := regexp.MatchString(command.Trigger, content)
            if match {
                executeCommand(s, m, command)
                return
            }
        case CommandTypeContains:
            if strings.Contains(content, command.Trigger) {
                executeCommand(s, m, command)
                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) {
    if isAdmin(message.Author) || // no restrictions for admins
            (!command.AdminOnly && (isDM(session, message) || !command.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) {
            command.UsersOnCooldown.Add(message.Author.ID)
            go removeCooldown(command, 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(command *Command, uid string) {
    time.Sleep(time.Duration(command.Cooldown) * time.Second)
    if command.UsersOnCooldown.Contains(uid) {
        command.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
}