diff --git a/storage/bots.go b/storage/bots.go index fc274dc..c942eda 100644 --- a/storage/bots.go +++ b/storage/bots.go @@ -23,26 +23,42 @@ func BotExists(id string) bool { return true } -func GetBots() []discord.SimpleUser { +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([]discord.SimpleUser, 0) + bots := make([]BotDataSimple, 0) for _, id := range botIDs { - bots = append(bots, GetBot(id).Simplify()) + bots = append(bots, BotDataSimple{ + User: GetBot(id).User.Simplify(), + TokenCount: len(TokensUnderBot(id)), + }) } return bots } -func GetBot(id string) discord.User { +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 bot + return BotData{ + User: bot, + TokenCount: len(TokensUnderBot(id)), + } } func ExistsBot(id string) bool { diff --git a/storage/tokens.go b/storage/tokens.go index 94d0ea6..717ebe1 100644 --- a/storage/tokens.go +++ b/storage/tokens.go @@ -164,7 +164,7 @@ func (t TokenBot) WithUser() TokenUserBot { return TokenUserBot{ ID: t.ID, End: t.End, - User: GetBot(t.BotID), + User: GetBot(t.BotID).User, } } @@ -226,4 +226,25 @@ func BotTokenFromToken(token string) (bool, string) { 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 } \ No newline at end of file