wireguard-dashboard/component/jwt.go

package component

import (
	"context"
	"errors"
	"fmt"
	"gitee.ltd/lxh/logger/log"
	"github.com/golang-jwt/jwt/v5"
	"github.com/google/uuid"
	"math/rand"
	"strings"
	"time"
	"wireguard-ui/config"
	"wireguard-ui/global/client"
	"wireguard-ui/global/constant"
	"wireguard-ui/utils"
)

type JwtComponent struct {
	ID string `json:"id"`
	jwt.RegisteredClaims
}

// JWT
// @description: 初始化JWT组件
// @return JwtComponent
func JWT() JwtComponent {
	return JwtComponent{}
}

// GenerateToken
// @description: 生成token
// @receiver JwtComponent
// @param userId
// @param password
// @return token
// @return expireTime
// @return err
func (JwtComponent) GenerateToken(userId, secret, source string, times ...time.Time) (token string, expireTime *jwt.NumericDate, err error) {
	var notBefore, issuedAt *jwt.NumericDate
	if len(times) != 0 {
		expireTime = jwt.NewNumericDate(times[0])
		notBefore = jwt.NewNumericDate(times[1])
		issuedAt = jwt.NewNumericDate(times[1])
	} else {
		timeNow := time.Now().Local()
		expireTime = jwt.NewNumericDate(timeNow.Add(7 * time.Hour))
		notBefore = jwt.NewNumericDate(timeNow)
		issuedAt = jwt.NewNumericDate(timeNow)
	}

	claims := JwtComponent{
		ID: userId,
		RegisteredClaims: jwt.RegisteredClaims{
			Issuer:    config.Config.Http.Endpoint,                   // 颁发站点
			Subject:   "you can you up,no can no bb",                 // 发布主题
			ExpiresAt: expireTime,                                    // 过期时间
			NotBefore: notBefore,                                     // token不得早于该时间
			IssuedAt:  issuedAt,                                      // token颁发时间
			ID:        strings.ReplaceAll(uuid.NewString(), "-", ""), // 该token的id
		},
	}

	t := jwt.NewWithClaims(jwt.SigningMethodHS512, claims)
	token, err = t.SignedString([]byte(secret))
	if err != nil {
		log.Errorf("生成token失败: %v", err.Error())
		return "", nil, errors.New("生成token失败")
	}

	switch source {
	case "http":
		client.Redis.Set(context.Background(),
			fmt.Sprintf("%s:%s", constant.UserToken, userId),
			token,
			time.Duration(expireTime.Sub(time.Now()).Abs().Seconds())*time.Second)
	case "tui":
		client.Redis.Set(context.Background(),
			fmt.Sprintf("%s:%s", constant.TUIUserToken, userId),
			token,
			time.Duration(expireTime.Sub(time.Now()).Abs().Seconds())*time.Second)
	}

	return
}

// ParseToken
// @description: 解析token
// @receiver JwtComponent
// @param token
// @return *JwtComponent
// @return error
func (JwtComponent) ParseToken(token, secret, source string) (*JwtComponent, error) {
	tokenStr := strings.Split(token, "Bearer ")[1]

	t, err := jwt.ParseWithClaims(tokenStr, &JwtComponent{}, func(token *jwt.Token) (any, error) {
		return []byte(secret), nil
	})

	if claims, ok := t.Claims.(*JwtComponent); ok && t.Valid {
		var userToken string
		switch source {
		case "http":
			userToken, err = client.Redis.Get(context.Background(), fmt.Sprintf("%s:%s", constant.UserToken, claims.ID)).Result()
			if err != nil {
				log.Errorf("缓存中用户[%s]的token查找失败: %v", claims.ID, err.Error())
				return nil, errors.New("token不存在")
			}
		case "tui":
			userToken, err = client.Redis.Get(context.Background(), fmt.Sprintf("%s:%s", constant.TUIUserToken, claims.ID)).Result()
			if err != nil {
				log.Errorf("缓存中用户[%s]的token查找失败: %v", claims.ID, err.Error())
				return nil, errors.New("token不存在")
			}
		}

		if userToken != tokenStr {
			log.Errorf("token不一致")
			return nil, errors.New("token错误")
		}

		return claims, nil
	} else {
		return nil, err
	}
}

// GenerateSecret
// @description: 生成token解析密钥【每个用户的secret不一样，提高安全性】
// @receiver JwtComponent
// @param secret
// @return string
func (JwtComponent) GenerateSecret(secret ...string) string {
	// 添加10个元素,增加随机性
	for i := 0; i <= 10; i++ {
		secret = append(secret, uuid.NewString())
	}
	// 混淆一下明文secret的顺序
	n := len(secret)
	for i := n - 1; i > 0; i-- {
		j := rand.Intn(i + 1)
		secret[i], secret[j] = secret[j], secret[i]
	}
	secretStr := strings.Join(secret, ".")
	return utils.Hash().MD5(utils.Hash().SHA256(utils.Hash().SHA512(secretStr)))
}

// Logout
// @description: 退出登陆
// @receiver JwtComponent
// @param userId
// @return error
func (JwtComponent) Logout(userId string) error {
	return client.Redis.Del(context.Background(), fmt.Sprintf("%s:%s", constant.UserToken, userId)).Err()
}