package otp

import (
        "errors"
        "strings"
)

var (
        ErrURIFormat           = errors.New("uri format error")
        ErrSecretDecode        = errors.New("secret base32 decode error")
        ErrSecretCannotBeEmpty = errors.New("secret cannot be empty")
)

var (
        minSkewNumber   = 0
        minPeriodNumber = 10
)

// Algorithms 支持的 HMAC 类型。
//
// 默认值：HMAC_SHA1，与 Google Authenticator 兼容。
//
// See https://github.com/google/google-authenticator/wiki/Key-Uri-Format
type Algorithms int

const (
        AlgorithmSHA1 Algorithms = iota + 1
        AlgorithmSHA256
        AlgorithmSHA512
)

// String 枚举值转换为字符串形式 - 该值可以放置在 uri 上。
func (h Algorithms) String() string {
        switch h {
        case AlgorithmSHA1:
                return "SHA1"
        case AlgorithmSHA256:
                return "SHA256"
        case AlgorithmSHA512:
                return "SHA512"
        default:
                panic("unreachable")
        }
}

// from 从字符串转换至 Algorithms 枚举
func (h Algorithms) from(str string) (Algorithms, error) {
        switch strings.ToUpper(str) {
        case "":
                return AlgorithmSHA1, nil
        case "SHA1":
                return AlgorithmSHA1, nil
        case "SHA256":
                return AlgorithmSHA256, nil
        case "SHA512":
                return AlgorithmSHA512, nil
        default:
                return 0, errors.New("unknown 'algorithm' string")
        }
}

// Digits 生成出来的一次性密码的长度。6 和 8 是最常见的值。
type Digits int

const (
        DigitsSix   Digits = 6
        DigitsEight Digits = 8
)

// from 从 int 类型转换至 Digits 枚举
func (d Digits) from(i int) (Digits, error) {
        switch i {
        case 6:
                return DigitsSix, nil
        case 8:
                return DigitsEight, nil
        default:
                return 0, errors.New("unknown 'digits' number")
        }
}

package otp

import (
        "crypto/hmac"
        "fmt"
        "net/url"
)

// HOTP 基于 RFC-4266 的 HOTP 算法
type HOTP struct {
        Otp
        // base32 encoded string
        Secret string
        // base32 decoded string
        decodedSecret []byte
}

// NewHOTP 创建一个 HOTP 结构体，可以使用 option 的模式传递参数。
//
// Params:
//
//        secret       : 必传，一个 base32 编码后的字符串，建议使用 RandomSecret 方法生成的。
//        WithCounter  : 设置初始计数器，该值仅用于 KeyURI 方法。
//        WithSkew     : 是否校验相邻的窗口。
//        WithAlgorithm: 设置 hmac 算法类型。
//
// Panic:
//   - secret base32 decode error
//   - secret is an empty string
//
// 注意: Google Authenticator 可能仅支持 Counter 这一个参数
//
// See https://github.com/google/google-authenticator/wiki/Key-Uri-Format
//
// Example:
//
//        secret := Base32Encode(RandomSecret(20))
//        hotp   := NewHOTP(secret, WithCounter(2))
func NewHOTP(secret string, options ...Option) *HOTP {
        if secret == "" {
                panic(ErrSecretCannotBeEmpty)
        }
        decodedSecret, err := Base32Decode(secret)
        if err != nil {
                panic(ErrSecretDecode)
        }
        otp := Otp{
                Skew:      0,
                Counter:   1,
                Period:    30,
                Algorithm: AlgorithmSHA1,
                Digits:    DigitsSix,
        }
        for _, opt := range options {
                opt(&otp)
        }
        return &HOTP{
                Otp:           otp,
                Secret:        secret,
                decodedSecret: decodedSecret,
        }
}

// At 通过指定的 Counter 生成一个 token。
//
// Example：
//
//        hotp  := NewHOTP(Base32Encode(RandomSecret(20)))
//        token := hotp.At(1)                 // 使用的 1 作为counter 生成 token
//        bool  := hotp.Verify(token, 1) // 校验 token 是否有效
func (h *HOTP) At(counter int64) string {
        s := intToByte(counter)
        hashFunc := hasher(h.Algorithm)
        mac := hmac.New(hashFunc, h.decodedSecret)
        mac.Write(s)
        hex := mac.Sum(nil)
        return truncate(hex, int(h.Digits))
}

