more

storage/tokens.go

@@ -0,0 +1,229 @@
+package storage
+
+import (
+	"encoding/json"
+
+	"git.zomo.dev/zomo/discord-retokenizer/discord"
+	"git.zomo.dev/zomo/discord-retokenizer/util"
+	"golang.org/x/crypto/bcrypt"
+)
+
+type TokenTiny struct {
+	ID        string `json:"id"`
+	TokenHash string `json:"token"`
+}
+
+type TokenBot struct {
+	ID        string `json:"id"`
+	BotID     string `json:"bot_id"`
+	End       string `json:"end"`
+	TokenHash string `json:"token"`
+}
+
+type SimpleTokenBot struct {
+	ID        string `json:"id"`
+	BotID     string `json:"bot_id"`
+	End       string `json:"end"`
+}
+
+func (t TokenBot) Simplify() SimpleTokenBot {
+	return SimpleTokenBot{
+		ID:    t.ID,
+		BotID: t.BotID,
+		End:   t.End,
+	}
+}
+
+func (t TokenTiny) MarshalBinary() ([]byte, error) {
+    return json.Marshal(t)
+}
+
+func (t *TokenTiny) UnmarshalBinary(data []byte) error {
+	return json.Unmarshal(data, t)
+}
+
+func (t TokenBot) MarshalBinary() ([]byte, error) {
+    return json.Marshal(t)
+}
+
+func (t *TokenBot) UnmarshalBinary(data []byte) error {
+	return json.Unmarshal(data, t)
+}
+
+func (t SimpleTokenBot) MarshalBinary() ([]byte, error) {
+    return json.Marshal(t)
+}
+
+func (t *SimpleTokenBot) UnmarshalBinary(data []byte) error {
+	return json.Unmarshal(data, t)
+}
+
+func GenerateToken(botID string) string {
+	if !ExistsBot(botID) {
+		return ""
+	}
+
+	// Generate a token
+	token := util.GenerateToken()
+	tokenID := util.GenerateID()
+	tokenHash, err := bcrypt.GenerateFromPassword([]byte(token), bcrypt.DefaultCost)
+	if err != nil {
+		panic(err)
+	}
+
+	tinyToken := TokenTiny{
+		ID:        tokenID,
+		TokenHash: string(tokenHash),
+	}
+
+	err = client.SAdd(ctx, "tokens", tinyToken).Err()
+	if err != nil {
+		panic(err)
+	}
+
+	botToken := TokenBot{
+		ID:        tokenID,
+		BotID:     botID,
+		End:       util.GetEnd(token),
+		TokenHash: string(tokenHash),
+	}
+	
+	err = client.Set(ctx, "token:" + tokenID, botToken, 0).Err()
+	if err != nil {
+		panic(err)
+	}
+
+	return token
+}
+
+func getTokenBot(token string) (bool, TokenBot) {
+	var tokens []TokenTiny
+	err := client.SMembers(ctx, "tokens").ScanSlice(&tokens)
+	if err != nil {
+		panic(err)
+	}
+	for _, t := range tokens {
+		err = bcrypt.CompareHashAndPassword([]byte(t.TokenHash), []byte(token))
+		if err == nil {
+			var botToken TokenBot
+			err := client.Get(ctx, "token:" + t.ID).Scan(&botToken)
+			if err != nil {
+				panic(err)
+			}
+			return true, botToken
+		}
+	}
+	return false, TokenBot{}
+}
+
+func DeleteToken(token string) {
+	foundToken, botToken := getTokenBot(token)
+	if !foundToken {
+		return
+	}
+	err := client.SRem(ctx, "tokens", botToken.TokenHash).Err()
+	if err != nil {
+		panic(err)
+	}
+	err = client.Del(ctx, "token:" + botToken.ID).Err()
+	if err != nil {
+		panic(err)
+	}
+}
+
+func UpdateToken(token string, botID string) bool {
+	if !ExistsBot(botID) {
+		return false
+	}
+
+	foundToken, botToken := getTokenBot(token)
+	if !foundToken {
+		return false
+	}
+
+	err := client.Set(ctx, "token:" + botToken.ID, botToken, 0).Err()
+	if err != nil {
+		panic(err)
+	}
+	return true
+}
+
+type TokenUserBot struct {
+	ID   string       `json:"id"`
+	End  string       `json:"token_end"`
+	User discord.User `json:"user"`
+}
+
+type SimpleTokenUserBot struct {
+	ID   string             `json:"id"`
+	End  string             `json:"token_end"`
+	User discord.SimpleUser `json:"user"`
+}
+
+func (t TokenBot) WithUser() TokenUserBot {
+	return TokenUserBot{
+		ID:   t.ID,
+		End:  t.End,
+		User: GetBot(t.BotID),
+	}
+}
+
+func (t TokenUserBot) Simplify() SimpleTokenUserBot {
+	return SimpleTokenUserBot{
+		ID:   t.ID,
+		End:  t.End,
+		User: t.User.Simplify(),
+	}
+}
+
+func (t TokenUserBot) MarshalBinary() ([]byte, error) {
+    return json.Marshal(t)
+}
+
+func (t *TokenUserBot) UnmarshalBinary(data []byte) error {
+	return json.Unmarshal(data, t)
+}
+
+func (t SimpleTokenUserBot) MarshalBinary() ([]byte, error) {
+    return json.Marshal(t)
+}
+
+func (t *SimpleTokenUserBot) UnmarshalBinary(data []byte) error {
+	return json.Unmarshal(data, t)
+}
+
+func GetTokens() []SimpleTokenUserBot {
+	var tokens []TokenTiny
+	err := client.SMembers(ctx, "tokens").ScanSlice(&tokens)
+	if err != nil {
+		panic(err)
+	}
+
+	simpleTokens := make([]SimpleTokenUserBot, 0)
+	for _, t := range tokens {
+		var tokenBot TokenBot
+		err := client.Get(ctx, "token:" + t.ID).Scan(&tokenBot)
+		if err != nil {
+			panic(err)
+		}
+		simpleTokens = append(simpleTokens, tokenBot.WithUser().Simplify())
+	}
+	return simpleTokens
+}
+
+func GetToken(id string) TokenUserBot {
+	var tokenBot TokenBot
+	err := client.Get(ctx, "token:" + id).Scan(&tokenBot)
+	if err != nil {
+		panic(err)
+	}
+	return tokenBot.WithUser()
+}
+
+func BotTokenFromToken(token string) (bool, string) {
+	foundBot, tokenBot := getTokenBot(token)
+	if !foundBot {
+		return false, ""
+	}
+	return true, getBotToken(tokenBot.BotID)
+}