Add agent => bot

agent/input/discord/discord.go

@@ -0,0 +1,153 @@
+package discord
+
+import (
+	"fmt"
+	"sync"
+
+	"errors"
+	"strings"
+
+	"github.com/bwmarrin/discordgo"
+	"github.com/micro/cli"
+	"github.com/micro/go-micro/agent/input"
+)
+
+func init() {
+	input.Inputs["discord"] = newInput()
+}
+
+func newInput() *discordInput {
+	return &discordInput{}
+}
+
+type discordInput struct {
+	token     string
+	whitelist []string
+	prefix    string
+	prefixfn  func(string) (string, bool)
+	botID     string
+
+	session *discordgo.Session
+
+	sync.Mutex
+	running bool
+	exit    chan struct{}
+}
+
+func (d *discordInput) Flags() []cli.Flag {
+	return []cli.Flag{
+		cli.StringFlag{
+			Name:   "discord_token",
+			EnvVar: "MICRO_DISCORD_TOKEN",
+			Usage:  "Discord token (prefix with Bot if it's for bot account)",
+		},
+		cli.StringFlag{
+			Name:   "discord_whitelist",
+			EnvVar: "MICRO_DISCORD_WHITELIST",
+			Usage:  "Discord Whitelist (seperated by ,)",
+		},
+		cli.StringFlag{
+			Name:   "discord_prefix",
+			Usage:  "Discord Prefix",
+			EnvVar: "MICRO_DISCORD_PREFIX",
+			Value:  "Micro ",
+		},
+	}
+}
+
+func (d *discordInput) Init(ctx *cli.Context) error {
+	token := ctx.String("discord_token")
+	whitelist := ctx.String("discord_whitelist")
+	prefix := ctx.String("discord_prefix")
+
+	if len(token) == 0 {
+		return errors.New("require token")
+	}
+
+	d.token = token
+	d.prefix = prefix
+
+	if len(whitelist) > 0 {
+		d.whitelist = strings.Split(whitelist, ",")
+	}
+
+	return nil
+}
+
+func (d *discordInput) Start() error {
+	if len(d.token) == 0 {
+		return errors.New("missing discord configuration")
+	}
+
+	d.Lock()
+	defer d.Unlock()
+
+	if d.running {
+		return nil
+	}
+
+	var err error
+	d.session, err = discordgo.New(d.token)
+	if err != nil {
+		return err
+	}
+
+	u, err := d.session.User("@me")
+	if err != nil {
+		return err
+	}
+
+	d.botID = u.ID
+	d.prefixfn = CheckPrefixFactory(fmt.Sprintf("<@%s> ", d.botID), fmt.Sprintf("<@!%s> ", d.botID), d.prefix)
+
+	d.exit = make(chan struct{})
+	d.running = true
+
+	return nil
+}
+
+func (d *discordInput) Stream() (input.Conn, error) {
+	d.Lock()
+	defer d.Unlock()
+	if !d.running {
+		return nil, errors.New("not running")
+	}
+
+	//Fire-n-forget close just in case...
+	d.session.Close()
+
+	conn := newConn(d)
+	if err := d.session.Open(); err != nil {
+		return nil, err
+	}
+	return conn, nil
+}
+
+func (d *discordInput) Stop() error {
+	d.Lock()
+	defer d.Unlock()
+
+	if !d.running {
+		return nil
+	}
+
+	close(d.exit)
+	d.running = false
+	return nil
+}
+
+func (d *discordInput) String() string {
+	return "discord"
+}
+
+// CheckPrefixFactory Creates a prefix checking function and stuff.
+func CheckPrefixFactory(prefixes ...string) func(string) (string, bool) {
+	return func(content string) (string, bool) {
+		for _, prefix := range prefixes {
+			if strings.HasPrefix(content, prefix) {
+				return strings.TrimPrefix(content, prefix), true
+			}
+		}
+		return "", false
+	}
+}