From 016a43b792f5e20a8af873b4a935c700bba28c96 Mon Sep 17 00:00:00 2001
From: bjdgyc <bjdgyc@163.com>
Date: Wed, 10 Mar 2021 17:00:16 +0800
Subject: [PATCH] =?UTF-8?q?=E5=AE=8C=E5=96=84=E6=B5=8B=E8=AF=95=E6=96=87?=
 =?UTF-8?q?=E4=BB=B6?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 server/dbdata/db_test.go    |  7 +++---
 server/dbdata/group_test.go | 33 ++++++++++++++++++++++++++++
 server/dbdata/user.go       |  4 ++--
 server/dbdata/user_test.go  | 43 +++++++++++++++++++++++++++++++++++++
 server/pkg/utils/util.go    |  2 +-
 5 files changed, 83 insertions(+), 6 deletions(-)
 create mode 100644 server/dbdata/group_test.go
 create mode 100644 server/dbdata/user_test.go

diff --git a/server/dbdata/db_test.go b/server/dbdata/db_test.go
index 9950f9d..412bb46 100644
--- a/server/dbdata/db_test.go
+++ b/server/dbdata/db_test.go
@@ -22,12 +22,13 @@ func closeIpdata() {
 }
 
 func TestDb(t *testing.T) {
-	assert := assert.New(t)
+	ast := assert.New(t)
 	preIpData()
 	defer closeIpdata()
 
 	u := User{Username: "a"}
-	_ = Save(&u)
+	err := Save(&u)
+	ast.Nil(err)
 
-	assert.Equal(u.Id, 1)
+	ast.Equal(u.Id, 1)
 }
diff --git a/server/dbdata/group_test.go b/server/dbdata/group_test.go
new file mode 100644
index 0000000..ee1dee1
--- /dev/null
+++ b/server/dbdata/group_test.go
@@ -0,0 +1,33 @@
+package dbdata
+
+import (
+	"testing"
+
+	"github.com/bjdgyc/anylink/pkg/utils"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestGetGroupNames(t *testing.T) {
+	ast := assert.New(t)
+
+	preIpData()
+	defer closeIpdata()
+
+	// 添加 group
+	g1 := Group{Name: "g1", ClientDns: []ValData{{Val: "114.114.114.114"}}}
+	err := SetGroup(&g1)
+	ast.Nil(err)
+	g2 := Group{Name: "g2", ClientDns: []ValData{{Val: "114.114.114.114"}}}
+	err = SetGroup(&g2)
+	ast.Nil(err)
+	g3 := Group{Name: "g3", ClientDns: []ValData{{Val: "114.114.114.114"}}}
+	err = SetGroup(&g3)
+	ast.Nil(err)
+
+	// 判断所有数据
+	gAll := []string{"g1", "g2", "g3"}
+	gs := GetGroupNames()
+	for _, v := range gs {
+		ast.Equal(true, utils.InArrStr(gAll, v))
+	}
+}
diff --git a/server/dbdata/user.go b/server/dbdata/user.go
index f020f5d..32ab53d 100644
--- a/server/dbdata/user.go
+++ b/server/dbdata/user.go
@@ -35,7 +35,7 @@ func SetUser(v *User) error {
 	planPass := v.PinCode
 	// 自动生成密码
 	if len(planPass) < 6 {
-		planPass = utils.RandomNum(8)
+		planPass = utils.RandomRunes(8)
 	}
 	v.PinCode = planPass
 
@@ -82,7 +82,7 @@ func CheckUser(name, pwd, group string) error {
 	groupData := &Group{}
 	err = One("Name", group, groupData)
 	if err != nil || groupData.Status != 1 {
-		return fmt.Errorf("%s %s", name, "用户组错误")
+		return fmt.Errorf("%s - %s", name, "用户组错误")
 	}
 
 	// 判断otp信息
diff --git a/server/dbdata/user_test.go b/server/dbdata/user_test.go
new file mode 100644
index 0000000..e076672
--- /dev/null
+++ b/server/dbdata/user_test.go
@@ -0,0 +1,43 @@
+package dbdata
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/xlzd/gotp"
+)
+
+func TestCheckUser(t *testing.T) {
+	ast := assert.New(t)
+
+	preIpData()
+	defer closeIpdata()
+
+	group := "group1"
+
+	// 添加一个组
+	dns := []ValData{{Val: "114.114.114.114"}}
+	route := []ValData{{Val: "192.168.1.1/24"}}
+	g := Group{Name: group, Status: 1, ClientDns: dns, RouteInclude: route}
+	err := SetGroup(&g)
+	ast.Nil(err)
+	// 判断 IpMask
+	ast.Equal(g.RouteInclude[0].IpMask, "192.168.1.1/255.255.255.0")
+
+	// 添加一个用户
+	u := User{Username: "aaa", Groups: []string{group}, Status: 1}
+	err = SetUser(&u)
+	ast.Nil(err)
+
+	// 验证 PinCode + OtpSecret
+	totp := gotp.NewDefaultTOTP(u.OtpSecret)
+	secret := totp.Now()
+	err = CheckUser("aaa", u.PinCode+secret, group)
+	ast.Nil(err)
+
+	// 单独验证密码
+	u.DisableOtp = true
+	_ = SetUser(&u)
+	err = CheckUser("aaa", u.PinCode, group)
+	ast.Nil(err)
+}
diff --git a/server/pkg/utils/util.go b/server/pkg/utils/util.go
index 0ea8b56..77c08d5 100644
--- a/server/pkg/utils/util.go
+++ b/server/pkg/utils/util.go
@@ -61,7 +61,7 @@ func HumanByte(bf interface{}) string {
 	return hb
 }
 
-func RandomNum(length int) string {
+func RandomRunes(length int) string {
 	letterRunes := []rune("abcdefghijklmnpqrstuvwxy1234567890")
 
 	bytes := make([]rune, length)