github.com/diamondburned/arikawa@v1.3.14/session/session.go (about)

     1  // Package session abstracts around the REST API and the Gateway, managing both
     2  // at once. It offers a handler interface similar to that in discordgo for
     3  // Gateway events.
     4  package session
     5  
     6  import (
     7  	"sync"
     8  
     9  	"github.com/pkg/errors"
    10  
    11  	"github.com/diamondburned/arikawa/api"
    12  	"github.com/diamondburned/arikawa/gateway"
    13  	"github.com/diamondburned/arikawa/utils/handler"
    14  )
    15  
    16  var ErrMFA = errors.New("account has 2FA enabled")
    17  
    18  // Closed is an event that's sent to Session's command handler. This works by
    19  // using (*Gateway).AfterClose. If the user sets this callback, no Closed events
    20  // would be sent.
    21  //
    22  // Usage
    23  //
    24  //    ses.AddHandler(func(*session.Closed) {})
    25  //
    26  type Closed struct {
    27  	Error error
    28  }
    29  
    30  // Session manages both the API and Gateway. As such, Session inherits all of
    31  // API's methods, as well has the Handler used for Gateway.
    32  type Session struct {
    33  	*api.Client
    34  	Gateway *gateway.Gateway
    35  
    36  	// Command handler with inherited methods.
    37  	*handler.Handler
    38  
    39  	// MFA only fields
    40  	MFA    bool
    41  	Ticket string
    42  
    43  	hstop chan struct{}
    44  	wstop sync.Once
    45  }
    46  
    47  func NewWithIntents(token string, intents ...gateway.Intents) (*Session, error) {
    48  	g, err := gateway.NewGatewayWithIntents(token, intents...)
    49  	if err != nil {
    50  		return nil, errors.Wrap(err, "failed to connect to Gateway")
    51  	}
    52  
    53  	return NewWithGateway(g), nil
    54  }
    55  
    56  // New creates a new session from a given token. Most bots should be using
    57  // NewWithIntents instead.
    58  func New(token string) (*Session, error) {
    59  	// Create a gateway
    60  	g, err := gateway.NewGateway(token)
    61  	if err != nil {
    62  		return nil, errors.Wrap(err, "failed to connect to Gateway")
    63  	}
    64  
    65  	return NewWithGateway(g), nil
    66  }
    67  
    68  // Login tries to log in as a normal user account; MFA is optional.
    69  func Login(email, password, mfa string) (*Session, error) {
    70  	// Make a scratch HTTP client without a token
    71  	client := api.NewClient("")
    72  
    73  	// Try to login without TOTP
    74  	l, err := client.Login(email, password)
    75  	if err != nil {
    76  		return nil, errors.Wrap(err, "failed to login")
    77  	}
    78  
    79  	if l.Token != "" && !l.MFA {
    80  		// We got the token, return with a new Session.
    81  		return New(l.Token)
    82  	}
    83  
    84  	// Discord requests MFA, so we need the MFA token.
    85  	if mfa == "" {
    86  		return nil, ErrMFA
    87  	}
    88  
    89  	// Retry logging in with a 2FA token
    90  	l, err = client.TOTP(mfa, l.Ticket)
    91  	if err != nil {
    92  		return nil, errors.Wrap(err, "failed to login with 2FA")
    93  	}
    94  
    95  	return New(l.Token)
    96  }
    97  
    98  func NewWithGateway(gw *gateway.Gateway) *Session {
    99  	return &Session{
   100  		Gateway: gw,
   101  		// Nab off gateway's token
   102  		Client:  api.NewClient(gw.Identifier.Token),
   103  		Handler: handler.New(),
   104  	}
   105  }
   106  
   107  func (s *Session) Open() error {
   108  	// Start the handler beforehand so no events are missed.
   109  	s.hstop = make(chan struct{})
   110  	s.wstop = sync.Once{}
   111  	go s.startHandler()
   112  
   113  	// Set the AfterClose's handler.
   114  	s.Gateway.AfterClose = func(err error) {
   115  		s.Handler.Call(&Closed{
   116  			Error: err,
   117  		})
   118  	}
   119  
   120  	if err := s.Gateway.Open(); err != nil {
   121  		return errors.Wrap(err, "failed to start gateway")
   122  	}
   123  
   124  	return nil
   125  }
   126  
   127  func (s *Session) startHandler() {
   128  	for {
   129  		select {
   130  		case <-s.hstop:
   131  			return
   132  		case ev := <-s.Gateway.Events:
   133  			s.Call(ev)
   134  		}
   135  	}
   136  }
   137  
   138  func (s *Session) Close() error {
   139  	// Stop the event handler
   140  	s.wstop.Do(func() { s.hstop <- struct{}{} })
   141  	// Close the websocket
   142  	return s.Gateway.Close()
   143  }