106 lines
2.3 KiB
Go
106 lines
2.3 KiB
Go
package totp
|
||
|
||
import (
|
||
"crypto/hmac"
|
||
"crypto/sha1"
|
||
"encoding/base32"
|
||
"encoding/binary"
|
||
"fmt"
|
||
"strings"
|
||
"sync"
|
||
"time"
|
||
)
|
||
|
||
// TOTPConfig TOTP配置
|
||
type TOTPConfig struct {
|
||
Secret string // TOTP密钥
|
||
Digits int // TOTP验证码长度
|
||
TimeStep time.Duration // TOTP时间步长
|
||
Algorithm string // TOTP算法
|
||
}
|
||
|
||
var (
|
||
defaultConfig = TOTPConfig{
|
||
Secret: "",
|
||
Digits: 6,
|
||
TimeStep: 30 * time.Second,
|
||
Algorithm: "SHA1",
|
||
}
|
||
mu sync.RWMutex
|
||
)
|
||
|
||
// SetTOTPSecret 设置TOTP密钥
|
||
func SetTOTPSecret(secret string) {
|
||
mu.Lock()
|
||
defer mu.Unlock()
|
||
defaultConfig.Secret = secret
|
||
}
|
||
|
||
// GetTOTPSecret 获取TOTP密钥
|
||
func GetTOTPSecret() string {
|
||
mu.RLock()
|
||
defer mu.RUnlock()
|
||
return defaultConfig.Secret
|
||
}
|
||
|
||
// GenerateTOTPCode 生成TOTP验证码
|
||
func GenerateTOTPCode() (string, error) {
|
||
mu.RLock()
|
||
config := defaultConfig
|
||
mu.RUnlock()
|
||
|
||
if config.Secret == "" {
|
||
return "", fmt.Errorf("TOTP密钥未设置")
|
||
}
|
||
|
||
// 确保密钥是Base32编码的
|
||
secret := strings.ToUpper(config.Secret)
|
||
secret = strings.ReplaceAll(secret, " ", "")
|
||
secretBytes, err := base32.StdEncoding.DecodeString(secret)
|
||
if err != nil {
|
||
return "", fmt.Errorf("解码TOTP密钥失败: %w", err)
|
||
}
|
||
|
||
// 获取当前时间戳并转换为TOTP时间计数器
|
||
timeCounter := uint64(time.Now().Unix()) / uint64(config.TimeStep.Seconds())
|
||
|
||
// 将时间计数器转换为字节数组
|
||
timeBytes := make([]byte, 8)
|
||
binary.BigEndian.PutUint64(timeBytes, timeCounter)
|
||
|
||
// 使用HMAC-SHA1计算TOTP值
|
||
h := hmac.New(sha1.New, secretBytes)
|
||
h.Write(timeBytes)
|
||
hash := h.Sum(nil)
|
||
|
||
// 根据RFC 6238，我们使用哈希的最后一个字节的低4位作为偏移
|
||
offset := hash[len(hash)-1] & 0x0f
|
||
|
||
// 从哈希中提取4字节并转换为整数
|
||
binary := binary.BigEndian.Uint32(hash[offset : offset+4])
|
||
|
||
// 屏蔽最高位并获取指定位数的数字
|
||
totp := binary & 0x7fffffff % uint32(pow10(config.Digits))
|
||
|
||
// 将数字格式化为字符串，填充前导零
|
||
return fmt.Sprintf("%0*d", config.Digits, totp), nil
|
||
}
|
||
|
||
// ValidateTOTPCode 验证TOTP验证码
|
||
func ValidateTOTPCode(code string) bool {
|
||
generatedCode, err := GenerateTOTPCode()
|
||
if err != nil {
|
||
return false
|
||
}
|
||
return code == generatedCode
|
||
}
|
||
|
||
// pow10 计算10的n次方
|
||
func pow10(n int) int {
|
||
result := 1
|
||
for i := 0; i < n; i++ {
|
||
result *= 10
|
||
}
|
||
return result
|
||
}
|