From 0b5a30aec96fb7f0cee7201af86d41947342278e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=88=98=E7=A5=A5=E8=B6=85?= <iwind.liu@gmail.com>
Date: Sun, 7 Feb 2021 09:48:36 +0800
Subject: [PATCH] =?UTF-8?q?=E6=95=B0=E6=8D=AE=E5=BA=93=E8=8A=82=E7=82=B9?=
 =?UTF-8?q?=E5=AF=86=E7=A0=81=E6=94=AF=E6=8C=81=E5=8A=A0=E5=AF=86=E5=AD=98?=
 =?UTF-8?q?=E5=82=A8?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/db/models/db_node_dao.go      | 41 ++++++++++++++++++++++++--
 internal/db/models/db_node_dao_test.go | 38 ++++++++++++++++++++++++
 2 files changed, 76 insertions(+), 3 deletions(-)

diff --git a/internal/db/models/db_node_dao.go b/internal/db/models/db_node_dao.go
index 96fa6966..85f696ef 100644
--- a/internal/db/models/db_node_dao.go
+++ b/internal/db/models/db_node_dao.go
@@ -1,11 +1,14 @@
 package models
 
 import (
+	"encoding/base64"
+	"github.com/TeaOSLab/EdgeAPI/internal/encrypt"
 	"github.com/TeaOSLab/EdgeAPI/internal/errors"
 	_ "github.com/go-sql-driver/mysql"
 	"github.com/iwind/TeaGo/Tea"
 	"github.com/iwind/TeaGo/dbs"
 	"github.com/iwind/TeaGo/types"
+	"strings"
 )
 
 const (
@@ -15,6 +18,8 @@ const (
 
 type DBNodeDAO dbs.DAO
 
+const DBNodePasswordEncodedPrefix = "EDGE_ENCODED:"
+
 func NewDBNodeDAO() *DBNodeDAO {
 	return dbs.NewDAO(&DBNodeDAO{
 		DAOObject: dbs.DAOObject{
@@ -61,7 +66,9 @@ func (this *DBNodeDAO) FindEnabledDBNode(tx *dbs.Tx, id int64) (*DBNode, error)
 	if result == nil {
 		return nil, err
 	}
-	return result.(*DBNode), err
+	node := result.(*DBNode)
+	node.Password = this.DecodePassword(node.Password)
+	return node, nil
 }
 
 // 根据主键查找名称
@@ -88,6 +95,9 @@ func (this *DBNodeDAO) ListEnabledNodes(tx *dbs.Tx, offset int64, size int64) (r
 		Slice(&result).
 		DescPk().
 		FindAll()
+	for _, node := range result {
+		node.Password = this.DecodePassword(node.Password)
+	}
 	return
 }
 
@@ -102,7 +112,7 @@ func (this *DBNodeDAO) CreateDBNode(tx *dbs.Tx, isOn bool, name string, descript
 	op.Port = port
 	op.Database = database
 	op.Username = username
-	op.Password = password
+	op.Password = this.EncodePassword(password)
 	op.Charset = charset
 	err := this.Save(tx, op)
 	if err != nil {
@@ -125,7 +135,7 @@ func (this *DBNodeDAO) UpdateNode(tx *dbs.Tx, nodeId int64, isOn bool, name stri
 	op.Port = port
 	op.Database = database
 	op.Username = username
-	op.Password = password
+	op.Password = this.EncodePassword(password)
 	op.Charset = charset
 	err := this.Save(tx, op)
 	return err
@@ -139,5 +149,30 @@ func (this *DBNodeDAO) FindAllEnabledAndOnDBNodes(tx *dbs.Tx) (result []*DBNode,
 		Slice(&result).
 		DescPk().
 		FindAll()
+	for _, node := range result {
+		node.Password = this.DecodePassword(node.Password)
+	}
 	return
 }
+
+// 加密密码
+func (this *DBNodeDAO) EncodePassword(password string) string {
+	if strings.HasPrefix(password, DBNodePasswordEncodedPrefix) {
+		return password
+	}
+	encodedString := base64.StdEncoding.EncodeToString(encrypt.MagicKeyEncode([]byte(password)))
+	return DBNodePasswordEncodedPrefix + encodedString
+}
+
+// 解密密码
+func (this *DBNodeDAO) DecodePassword(password string) string {
+	if !strings.HasPrefix(password, DBNodePasswordEncodedPrefix) {
+		return password
+	}
+	dataString := password[len(DBNodePasswordEncodedPrefix):]
+	data, err := base64.StdEncoding.DecodeString(dataString)
+	if err != nil {
+		return password
+	}
+	return string(encrypt.MagicKeyDecode(data))
+}
diff --git a/internal/db/models/db_node_dao_test.go b/internal/db/models/db_node_dao_test.go
index 97c24b56..a39114ff 100644
--- a/internal/db/models/db_node_dao_test.go
+++ b/internal/db/models/db_node_dao_test.go
@@ -2,4 +2,42 @@ package models
 
 import (
 	_ "github.com/go-sql-driver/mysql"
+	"testing"
 )
+
+func TestDBNodeDAO_EncodePassword(t *testing.T) {
+	dao := NewDBNodeDAO()
+	for _, password := range []string{
+		"123456",
+		"abcdefxyz",
+		"123abc$*&^%",
+		"$%#@!@(*))*&^&=]{|",
+		"中文",
+	} {
+		encoded := dao.EncodePassword(password)
+		decoded := dao.DecodePassword(encoded)
+		if decoded != password {
+			t.Fatal(decoded, password)
+		}
+	}
+}
+
+func TestDBNodeDAO_EncodePassword_Encoded(t *testing.T) {
+	dao := NewDBNodeDAO()
+	password := DBNodePasswordEncodedPrefix + "123456"
+	encoded := dao.EncodePassword(password)
+	if encoded != password {
+		t.Fatal()
+	}
+	t.Log(encoded)
+}
+
+func TestDBNodeDAO_EncodePassword_Decoded(t *testing.T) {
+	dao := NewDBNodeDAO()
+	password := "123456"
+	decoded := dao.DecodePassword(password)
+	if decoded != password {
+		t.Fatal()
+	}
+	t.Log(decoded)
+}