141 lines
2.6 KiB
Go
141 lines
2.6 KiB
Go
package storage
|
|
|
|
import (
|
|
"fmt"
|
|
|
|
"git.zomo.dev/zomo/discord-retokenizer/discord"
|
|
"github.com/go-redis/redis/v9"
|
|
)
|
|
|
|
func BotExists(id string) bool {
|
|
exist, err := client.SMIsMember(ctx, "bots", id).Result()
|
|
if err != nil {
|
|
if err == redis.Nil {
|
|
return false
|
|
}
|
|
panic(err)
|
|
}
|
|
for _, b := range exist {
|
|
if !b {
|
|
return false
|
|
}
|
|
}
|
|
return true
|
|
}
|
|
|
|
type BotData struct {
|
|
User discord.User `json:"user"`
|
|
TokenCount int `json:"token_count"`
|
|
}
|
|
|
|
type BotDataSimple struct {
|
|
User discord.SimpleUser `json:"user"`
|
|
TokenCount int `json:"token_count"`
|
|
}
|
|
|
|
func GetBots() []BotDataSimple {
|
|
botIDs, err := client.SMembers(ctx, "bots").Result()
|
|
if err != nil && err != redis.Nil {
|
|
panic(err)
|
|
}
|
|
bots := make([]BotDataSimple, 0)
|
|
for _, id := range botIDs {
|
|
bots = append(bots, BotDataSimple{
|
|
User: GetBot(id).User.Simplify(),
|
|
TokenCount: len(TokensUnderBot(id)),
|
|
})
|
|
}
|
|
return bots
|
|
}
|
|
|
|
func GetBot(id string) BotData {
|
|
var bot discord.User
|
|
key := fmt.Sprintf("bot:%s:data", id)
|
|
err := client.Get(ctx, key).Scan(&bot)
|
|
if err != nil && err != redis.Nil {
|
|
panic(err)
|
|
}
|
|
return BotData{
|
|
User: bot,
|
|
TokenCount: len(TokensUnderBot(id)),
|
|
}
|
|
}
|
|
|
|
func ExistsBot(id string) bool {
|
|
exist, err := client.SIsMember(ctx, "bots", id).Result()
|
|
if err != nil {
|
|
if err == redis.Nil {
|
|
return false
|
|
}
|
|
panic(err)
|
|
}
|
|
return exist
|
|
}
|
|
|
|
func RemoveBot(id string) {
|
|
err := client.SRem(ctx, "bots", id).Err()
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
keyPrefix := fmt.Sprintf("bot:%s:", id)
|
|
err = client.Del(ctx, keyPrefix + "token").Err()
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
err = client.Del(ctx, keyPrefix + "data").Err()
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
}
|
|
|
|
func getBotToken(id string) string {
|
|
if !BotExists(id) {
|
|
return ""
|
|
}
|
|
key := fmt.Sprintf("bot:%s:token", id)
|
|
token, err := client.Get(ctx, key).Result()
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
return token
|
|
}
|
|
|
|
func RefreshBot(id string) discord.User {
|
|
token := getBotToken(id)
|
|
bot := discord.GetDiscordUser(token)
|
|
|
|
key := fmt.Sprintf("bot:%s:data", id)
|
|
err := client.Set(ctx, key, bot, 0).Err()
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
return bot
|
|
}
|
|
|
|
func AddBot(token string) bool {
|
|
user := discord.GetDiscordUser(token)
|
|
|
|
if user.ID == "" {
|
|
return false
|
|
}
|
|
|
|
err := client.SAdd(ctx, "bots", user.ID).Err()
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
userID := user.ID
|
|
keyPrefix := fmt.Sprintf("bot:%s:", userID)
|
|
|
|
err = client.Set(ctx, keyPrefix + "token", token, 0).Err()
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
err = client.Set(ctx, keyPrefix + "data", user, 0).Err()
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
|
|
return true
|
|
} |