feat: [totp] Fork https://github.com/dgryski/dgoogauth function and u…

…nit testing

totp/totp.go

@@ -0,0 +1,120 @@
+package totp
+
+import (
+	"crypto/hmac"
+	"crypto/sha1"
+	"encoding/base32"
+	"encoding/binary"
+	"errors"
+	"sort"
+	"strconv"
+	"time"
+)
+
+var (
+	ErrNotMatch = errors.New("code does not match")
+)
+
+type Totp struct {
+	Secret        string
+	Window        int
+	Counter       int
+	DisallowReuse []int
+	ScratchCodes  []int
+}
+
+func (x *Totp) Authenticate(password string) (bool, error) {
+	var scratch bool
+	switch {
+	case len(password) == 6 && password[0] >= '0' && password[0] <= '9':
+		break
+	case len(password) == 8 && password[0] >= '1' && password[0] <= '9':
+		scratch = true
+		break
+	default:
+		return false, ErrNotMatch
+	}
+	code, err := strconv.Atoi(password)
+	if err != nil {
+		return false, ErrNotMatch
+	}
+	if scratch {
+		return x.CheckScratchCodes(code), nil
+	}
+	if x.Counter > 0 {
+		return x.CheckCode(code), nil
+	}
+	ts := int(time.Now().UTC().Unix() / 30)
+	return x.CheckTotpCode(ts, code), nil
+}
+
+func (x *Totp) CheckScratchCodes(code int) bool {
+	for i, v := range x.ScratchCodes {
+		if code == v {
+			l := len(x.ScratchCodes) - 1
+			x.ScratchCodes[i] = x.ScratchCodes[l]
+			x.ScratchCodes = x.ScratchCodes[0:l]
+			return true
+		}
+	}
+	return false
+}
+
+func (x *Totp) CheckCode(code int) bool {
+	for i := 0; i < x.Window; i++ {
+		if Compute(x.Secret, int64(x.Counter+i)) == code {
+			x.Counter += i + 1
+			return true
+		}
+	}
+	x.Counter++
+	return false
+}
+
+func (x *Totp) CheckTotpCode(ts, code int) bool {
+	minT := ts - (x.Window / 2)
+	maxT := ts + (x.Window / 2)
+	for t := minT; t <= maxT; t++ {
+		if Compute(x.Secret, int64(t)) == code {
+			if x.DisallowReuse != nil {
+				for _, timeCode := range x.DisallowReuse {
+					if timeCode == t {
+						return false
+					}
+				}
+				x.DisallowReuse = append(x.DisallowReuse, t)
+				sort.Ints(x.DisallowReuse)
+				m := 0
+				for x.DisallowReuse[m] < minT {
+					m++
+				}
+				x.DisallowReuse = x.DisallowReuse[m:]
+			}
+			return true
+		}
+	}
+	return false
+}
+
+func Compute(secret string, value int64) int {
+	key, err := base32.StdEncoding.DecodeString(secret)
+	if err != nil {
+		return -1
+	}
+
+	hash := hmac.New(sha1.New, key)
+	err = binary.Write(hash, binary.BigEndian, value)
+	if err != nil {
+		return -1
+	}
+	h := hash.Sum(nil)
+
+	offset := h[19] & 0x0f
+
+	truncated := binary.BigEndian.Uint32(h[offset : offset+4])
+
+	truncated &= 0x7fffffff
+	code := truncated % 1000000
+
+	return int(code)
+}

totp/totp_test.go

@@ -0,0 +1,104 @@
+package totp_test
+
+import (
+	"fmt"
+	"github.com/stretchr/testify/assert"
+	"github.com/weplanx/go/totp"
+	"testing"
+	"time"
+)
+
+type Value1 struct {
+	code   string
+	result bool
+}
+
+func TestAuthenticate(t *testing.T) {
+	x := &totp.Totp{
+		Secret:       "2SH3V3GDW7ZNMGYE",
+		Window:       3,
+		Counter:      1,
+		ScratchCodes: []int{11112222, 22223333},
+	}
+	values := []Value1{
+		{"foobar", false},
+		{"1fooba", false},
+		{"1111111", false},
+		{"293240", true},
+		{"293240", false},
+		{"33334444", false},
+		{"11112222", true},
+		{"11112222", false},
+	}
+	for _, v := range values {
+		r, _ := x.Authenticate(v.code)
+		assert.Equal(t, v.result, r)
+	}
+	x.Counter = 0
+	ts := time.Now().UTC().Unix() / 30
+	code := fmt.Sprintf("%06d", totp.Compute(x.Secret, ts))
+	values = []Value1{
+		{code + "1", false},
+		{code, true},
+	}
+	for _, v := range values {
+		r, _ := x.Authenticate(v.code)
+		assert.Equal(t, v.result, r)
+	}
+}
+
+type Value2 struct {
+	code   int
+	ts     int
+	result bool
+}
+
+type Value3 struct {
+	code       int
+	ts         int
+	result     bool
+	disallowed []int
+}
+
+func TestTotpCode(t *testing.T) {
+	var x totp.Totp
+	x.Secret = "2SH3V3GDW7ZNMGYE"
+	x.Window = 5
+	values := []Value2{
+		{50548, 9997, false},
+		{50548, 9998, true},
+		{50548, 9999, true},
+		{50548, 10000, true},
+		{50548, 10001, true},
+		{50548, 10002, true},
+		{50548, 10003, false},
+	}
+
+	for _, v := range values {
+		r := x.CheckTotpCode(v.ts, v.code)
+		assert.Equal(t, v.result, r)
+	}
+
+	x.DisallowReuse = make([]int, 0)
+	var noreuses = []Value3{
+		{50548 /* 10000 */, 9997, false, []int{}},
+		{50548 /* 10000 */, 9998, true, []int{10000}},
+		{50548 /* 10000 */, 9999, false, []int{10000}},
+		{478726 /* 10001 */, 10001, true, []int{10000, 10001}},
+		{646986 /* 10002 */, 10002, true, []int{10000, 10001, 10002}},
+		{842639 /* 10003 */, 10003, true, []int{10001, 10002, 10003}},
+	}
+
+	for _, v := range noreuses {
+		r := x.CheckTotpCode(v.ts, v.code)
+		assert.Equal(t, v.result, r)
+		assert.Equal(t, len(x.DisallowReuse), len(v.disallowed))
+		same := true
+		for i := range v.disallowed {
+			if v.disallowed[i] != x.DisallowReuse[i] {
+				same = false
+			}
+		}
+		assert.True(t, same)
+	}
+}