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).User, } } 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) } func TokensUnderBot(botID string) []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 token TokenBot err := client.Get(ctx, "token:" + t.ID).Scan(&token) if err != nil { panic(err) } if token.BotID == botID { simpleTokens = append(simpleTokens, token.WithUser().Simplify()) } } return simpleTokens }