// Verify 校验token是否有效，窗口内的所有结果都认为有效。
//
// Params:
//
//        token  : 需要进行校验的参数，一个字符串，如果字符串为空将会返回 false
//        counter: 计数器
//
// Example:
//
//        hotp  := NewHOTP(Base32Encode(RandomSecret(20)), WithSkew(1))
//        token := hotp.At(2)                     // 使用的 2 作为counter 生成 token
//        bool  := hotp.Verify(token, 2) // 通过 WithSkew 方法指定 skew 参数为1，那么这里将会校验 counter 为 1、2、3 的token
func (h *HOTP) Verify(token string, counter int64) bool {
        if token == "" {
                return false
        }
        c := counter
        for i := c - int64(h.Skew); i <= c+int64(h.Skew); i++ {
                if h.At(i) == token {
                        return true
                }
        }
        return false
}

// KeyURI 返回一个 KeyURI 结构体，其包含转换至 URI 和生成二维码的方法。
func (h *HOTP) KeyURI(account, issuer string) *KeyURI {
        ret := &KeyURI{
                Type:      "hotp",
                Label:     url.PathEscape(fmt.Sprintf("%s:%s", issuer, account)),
                Counter:   h.Counter,
                Digits:    int(h.Digits),
                Algorithm: h.Algorithm.String(),
                Issuer:    url.QueryEscape(issuer),
                Secret:    h.Secret,
        }
        return ret
}

package otp

import (
        "fmt"
        "github.com/skip2/go-qrcode"
        "net/url"
        "strconv"
        "strings"
)

// KeyURI TOTP 或 HOTP 的 URI 包含的参数。
//
// URI 的格式可以参考：https://github.com/google/google-authenticator/wiki/Key-Uri-Format
//
// 部分属性 Google Authenticator 可能不会采用仅支持默认值。具体细节可以查看上面链接的文档。
type KeyURI struct {
        // otp 算法的类型只能是 totp or hotp
        Type string
        // 标签，用于识别密钥与哪个帐户关联。它包含一个帐户名称，该名称是一个 URI 编码的字符串，可以选择以标识管理该帐户的提供商或服务的发行者字符串为前缀。
        // 发行者前缀和帐户名称应使用文字或 URL 编码的冒号分隔，并且帐户名称之前可以有可选空格。发行人或账户名称本身都不能包含冒号。
        // 根据 Google Authenticator 的建议，应该拼接发行商字符串为前缀。
        // 需要已被 url.QueryEscape 方法处理过。
        Label string
        // hotp 或 totp 采用的哈希算法类型
        // Google Authenticator 可能会忽略此参数，而采用默认值：HMAC-SHA1。
        Algorithm string
        // 向用户显示一次性密码的长度。默认值为 6。
        // Google Authenticator 可能会忽略此参数，而采用默认值 6。
        Digits int
        // 当 type 为 hotp 时必选，它将设置初始计数器值。
        Counter int64
        // 仅当 type 为 totp 时可选，该 period 参数定义 TOTP 密码的有效期限（以秒为单位）。默认值为 30。
        // Google Authenticator 可能会忽略此参数，而采用默认值 30。
        Period int
        // 发行商，使用 URL 编码进行编码的字符串
        // 需要已被 url.QueryEscape 方法处理过。
        Issuer string
        // base32 编码的任意字符，不应该填充。
        Secret string
}

// URI 生成 otpauth 的 URI 形式，可以将其作为二维码的内容供 Google Authenticator 扫码导入。
// params 顺序：secret、issuer、algorithm、digits、period、counter
func (p KeyURI) URI() *url.URL {
        u := url.URL{}
        u.Scheme = "otpauth"
        u.Host = p.Type
        u.Path = p.Label
        params := "secret=" + p.Secret
        params += "&issuer=" + p.Issuer

        if p.Algorithm != "SHA1" {
                params += "&algorithm=" + p.Algorithm
        }
        if p.Digits != 6 {
                params += "&digits=" + strconv.Itoa(p.Digits)
        }
        if p.Type == "totp" {
                if p.Period != 30 {
                        params += "&period=" + strconv.Itoa(p.Period)
                }
        } else {
                params += "&counter=" + strconv.FormatInt(p.Counter, 10)
        }
        u.RawQuery = params
        return &u
}

