This commit is contained in:
2022-12-09 23:26:38 -06:00
parent 8df6231b71
commit d8a6c8acd1
15 changed files with 643 additions and 56 deletions

View File

@@ -1,27 +1,125 @@
package storage
type SimpleDiscordUser struct {
Username string `json:"username"`
Discriminator string `json:"discriminator"`
Avatar string `json:"avatar"`
Banner string `json:"banner"`
Accent_color int `json:"accent_color"`
Verified bool `json:"verified"`
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 {
Token string
ID string
Loaded SimpleDiscordUser
func GetBots() []discord.SimpleUser {
botIDs, err := client.SMembers(ctx, "bots").Result()
if err != nil && err != redis.Nil {
panic(err)
}
bots := make([]discord.SimpleUser, 0)
for _, id := range botIDs {
bots = append(bots, GetBot(id).Simplify())
}
return bots
}
// one array of just bot ids
// `bot:(id):token`
// `bot:(id):data`
// SimpleDiscordUser
func GetBot(id string) discord.User {
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
}
func GetBots() {
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
}

View File

@@ -41,8 +41,6 @@ func CheckLogin(username string, password string, ip string) (bool, string) {
panic(err)
}
fmt.Println(user, username)
if user != username {
return false, ""
}
@@ -56,11 +54,27 @@ func CheckLogin(username string, password string, ip string) (bool, string) {
}
type LoginToken struct {
Token string `json:"token"`
ID string `json:"id"`
TokenHash string `json:"token"`
IP string `json:"ip"`
End string `json:"end"`
}
type LoginTokenSimple struct {
ID string `json:"id"`
IP string `json:"ip"`
End string `json:"end"`
}
func (t LoginToken) Simplify() LoginTokenSimple {
return LoginTokenSimple{
ID: t.ID,
IP: t.IP,
End: t.End,
}
}
func (t LoginToken) MarshalBinary() ([]byte, error) {
return json.Marshal(t)
}
@@ -70,21 +84,22 @@ func (t *LoginToken) UnmarshalBinary(data []byte) error {
}
func createLoginToken(ip string) string {
token := util.GeneratePassword(32)
token := util.GenerateToken()
tokenHash, err := bcrypt.GenerateFromPassword([]byte(token), bcrypt.DefaultCost)
if err != nil {
panic(err)
}
tokenData := LoginToken{
Token: string(tokenHash),
IP: ip,
End: token[len(token) - 4:],
ID: util.GenerateID(),
TokenHash: string(tokenHash),
IP: ip,
End: util.GetEnd(token),
}
member := redis.Z{
Score: float64(time.Now().Unix() + 4 * 60 * 60),
Score: float64(time.Now().Unix() + 4 * 60 * 60),
Member: tokenData,
}
@@ -96,8 +111,7 @@ func createLoginToken(ip string) string {
return token
}
func CheckLoginToken(token string, ip string) bool {
func getLoginTokens() []LoginToken {
expired, err := client.ZRangeByScore(ctx, "loginTokens", &redis.ZRangeBy{
Min: "-inf",
Max: fmt.Sprintf("%d", time.Now().Unix()),
@@ -118,10 +132,28 @@ func CheckLoginToken(token string, ip string) bool {
panic(err)
}
fmt.Printf("%v\n", current)
return current
}
func GetLoginTokensSimple() []LoginTokenSimple {
simple := make([]LoginTokenSimple, 0)
for _, c := range getLoginTokens() {
simple = append(simple, c.Simplify())
}
return simple
}
func ClearLoginTokens() {
client.Del(ctx, "loginTokens")
}
func CheckLoginToken(token string, ip string) bool {
current := getLoginTokens()
for _, c := range current {
err = bcrypt.CompareHashAndPassword([]byte(c.Token), []byte(token))
err := bcrypt.CompareHashAndPassword([]byte(c.TokenHash), []byte(token))
if err == nil && ip == c.IP {
return true
}

View File

@@ -15,9 +15,9 @@ var client *redis.Client = nil
func initializeRedis() {
username := "default"
UpdateUsername( username)
UpdateUsername(username)
password := util.GeneratePassword(16)
UpdatePassword( password)
UpdatePassword(password)
fmt.Printf("FIRST TIME SETUP\nusername: %s\npassword: %s\n\n\n", username, password)
}

229
storage/tokens.go Normal file
View File

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