154 lines
		
	
	
		
			2.7 KiB
		
	
	
	
		
			Go
		
	
	
	
	
	
			
		
		
	
	
			154 lines
		
	
	
		
			2.7 KiB
		
	
	
	
		
			Go
		
	
	
	
	
	
| 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
 | |
| 	}
 | |
| }
 |