github.com/annwntech/go-micro/v2@v2.9.5/agent/input/discord/discord.go (about)

     1  package discord
     2  
     3  import (
     4  	"fmt"
     5  	"sync"
     6  
     7  	"errors"
     8  	"strings"
     9  
    10  	"github.com/bwmarrin/discordgo"
    11  	"github.com/micro/cli/v2"
    12  	"github.com/annwntech/go-micro/v2/agent/input"
    13  )
    14  
    15  func init() {
    16  	input.Inputs["discord"] = newInput()
    17  }
    18  
    19  func newInput() *discordInput {
    20  	return &discordInput{}
    21  }
    22  
    23  type discordInput struct {
    24  	token     string
    25  	whitelist []string
    26  	prefix    string
    27  	prefixfn  func(string) (string, bool)
    28  	botID     string
    29  
    30  	session *discordgo.Session
    31  
    32  	sync.Mutex
    33  	running bool
    34  	exit    chan struct{}
    35  }
    36  
    37  func (d *discordInput) Flags() []cli.Flag {
    38  	return []cli.Flag{
    39  		&cli.StringFlag{
    40  			Name:    "discord_token",
    41  			EnvVars: []string{"MICRO_DISCORD_TOKEN"},
    42  			Usage:   "Discord token (prefix with Bot if it's for bot account)",
    43  		},
    44  		&cli.StringFlag{
    45  			Name:    "discord_whitelist",
    46  			EnvVars: []string{"MICRO_DISCORD_WHITELIST"},
    47  			Usage:   "Discord Whitelist (seperated by ,)",
    48  		},
    49  		&cli.StringFlag{
    50  			Name:    "discord_prefix",
    51  			Usage:   "Discord Prefix",
    52  			EnvVars: []string{"MICRO_DISCORD_PREFIX"},
    53  			Value:   "Micro ",
    54  		},
    55  	}
    56  }
    57  
    58  func (d *discordInput) Init(ctx *cli.Context) error {
    59  	token := ctx.String("discord_token")
    60  	whitelist := ctx.String("discord_whitelist")
    61  	prefix := ctx.String("discord_prefix")
    62  
    63  	if len(token) == 0 {
    64  		return errors.New("require token")
    65  	}
    66  
    67  	d.token = token
    68  	d.prefix = prefix
    69  
    70  	if len(whitelist) > 0 {
    71  		d.whitelist = strings.Split(whitelist, ",")
    72  	}
    73  
    74  	return nil
    75  }
    76  
    77  func (d *discordInput) Start() error {
    78  	if len(d.token) == 0 {
    79  		return errors.New("missing discord configuration")
    80  	}
    81  
    82  	d.Lock()
    83  	defer d.Unlock()
    84  
    85  	if d.running {
    86  		return nil
    87  	}
    88  
    89  	var err error
    90  	d.session, err = discordgo.New("Bot " + d.token)
    91  	if err != nil {
    92  		return err
    93  	}
    94  
    95  	u, err := d.session.User("@me")
    96  	if err != nil {
    97  		return err
    98  	}
    99  
   100  	d.botID = u.ID
   101  	d.prefixfn = CheckPrefixFactory(fmt.Sprintf("<@%s> ", d.botID), fmt.Sprintf("<@!%s> ", d.botID), d.prefix)
   102  
   103  	d.exit = make(chan struct{})
   104  	d.running = true
   105  
   106  	return nil
   107  }
   108  
   109  func (d *discordInput) Stream() (input.Conn, error) {
   110  	d.Lock()
   111  	defer d.Unlock()
   112  	if !d.running {
   113  		return nil, errors.New("not running")
   114  	}
   115  
   116  	//Fire-n-forget close just in case...
   117  	d.session.Close()
   118  
   119  	conn := newConn(d)
   120  	if err := d.session.Open(); err != nil {
   121  		return nil, err
   122  	}
   123  	return conn, nil
   124  }
   125  
   126  func (d *discordInput) Stop() error {
   127  	d.Lock()
   128  	defer d.Unlock()
   129  
   130  	if !d.running {
   131  		return nil
   132  	}
   133  
   134  	close(d.exit)
   135  	d.running = false
   136  	return nil
   137  }
   138  
   139  func (d *discordInput) String() string {
   140  	return "discord"
   141  }
   142  
   143  // CheckPrefixFactory Creates a prefix checking function and stuff.
   144  func CheckPrefixFactory(prefixes ...string) func(string) (string, bool) {
   145  	return func(content string) (string, bool) {
   146  		for _, prefix := range prefixes {
   147  			if strings.HasPrefix(content, prefix) {
   148  				return strings.TrimPrefix(content, prefix), true
   149  			}
   150  		}
   151  		return "", false
   152  	}
   153  }