// QRCode 将此 URI 信息生成一个二维码，可供 Google Authenticator 扫码导入。
func (p KeyURI) QRCode() ([]byte, error) {
        uri := p.URI().String()
        code, err := qrcode.New(uri, qrcode.Highest)
        if err != nil {
                return nil, err
        }
        png, err := code.PNG(256)
        if err != nil {
                return nil, err
        }
        return png, nil
}

// FromURI 解析 URI 创建一个 KeyURI 结构体。
func FromURI(uri string) (*KeyURI, error) {
        u, err := url.Parse(uri)
        if err != nil {
                return nil, ErrURIFormat
        }
        if u.Scheme != "otpauth" {
                return nil, ErrURIFormat
        }
        if u.Host != "hotp" && u.Host != "totp" {
                return nil, ErrURIFormat
        }
        query := u.Query()
        issuer := query.Get("issuer")
        secret := query.Get("secret")
        if secret == "" {
                return nil, ErrURIFormat
        }
        digits, err := atoi(query.Get("digits"), 6)
        if err != nil {
                return nil, ErrURIFormat
        }
        digitsEnum, err := Digits.from(DigitsSix, digits)
        if err != nil {
                return nil, ErrURIFormat
        }
        period, err := atoi(query.Get("period"), 30)
        if err != nil || period < minPeriodNumber {
                return nil, ErrURIFormat
        }
        counter, err := parseInt(query.Get("counter"), 1, 10, 64)
        if err != nil {
                return nil, ErrURIFormat
        }
        algorithm, err := Algorithms.from(AlgorithmSHA1, query.Get("algorithm"))
        if err != nil {
                return nil, ErrURIFormat
        }

        if u.Host == "hotp" {
                period = 0
        } else {
                counter = 0
        }

        // 按照规则 issuer 和 account 都不能包含 ":"
        path := strings.Split(u.Path, ":")
        // 如果 label 存在 issuer 但是 params 中不存在 issuer
        if issuer == "" {
                if len(path) > 1 {
                        issuer = path[0][1:]
                }
        }
        var label = u.Path[1:]
        // 如果 label 不存在 issuer 但是 params 中存在 issuer
        if len(path) == 1 && issuer != "" {
                label = fmt.Sprintf("%s:%s", issuer, u.Path[1:])
        }

        key := &KeyURI{
                Type:      u.Host,
                Label:     label,
                Algorithm: algorithm.String(),
                Digits:    int(digitsEnum),
                Counter:   counter,
                Period:    period,
                Issuer:    issuer,
                Secret:    secret,
        }
        return key, nil
}

package otp

type Otp struct {
        // 指定时间窗口，默认 30 秒有效期。
        // Google Authenticator 可能仅支持默认参数。
        Period int
        // 初始计数器数值，默认为 1。
        // 该参数仅用来指定 otpauth uri 上的 counter 参数，不会使用它来生成 token
        Counter int64
        // 指定一次性密码的长度，默认 6 位数字。
        // Google Authenticator 可能仅支持默认参数。
        Digits Digits
        // 是否校验相邻的时间窗口，默认为 0。
        // 有些时候服务端的时间和客户端的时间并不是同步的，存在时间误差，再加上网络延时，一次性密码的剩余有效期等等，密码刚到达服务端可能就过期了，
        // 这时候可以通过此参数为相邻几个时间窗口进行校验，加强用户体验，但是安全性降低了。
        // 如果此参数为1，那么会同时校验当前时间窗口、上个时间窗口以及下个时间窗口。如果是 HOTP 那么就是相邻的计数器。
        Skew int
        // 指定 hmac 算法，默认 hmac-sha1
        // Google Authenticator 可能仅支持默认参数。
        Algorithm Algorithms
}

type Option func(opt *Otp)

// WithSkew 配置同时校验的窗口数，默认为 0 仅校验当前时间窗口。
//
// 取值范围是：skew >=0 如果传入的值小于 0 将会设置为 0。
func WithSkew(skew int) Option {
        return func(opt *Otp) {
                if skew < minSkewNumber {
                        skew = minSkewNumber
                }
                opt.Skew = skew
        }
}

