discord-retokenizer/storage/tokens.go

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)
}