From b61d8aaabce23d3fa5b3ec7186ed39f31cf9d915 Mon Sep 17 00:00:00 2001
From: ssongliu <73214554+ssongliu@users.noreply.github.com>
Date: Wed, 5 Jul 2023 14:02:12 +0800
Subject: [PATCH] =?UTF-8?q?fix:=20=E4=B8=BB=E6=9C=BA=E5=AF=86=E7=A0=81?=
 =?UTF-8?q?=E5=9B=9E=E6=98=BE=E8=A7=A3=E5=AF=86=20(#1540)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 backend/app/api/v1/host.go                | 29 +++++++++++++++++++---
 backend/app/service/host.go               | 30 +++++++++++++++++++++++
 backend/init/migration/migrations/init.go |  9 +++++++
 3 files changed, 65 insertions(+), 3 deletions(-)

diff --git a/backend/app/api/v1/host.go b/backend/app/api/v1/host.go
index 588aa545d..ac266a714 100644
--- a/backend/app/api/v1/host.go
+++ b/backend/app/api/v1/host.go
@@ -42,6 +42,8 @@ func (b *BaseApi) CreateHost(c *gin.Context) {
 			return
 		}
 		req.Password = passwordItem
+		req.PrivateKey = ""
+		req.PassPhrase = ""
 	}
 	if req.AuthMode == "key" && len(req.PrivateKey) != 0 {
 		privateKey, err := base64.StdEncoding.DecodeString(req.PrivateKey)
@@ -55,6 +57,16 @@ func (b *BaseApi) CreateHost(c *gin.Context) {
 			return
 		}
 		req.Password = keyItem
+
+		if len(req.PassPhrase) != 0 {
+			pass, err := encrypt.StringEncrypt(req.PassPhrase)
+			if err != nil {
+				helper.ErrorWithDetail(c, constant.CodeErrBadRequest, constant.ErrTypeInvalidParams, err)
+				return
+			}
+			req.PassPhrase = pass
+		}
+		req.Password = ""
 	}
 
 	host, err := hostService.Create(req)
@@ -229,6 +241,15 @@ func (b *BaseApi) UpdateHost(c *gin.Context) {
 			return
 		}
 		req.PrivateKey = keyItem
+
+		if len(req.PassPhrase) != 0 {
+			pass, err := encrypt.StringEncrypt(req.PassPhrase)
+			if err != nil {
+				helper.ErrorWithDetail(c, constant.CodeErrBadRequest, constant.ErrTypeInvalidParams, err)
+				return
+			}
+			req.PassPhrase = pass
+		}
 	}
 
 	upMap := make(map[string]interface{})
@@ -239,10 +260,12 @@ func (b *BaseApi) UpdateHost(c *gin.Context) {
 	upMap["user"] = req.User
 	upMap["auth_mode"] = req.AuthMode
 	upMap["remember_password"] = req.RememberPassword
-	if len(req.Password) != 0 {
+	if req.AuthMode == "password" {
 		upMap["password"] = req.Password
-	}
-	if len(req.PrivateKey) != 0 {
+		upMap["private_key"] = ""
+		upMap["pass_phrase"] = ""
+	} else {
+		upMap["password"] = ""
 		upMap["private_key"] = req.PrivateKey
 		upMap["pass_phrase"] = req.PassPhrase
 	}
diff --git a/backend/app/service/host.go b/backend/app/service/host.go
index 7efcde510..6db0d7972 100644
--- a/backend/app/service/host.go
+++ b/backend/app/service/host.go
@@ -105,6 +105,10 @@ func (u *HostService) TestLocalConn(id uint) bool {
 		connInfo.PrivateKey = []byte(host.PrivateKey)
 	}
 	if len(host.PassPhrase) != 0 {
+		host.PassPhrase, err = encrypt.StringDecrypt(host.PassPhrase)
+		if err != nil {
+			return false
+		}
 		connInfo.PassPhrase = []byte(host.PassPhrase)
 	}
 	client, err := connInfo.NewClient()
@@ -133,6 +137,13 @@ func (u *HostService) GetHostInfo(id uint) (*model.Host, error) {
 			return nil, err
 		}
 	}
+
+	if len(host.PassPhrase) != 0 {
+		host.PassPhrase, err = encrypt.StringDecrypt(host.PassPhrase)
+		if err != nil {
+			return nil, err
+		}
+	}
 	return &host, err
 }
 
@@ -153,6 +164,25 @@ func (u *HostService) SearchWithPage(search dto.SearchHostWithPage) (int64, inte
 			item.Password = ""
 			item.PrivateKey = ""
 			item.PassPhrase = ""
+		} else {
+			if len(host.Password) != 0 {
+				item.Password, err = encrypt.StringDecrypt(host.Password)
+				if err != nil {
+					return 0, nil, err
+				}
+			}
+			if len(host.PrivateKey) != 0 {
+				item.PrivateKey, err = encrypt.StringDecrypt(host.PrivateKey)
+				if err != nil {
+					return 0, nil, err
+				}
+			}
+			if len(host.PassPhrase) != 0 {
+				item.PassPhrase, err = encrypt.StringDecrypt(host.PassPhrase)
+				if err != nil {
+					return 0, nil, err
+				}
+			}
 		}
 		dtoHosts = append(dtoHosts, item)
 	}
diff --git a/backend/init/migration/migrations/init.go b/backend/init/migration/migrations/init.go
index 44fdc3802..2caab3669 100644
--- a/backend/init/migration/migrations/init.go
+++ b/backend/init/migration/migrations/init.go
@@ -460,6 +460,15 @@ var EncryptHostPassword = &gormigrate.Migration{
 					return err
 				}
 			}
+			if len(host.PassPhrase) != 0 {
+				pass, err := encrypt.StringEncrypt(host.PassPhrase)
+				if err != nil {
+					return err
+				}
+				if err := tx.Model(&model.Host{}).Update("pass_phrase", pass).Error; err != nil {
+					return err
+				}
+			}
 		}
 		return nil
 	},