// WithDigits 配置一次性密码的显示长度，默认为 6, Google Authenticator 可能不支持其他的长度。
func WithDigits(digits Digits) Option {
        return func(opt *Otp) {
                opt.Digits = digits
        }
}

// WithPeriod 配置时间一次性密码的有效期，默认 30 秒，仅支持 TOTP 类型。
//
// 取值范围是：period >=10 如果传入的值小于 10 将会设置为 10。
func WithPeriod(period int) Option {
        return func(opt *Otp) {
                if period < minPeriodNumber {
                        period = minPeriodNumber
                }
                opt.Period = period
        }
}

// WithCounter 配置计数器的值，默认为 1 (Google 的默认就是 1)，仅支持 HOTP 类型。
func WithCounter(counter int64) Option {
        return func(opt *Otp) {
                opt.Counter = counter
        }
}

// WithAlgorithm 配置哈希算法类型。
func WithAlgorithm(algorithm Algorithms) Option {
        return func(opt *Otp) {
                opt.Algorithm = algorithm
        }
}

package otp

import (
        "crypto/rand"
        "crypto/sha1"
        "crypto/sha256"
        "crypto/sha512"
        "encoding/base32"
        "hash"
        "math"
        "strconv"
        "strings"
)

// RandomSecret 获取一个给定长度(字节数)的随机秘钥，如果生成失败将会 panic。
//
// 建议存储时将其转换至 base32 或其他的编码，直接转换成字符串可能会存在换行符等奇怪的字符。
//
// 内部使用 rand.Read 方法，如果此方法报错将会 panic
//
// rfc4266 中建议 secret 最少为 160 位也就是 20 个字节。
//
//        https://datatracker.ietf.org/doc/html/rfc4226
//
// 也可看下此文档解释自行选择合适长度：
//
//        https://github.com/darrenedale/php-totp/blob/HEAD/Secrets.md
func RandomSecret(length int) []byte {
        // 建议选择适合对应 hmac 算法的长度。
        // HMAC-SHA1   建议选择 20 字节长度
        // HMAC-SHA256 建议选择 32 字节长度
        // HMAC-SHA512 建议选择 64 字节长度
        randomBytes := make([]byte, length)
        _, err := rand.Read(randomBytes)
        if err != nil {
                panic(err)
        }
        return randomBytes
}

// Base32Decode 对一个字符串进行 base32 解码
func Base32Decode(str string) ([]byte, error) {
        // base32 只包含大小字母
        upper := strings.ToUpper(str)
        return base32.StdEncoding.WithPadding(base32.NoPadding).DecodeString(upper)
}

// Base32Encode 对一个字符串进行 base32 编码
func Base32Encode(str []byte) string {
        return base32.StdEncoding.WithPadding(base32.NoPadding).EncodeToString(str)
}

// padZero 在字符串的签名填充数字0
func padZero(value string, size int) string {
        if len(value) >= size {
                return value
        }
        return strings.Repeat("0", size-len(value)) + value
}

// intToByte 数字转换成二进制字节格式
func intToByte(number int64) []byte {
        result := make([]byte, 0, 8)
        shifts := []uint{56, 48, 40, 32, 24, 16, 8, 0}
        for _, shift := range shifts {
                result = append(result, byte((number>>shift)&0xff))
        }
        return result
}

// truncate 计算出指定位数的数字字符串(不足位数前面补0)
func truncate(h []byte, digits int) string {
        offset := h[len(h)-1] & 0xf
        bits := uint32(h[offset]&0x7f)<<24 |
                uint32(h[offset+1]&0xff)<<16 |
                uint32(h[offset+2]&0xff)<<8 |
                uint32(h[offset+3]&0xff)
        value := bits % uint32(math.Pow10(digits))
        return padZero(strconv.Itoa(int(value)), digits)
}

func hasher(algorithm Algorithms) func() hash.Hash {
        switch algorithm {
        case AlgorithmSHA1:
                return sha1.New
        case AlgorithmSHA256:
                return sha256.New
        case AlgorithmSHA512:
                return sha512.New
        default:
                panic("unreachable")
        }
}

func atoi(str string, def int) (int, error) {
        if str == "" {
                return def, nil
        }
        val, err := strconv.Atoi(str)
        if err != nil {
                return 0, err
        }
        return val, nil
}

