Pārlūkot izejas kodu

Added bot handling and routing to the library

Thisnthat 5 gadi atpakaļ
vecāks
revīzija
afa874101b
4 mainītis faili ar 326 papildinājumiem un 0 dzēšanām
  1. 5 0
      bot.go
  2. 75 0
      bot/event/mux/help.go
  3. 95 0
      bot/event/mux/message.go
  4. 151 0
      bot/event/mux/router.go

+ 5 - 0
bot.go

@@ -7,6 +7,7 @@ import (
 	"sync"
 	"syscall"
 
+	"git.mgmcomp.net/thisnthat/discord/bot/event/mux"
 	"github.com/Sirupsen/logrus"
 	"github.com/thisnthat-dev/discordgo"
 )
@@ -29,6 +30,10 @@ func (bot *Bot) AddHandler(handler interface{}) error {
 	return nil
 }
 
+func (bot *Bot) SetEventRouter(router *mux.Router) {
+	bot.AddHandler(router.OnMessageCreate)
+}
+
 func (bot *Bot) Start(wg *sync.WaitGroup) error {
 	defer wg.Done()
 

+ 75 - 0
bot/event/mux/help.go

@@ -0,0 +1,75 @@
+package mux
+
+import (
+	"fmt"
+	"sort"
+	"strconv"
+
+	"github.com/thisnthat-dev/discordgo"
+)
+
+func (r *Router) helpCommandHandler(s *discordgo.Session, m *discordgo.Message, ctx *Context) {
+	displayPrefix := ""
+
+	switch ctx.Method {
+	case DirectMethod:
+		displayPrefix = ""
+	case PrefixMethod:
+		displayPrefix = r.prefix
+	case MentionMethod:
+		displayPrefix = fmt.Sprintf("@%s ", s.State.User.Username)
+	}
+
+	// Sort commands
+	maxDisplayLen := 0
+	keys := make([]string, 0, len(r.routes))
+	cmdmap := make(map[string]*route)
+
+	for _, v := range r.routes {
+		// Only display commands with a description
+		if v.Description == "" {
+			continue
+		}
+
+		// Calculate the max length of command+args string
+		l := len(v.Name) // TODO: Add the +args part :)
+		if l > maxDisplayLen {
+			maxDisplayLen = l
+		}
+
+		cmdmap[v.Name] = v
+
+		// help and about are added separately below.
+		if v.Name == "help" || v.Name == "about" {
+			continue
+		}
+
+		keys = append(keys, v.Name)
+	}
+
+	sort.Strings(keys)
+
+	// TODO: Learn more link needs to be configurable
+	resp := "```autoit\n"
+
+	v, ok := cmdmap["help"]
+	if ok {
+		keys = append([]string{v.Name}, keys...)
+	}
+
+	v, ok = cmdmap["about"]
+	if ok {
+		keys = append([]string{v.Name}, keys...)
+	}
+
+	// Add sorted result to help msg
+	for _, k := range keys {
+		v := cmdmap[k]
+		resp += fmt.Sprintf("%s%-"+strconv.Itoa(maxDisplayLen)+"s # %s\n", displayPrefix, v.Name, v.Description)
+	}
+
+	resp += "```\n"
+	s.ChannelMessageSend(m.ChannelID, resp)
+
+	return
+}

+ 95 - 0
bot/event/mux/message.go

@@ -0,0 +1,95 @@
+package mux
+
+import (
+	"fmt"
+	"regexp"
+	"strings"
+
+	"github.com/Sirupsen/logrus"
+	"github.com/thisnthat-dev/discordgo"
+)
+
+// OnMessageCreate is a DiscordGo Event Handler function.  This must be
+// registered using the DiscordGo.Session.AddHandler function.  This function
+// will receive all Discord messages and parse them for matches to registered
+// routes.
+func (r *Router) OnMessageCreate(s *discordgo.Session, m *discordgo.MessageCreate) {
+	logrus.Info("Got message")
+	var err error
+
+	// Ignore all messages created by the Bot account itself
+	if m.Author.ID == s.State.User.ID {
+		return
+	}
+
+	// Start building the command context
+	ctx := &Context{
+		Content: strings.TrimSpace(m.Content),
+	}
+
+	var channel *discordgo.Channel
+	channel, err = s.State.Channel(m.ChannelID)
+	if err != nil {
+		channel, err = s.Channel(m.ChannelID)
+		if err != nil {
+			logrus.Warnf("Failed to get the channel for message: %s", err)
+		} else {
+			err = s.State.ChannelAdd(channel)
+			if err != nil {
+				logrus.Warnf("Failed to add the channel to the session state: %s", err)
+			}
+		}
+	}
+
+	var isDirectMessage bool // The message was received in a direct message
+	var isCommand bool       // The message is intended as a command to the bot
+
+	if channel != nil {
+		if channel.Type == discordgo.ChannelTypeDM {
+			isDirectMessage = true
+			isCommand = true
+			ctx.Method = DirectMethod
+		}
+	}
+
+	if !isDirectMessage {
+		reg := regexp.MustCompile(fmt.Sprintf("<@!?(%s)>", s.State.User.ID))
+
+		for _, mention := range m.Mentions {
+			if mention.ID == s.State.User.ID {
+				// Was the @mention the first part of the string?
+				if reg.FindStringIndex(ctx.Content)[0] == 0 {
+					isCommand = true
+					ctx.Method = MentionMethod
+				}
+
+				ctx.Content = reg.ReplaceAllString(ctx.Content, "")
+			}
+		}
+
+		// Check to see if the context string starts with the command prefix
+		if len(r.prefix) > 0 {
+			// If the content start with the command prefix, then this is intended as a command
+			if strings.HasPrefix(ctx.Content, r.prefix) {
+				isCommand = true
+				ctx.Method = PrefixMethod
+				// Strip off the command prefix
+				ctx.Content = strings.TrimPrefix(ctx.Content, r.prefix)
+			}
+		}
+	}
+
+	// If this is not a command do nothing
+	if !isCommand {
+		return
+	}
+
+	route, fields := r.findRoute(ctx.Content, MessageRoute)
+	if route != nil {
+		ctx.Fields = fields
+		route.Run(s, m.Message, ctx)
+		return
+	}
+
+	r.helpRoute.Run(s, m.Message, ctx)
+}

+ 151 - 0
bot/event/mux/router.go

@@ -0,0 +1,151 @@
+package mux
+
+import (
+	"errors"
+	"strings"
+
+	"github.com/Sirupsen/logrus"
+
+	"github.com/thisnthat-dev/discordgo"
+)
+
+// RouteOptions holds information about a specific message route handler
+type RouteOptions struct {
+	Name        string    // match name that should trigger this route handler
+	Description string    // short description of this route
+	Type        RouteType // The type of route this is
+	Callback    Handler   // route handler function to call
+}
+
+type route struct {
+	Name        string
+	Description string
+	Type        RouteType
+	Run         Handler
+}
+
+// RouteType is the type of route being created
+type RouteType int
+
+const (
+	// MessageRoute is a route for handling OnMessageCreate events
+	MessageRoute RouteType = 1001
+)
+
+// Handler is the function signature required for a message route handler.
+type Handler func(*discordgo.Session, *discordgo.Message, *Context)
+
+// Context holds a bit of extra data we pass along to route handlers
+// This way processing some of this only needs to happen once.
+type Context struct {
+	Fields          []string
+	Content         string
+	isDirectMessage bool // Indicates the message was sent via direct message
+	Method          Method
+}
+
+// Method is the method a command was received
+type Method int
+
+const (
+	// DirectMethod Discord DM to the bot
+	DirectMethod Method = 1001
+	// MentionMethod is a message started with a @mention to the bot
+	MentionMethod Method = 1002
+	// PrefixMethod is a message started with a command prefix
+	PrefixMethod Method = 1003
+)
+
+// Router is the main struct for all mux methods.
+type Router struct {
+	routes    []*route
+	helpRoute *route
+	prefix    string
+}
+
+// New returns a new Discord message route mux
+func New() *Router {
+	r := &Router{}
+
+	helpRoute := &route{}
+	helpRoute.Name = "help"
+	helpRoute.Description = "Get Help"
+	helpRoute.Run = r.helpCommandHandler
+
+	r.helpRoute = helpRoute
+
+	return r
+}
+
+// SetCommandPrefix will set the prefix used for message commands
+func (r *Router) SetCommandPrefix(cmdPrefix string) {
+	r.prefix = cmdPrefix + " "
+}
+
+// ErrRegisterRouteNameRequired is returned when registering a new route without a name
+var ErrRegisterRouteNameRequired = errors.New("All routes must have a name")
+
+// ErrRegisterInvalidRouteType is returned when registering a new route with an invalid type
+var ErrRegisterInvalidRouteType = errors.New("Invalid route type")
+
+// ErrRegisterCallbackRequired is returned when registering a new route without a callback
+var ErrRegisterCallbackRequired = errors.New("A valid callback must be provided")
+
+// Register allows you to register a route
+//func (r *Router) Register(routeType RouteType, name, desc string, callback Handler) (*Route, error) {
+func (r *Router) Register(name string, options RouteOptions) error {
+	if name == "" {
+		return ErrRegisterRouteNameRequired
+	}
+
+	if !isValidRouteType(options.Type) {
+		return ErrRegisterInvalidRouteType
+	}
+
+	if options.Callback == nil {
+		return ErrRegisterCallbackRequired
+	}
+
+	route := route{}
+	route.Type = options.Type
+	route.Name = name
+	route.Description = options.Description
+	route.Run = options.Callback
+	r.routes = append(r.routes, &route)
+
+	return nil
+}
+
+func isValidRouteType(routeType RouteType) bool {
+	switch routeType {
+	case MessageRoute:
+		return true
+	}
+
+	return false
+}
+
+func (r *Router) findRoute(searchString string, routeType RouteType) (*route, []string) {
+	logrus.Printf("Find Route: %s", searchString)
+	fields := strings.Fields(searchString)
+
+	if len(fields) == 0 {
+		logrus.Printf("No route Found")
+		return nil, nil
+	}
+
+	commandKey := fields[0]
+
+	for _, route := range r.routes {
+		if route.Type != routeType {
+			continue
+		}
+
+		if strings.ToLower(commandKey) == strings.ToLower(route.Name) {
+			logrus.Printf("Route Found: %s", route.Name)
+			return route, fields[0:]
+		}
+	}
+
+	return nil, nil
+}