mautrix-discord/remoteauth/client.go
2022-09-13 14:45:50 +03:00

125 lines
2.1 KiB
Go

package remoteauth
import (
"context"
"crypto/rand"
"crypto/rsa"
"crypto/sha256"
"encoding/base64"
"encoding/json"
"net/http"
"sync"
"github.com/gorilla/websocket"
"github.com/bwmarrin/discordgo"
)
type Client struct {
sync.Mutex
URL string
conn *websocket.Conn
qrChan chan string
doneChan chan struct{}
user User
err error
heartbeats int
closed bool
privateKey *rsa.PrivateKey
}
// New creates a new Discord remote auth client. qrChan is a channel that will
// receive the qrcode once it is available.
func New() (*Client, error) {
privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
if err != nil {
return nil, err
}
return &Client{
URL: "wss://remote-auth-gateway.discord.gg/?v=2",
privateKey: privateKey,
}, nil
}
// Dial will start the QRCode login process. ctx may be used to abandon the
// process.
func (c *Client) Dial(ctx context.Context, qrChan chan string, doneChan chan struct{}) error {
c.Lock()
defer c.Unlock()
header := http.Header{}
for key, value := range discordgo.DroidWSHeaders {
header.Set(key, value)
}
c.qrChan = qrChan
c.doneChan = doneChan
conn, _, err := websocket.DefaultDialer.DialContext(ctx, c.URL, header)
if err != nil {
return err
}
c.conn = conn
go c.processMessages()
return nil
}
func (c *Client) Result() (User, error) {
c.Lock()
defer c.Unlock()
return c.user, c.err
}
func (c *Client) close() error {
c.Lock()
defer c.Unlock()
if c.closed {
return nil
}
c.conn.WriteMessage(
websocket.CloseMessage,
websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""),
)
c.closed = true
defer close(c.doneChan)
return c.conn.Close()
}
func (c *Client) write(p clientPacket) error {
c.Lock()
defer c.Unlock()
payload, err := json.Marshal(p)
if err != nil {
return err
}
return c.conn.WriteMessage(websocket.TextMessage, payload)
}
func (c *Client) decrypt(payload string) ([]byte, error) {
// Decode the base64 string.
raw, err := base64.StdEncoding.DecodeString(payload)
if err != nil {
return []byte{}, err
}
// Decrypt the data.
return rsa.DecryptOAEP(sha256.New(), nil, c.privateKey, raw, nil)
}