Device/SERVER/Meter_Service/core/utils/security.go

package utils

import (
	"crypto/rand"
	"encoding/hex"
	"errors"
	"log"

	"golang.org/x/crypto/scrypt"
)

const (
	SymbolSet = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789!@#$%^&*()-_=+,.?/:;{}[]`~"
	LetterSet = "ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
	NumberSet = "0123456789"
)

// ErrInvalidCharsetLength 表示字符集长度错误
var ErrInvalidCharsetLength = errors.New("字符集长度错误")

func generateRandString(length int, charset string) (string, error) {
	var chars = []byte(charset)
	clen := len(chars)
	if clen < 2 || clen > 256 {
		return "", ErrInvalidCharsetLength
	}
	maxRb := 255 - (256 % clen)
	b := make([]byte, length)
	r := make([]byte, length+(length/4)) // storage for random bytes.
	i := 0
	for {
		n, err := rand.Read(r)
		if err != nil {
			log.Printf("读取随机字节出错: %v", err)
			return "", err
		}
		if n != len(r) {
			log.Printf("未预期的随机字节数量: %d", n)
			return "", errors.New("随机字节长度不正确")
		}
		for _, rb := range r {
			c := int(rb)
			if c > maxRb {
				continue // Skip this number to avoid modulo bias.
			}
			b[i] = chars[c%clen]
			i++
			if i == length {
				return string(b), nil
			}
		}
	}
}

// GenerateRandomString 生成指定长度的随机字符串
func GenerateRandomString(length int, charset string) string {
	str, err := generateRandString(length, charset)
	if err != nil {
		log.Printf("生成随机字符串失败: %v", err)
		return ""
	}
	return str
}

// GenerateRandomKey 生成指定长度的随机字符串，使用预定义的字符集
func GenerateRandomKey(length int, charsetType int) string {
	var charset string
	switch charsetType {
	case 1: // Symbol set
		charset = SymbolSet
	case 2: // Letter set
		charset = LetterSet
	case 3: // Number set
		charset = NumberSet
	default:
		log.Println("无效的字符集类型")
		return ""
	}
	return GenerateRandomString(length, charset)
}

// GenerateRandomKey20 生成20位随机字符串
func GenerateRandomKey20() string {
	return GenerateRandomKey(20, 1)
}

// GenerateRandomKey16 生成16位随机字符串
func GenerateRandomKey16() string {
	return GenerateRandomKey(16, 1)
}

// GenerateRandomKey6 生成6位随机字符串
func GenerateRandomKey6() string {
	return GenerateRandomKey(6, 2)
}

// GenerateRandomKey4 生成4位随机字符串
func GenerateRandomKey4() string {
	return GenerateRandomKey(4, 3)
}

// SetPassword 根据明文密码和加盐值生成密码
func SetPassword(password string, salt string) (verify string, err error) {
	var rb []byte
	rb, err = scrypt.Key([]byte(password), []byte(salt), 16384, 8, 1, 32)
	if err != nil {
		log.Printf("生成密码哈希失败: %v", err)
		return
	}
	verify = hex.EncodeToString(rb)
	return
}