discord-retokenizer/storage/login.go

package storage

import (
	"encoding/json"
	"fmt"
	"time"

	"git.zomo.dev/zomo/discord-retokenizer/util"
	"github.com/mileusna/useragent"
	"golang.org/x/crypto/bcrypt"
)

type UserAgentSimple struct {
	Name      string `json:"name"`
	Version   string `json:"version"`
	OS        string `json:"os"`
	OSVersion string `json:"os_version"`
	Mobile    bool   `json:"mobile"`
	Tablet    bool   `json:"tablet"`
	Desktop   bool   `json:"desktop"`
}

func (ua UserAgentSimple) Compare(ua2 UserAgentSimple) bool {
	return ua.Name == ua2.Name &&
		ua.OS == ua2.OS &&
		ua.Mobile == ua2.Mobile &&
		ua.Tablet == ua2.Tablet &&
		ua.Desktop == ua2.Desktop
}

func ParseUA(userAgentString string) UserAgentSimple {
	ua := useragent.Parse(userAgentString)
	return UserAgentSimple{
		Name:      ua.Name,
		Version:   ua.Version,
		OS:        ua.OS,
		OSVersion: ua.OSVersion,
		Mobile:    ua.Mobile,
		Tablet:    ua.Tablet,
		Desktop:   ua.Desktop,
	}
}

func UpdateUsername(username string) {
	if username != "" {
		client.Set(ctx, "username", username, 0)
	}
}

func UpdatePassword(password string) {
	if password != "" {
		passHash, err := bcrypt.GenerateFromPassword([]byte(password), bcrypt.DefaultCost)
		if err != nil {
			panic(err)
		}
		client.Set(ctx, "password", string(passHash), 0)
	}
}

func CheckLoginPassword(username string, password string, ip string, userAgent UserAgentSimple) (bool, string) {
	if username == "" || password == "" {
		return false, ""
	}

	user, err := client.Get(ctx, "username").Result()
	if err != nil {
		panic(err)
	}

	pass, err := client.Get(ctx, "password").Result()
	if err != nil {
		panic(err)
	}

	if user != username {
		return false, ""
	}

	err = bcrypt.CompareHashAndPassword([]byte(pass), []byte(password))
	if err != nil {
		return false, ""
	}
	
	//return existing token if it exists
	tokens := getLoginTokens()
	fmt.Printf("There are %d tokens", len(tokens))
	for _, token := range tokens {
		// fmt.Printf("Checking token %s\n", token.ID)
		// fmt.Printf("IP:\n  Given: %s\n  Token: %s\n", ip, token.IP)
		// fmt.Printf("UA:\n  Given: %+v\n  Token: %+v\n", userAgent, token.UserAgent)
		// fmt.Printf("Compare UA: %t\n", token.UserAgent.Compare(userAgent))
		if token.IP == ip && userAgent.Compare(token.UserAgent) {
			return true, token.Token
		}
	}
	
	return true, createLoginToken(ip, userAgent)
}

type LoginToken struct {
	ID        string          `json:"id"`
	Token     string          `json:"token"`
	IP        string          `json:"ip"`
	End       string          `json:"end"`
	UserAgent UserAgentSimple `json:"user_agent"`
	CreatedAt string          `json:"created_at"`
	LastLogin string          `json:"last_login"`
}


type LoginTokenSimple struct {
	ID        string 		  `json:"id"`
	IP        string 		  `json:"ip"`
	End       string 		  `json:"end"`
	UserAgent UserAgentSimple `json:"user_agent"`
	CreatedAt string          `json:"created_at"`
	LastLogin string 		  `json:"last_login"`
}

func (t LoginToken) Simplify() LoginTokenSimple {
	return LoginTokenSimple{
		ID: t.ID,
		IP: t.IP,
		End: t.End,
		UserAgent: t.UserAgent,
		CreatedAt: t.CreatedAt,
		LastLogin: t.LastLogin,
	}
}

func (t LoginToken) MarshalBinary() ([]byte, error) {
    return json.Marshal(t)
}

func (t *LoginToken) UnmarshalBinary(data []byte) error {
	return json.Unmarshal(data, t)
}

func createLoginToken(ip string, ua UserAgentSimple) string {
	token := util.GenerateToken()
	
	tokenData := LoginToken{
		ID: 	   util.GenerateID(),
		Token:     token,
		IP:		   ip,
		End:       util.GetEnd(token),
		UserAgent: ua,
		CreatedAt: time.Now().Format(time.RFC3339),
		LastLogin: time.Now().Format(time.RFC3339),
	}

	err := client.RPush(ctx, "loginTokens", tokenData.ID).Err()
	if err != nil {
		panic(err)
	}

	err = client.Set(ctx, "loginToken:"+tokenData.ID, tokenData, 0).Err()
	if err != nil {
		panic(err)
	}

	return token
}

func getLoginTokens() []LoginToken {
	var ids []string
	err := client.LRange(ctx, "loginTokens", 0, -1).ScanSlice(&ids)
	if err != nil {
		panic(err)
	}

	var tokens []LoginToken
	for _, id := range ids {
		var token LoginToken
		err = client.Get(ctx, "loginToken:"+id).Scan(&token)
		if err != nil {
			panic(err)
		}
		tokens = append(tokens, token)
	}
	return tokens
}

func updateLastLoginToken(token LoginToken) {
	token.CreatedAt = time.Now().Format(time.RFC3339)
	err := client.Set(ctx, "loginToken:"+token.ID, token, 0).Err()
	if err != nil {
		panic(err)
	}
}

func GetLoginTokensSimple() []LoginTokenSimple {
	simple := make([]LoginTokenSimple, 0)

	for _, c := range getLoginTokens() {
		simple = append(simple, c.Simplify())
	}

	return simple
}

func ClearLoginTokens() {
	var ids []string
	err := client.LRange(ctx, "loginTokens", 0, -1).ScanSlice(&ids)
	if err != nil {
		panic(err)
	}

	err = client.Del(ctx, "loginTokens").Err()
	if err != nil {
		panic(err)
	}

	for _, id := range ids {
		err = client.Del(ctx, "loginToken:"+id).Err()
		if err != nil {
			panic(err)
		}
	}
}

func CheckLoginToken(token string, ip string, userAgent UserAgentSimple) bool {
	current := getLoginTokens()

	for _, c := range current {
		if token == c.Token && ip == c.IP && userAgent.Compare(c.UserAgent) {
			fmt.Printf("Checking token %s\n", c.ID)
			fmt.Printf("IP:\n  Given: %s\n  Token: %s\n", ip, c.IP)
			fmt.Printf("UA:\n  Given: %+v\n  Token: %+v\n", userAgent, c.UserAgent)
			updateLastLoginToken(c)			
			return true
		}
	}

	return false
}