func parseInt(str string, def int64, base int, bitSize int) (int64, error) {
        if str == "" {
                return def, nil
        }
        val, err := strconv.ParseInt(str, base, bitSize)
        if err != nil {
                return 0, err
        }
        return val, nil
}

package otp

import (
        "crypto/hmac"
        "fmt"
        "net/url"
        "time"
)

// TOTP 基于 RFC-6238 的 TOTP 算法
type TOTP struct {
        Otp
        // base32 encoded string
        Secret string
        // base32 decoded string
        decodedSecret []byte
}

// NewTOTP 创建一个 TOTP 结构体，可以使用 option 的模式传递参数。
//
// Params:
//
//        secret       : 必传，一个 base32 编码后的字符串，建议使用 RandomSecret 方法生成的。
//        WithPeriod   : 设置 token 有效期长度。
//        WithSkew     : 是否校验相邻的窗口。
//        WithAlgorithm: 设置 hmac 算法类型。
//
// Panic:
//   - secret base32 decode error
//   - secret is an empty string
//
// 默认参数才是 Google Authenticator 兼容的，自定义参数的话 Google Authenticator 可能不会识别。
//
// See https://github.com/google/google-authenticator/wiki/Key-Uri-Format
//
// Example:
//
//        secret := Base32Encode(RandomSecret(20))
//        totp   := NewTOTP(secret, WithDigits(DigitsEight))
func NewTOTP(secret string, options ...Option) *TOTP {
        if secret == "" {
                panic(ErrSecretCannotBeEmpty)
        }
        decodedSecret, err := Base32Decode(secret)
        if err != nil {
                fmt.Println(err, secret)
                panic(ErrSecretDecode)
        }
        otp := Otp{
                Skew:      0,
                Counter:   1,
                Period:    30,
                Algorithm: AlgorithmSHA1,
                Digits:    DigitsSix,
        }
        for _, opt := range options {
                opt(&otp)
        }
        return &TOTP{
                Otp:           otp,
                Secret:        secret,
                decodedSecret: decodedSecret,
        }
}

// Now 基于当前时间点生成 token。
func (o *TOTP) Now() string {
        return o.At(time.Now())
}

// At 生成某个时间点的 token。
func (o *TOTP) At(t time.Time) string {
        key := intToByte(t.Unix() / int64(o.Period))
        hashFunc := hasher(o.Algorithm)
        mac := hmac.New(hashFunc, o.decodedSecret)
        mac.Write(key)
        h := mac.Sum(nil)
        return truncate(h, int(o.Digits))
}

// WithExpiration 获取指定时间的 token 和对应的剩余有效时间。
func (o *TOTP) WithExpiration(t time.Time) (string, int) {
        token := o.At(t)
        expiration := o.Expiration(t)
        return token, expiration
}

// Expiration 获取指定时间窗口的 token 剩余有效时间。
func (o *TOTP) Expiration(t time.Time) int {
        return int(int64(o.Period) - t.Unix()%int64(o.Period))
}

// Verify 校验 token 是否在指定的时间有效。
//
// Params:
//
//        token: 需要进行校验的参数，一个字符串，如果字符串为空将会返回 false。
//        t    : 指定的时间，用以校验 token 在这个时间点是否仍有效。
func (o *TOTP) Verify(token string, t time.Time) bool {
        if token == "" {
                return false
        }
        givenTime := t
        sec := t.Unix()
        for i := o.Skew * -1; i <= o.Skew; i++ {
                givenTime = time.Unix(sec, 0).Add(time.Second * time.Duration(o.Period*i))
                if o.At(givenTime) == token {
                        return true
                }
        }
        return false
}

// KeyURI 返回一个 KeyURI 结构体，其包含转换至 URI 和生成二维码的方法。
func (o *TOTP) KeyURI(account, issuer string) *KeyURI {
        ret := &KeyURI{
                Type:      "totp",
                Label:     url.PathEscape(fmt.Sprintf("%s:%s", issuer, account)),
                Algorithm: o.Algorithm.String(),
                Digits:    int(o.Digits),
                Period:    o.Period,
                Issuer:    url.QueryEscape(issuer),
                Secret:    o.Secret,
        }
        return ret
}