From 9fb065f65fa427f44fc2d221210c83fe6d7a502c Mon Sep 17 00:00:00 2001
From: CaptainB <bin@fit2cloud.com>
Date: Wed, 6 Sep 2023 19:23:47 +0800
Subject: [PATCH] =?UTF-8?q?ci:=20totp=20=E9=85=8D=E7=BD=AE,=20=E4=B8=BB?=
 =?UTF-8?q?=E5=BA=94=E7=94=A8=E4=B8=8B=E5=8F=91=E6=89=A7=E8=A1=8C=E4=BB=BB?=
 =?UTF-8?q?=E5=8A=A1=E6=97=B6=E5=B8=A6=E4=B8=8Aotp-token,=20=E4=BF=9D?=
 =?UTF-8?q?=E8=AF=81=E5=AE=89=E5=85=A8=E6=80=A7?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../io/metersphere/sdk/config/TOTPConfig.java | 25 +++++++++++++++++++
 .../sdk/constants/MsHttpHeaders.java          |  5 ++++
 .../sdk/service/NodeResourcePoolService.java  | 17 +++++++++++--
 3 files changed, 45 insertions(+), 2 deletions(-)
 create mode 100644 backend/framework/sdk/src/main/java/io/metersphere/sdk/config/TOTPConfig.java
 create mode 100644 backend/framework/sdk/src/main/java/io/metersphere/sdk/constants/MsHttpHeaders.java

diff --git a/backend/framework/sdk/src/main/java/io/metersphere/sdk/config/TOTPConfig.java b/backend/framework/sdk/src/main/java/io/metersphere/sdk/config/TOTPConfig.java
new file mode 100644
index 0000000000..d43227e026
--- /dev/null
+++ b/backend/framework/sdk/src/main/java/io/metersphere/sdk/config/TOTPConfig.java
@@ -0,0 +1,25 @@
+package io.metersphere.sdk.config;
+
+import com.bastiaanjansen.otp.HMACAlgorithm;
+import com.bastiaanjansen.otp.TOTPGenerator;
+import org.springframework.beans.factory.annotation.Value;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+
+import java.nio.charset.StandardCharsets;
+
+@Configuration
+public class TOTPConfig {
+    @Value("${totp.secret:secret}")
+    private String secret;
+
+    @Bean
+    public TOTPGenerator totpGenerator() {
+        return new TOTPGenerator.Builder(secret.getBytes(StandardCharsets.UTF_8))
+                .withHOTPGenerator(builder -> {
+                    builder.withPasswordLength(6);
+                    builder.withAlgorithm(HMACAlgorithm.SHA256); // SHA256 and SHA512 are also supported
+                })
+                .build();
+    }
+}
diff --git a/backend/framework/sdk/src/main/java/io/metersphere/sdk/constants/MsHttpHeaders.java b/backend/framework/sdk/src/main/java/io/metersphere/sdk/constants/MsHttpHeaders.java
new file mode 100644
index 0000000000..1987687139
--- /dev/null
+++ b/backend/framework/sdk/src/main/java/io/metersphere/sdk/constants/MsHttpHeaders.java
@@ -0,0 +1,5 @@
+package io.metersphere.sdk.constants;
+
+public interface MsHttpHeaders {
+    String OTP_TOKEN = "otp-token";
+}
diff --git a/backend/framework/sdk/src/main/java/io/metersphere/sdk/service/NodeResourcePoolService.java b/backend/framework/sdk/src/main/java/io/metersphere/sdk/service/NodeResourcePoolService.java
index f91175a58c..302e620c9d 100644
--- a/backend/framework/sdk/src/main/java/io/metersphere/sdk/service/NodeResourcePoolService.java
+++ b/backend/framework/sdk/src/main/java/io/metersphere/sdk/service/NodeResourcePoolService.java
@@ -1,15 +1,21 @@
 package io.metersphere.sdk.service;
 
 
+import com.bastiaanjansen.otp.TOTPGenerator;
+import io.metersphere.sdk.constants.MsHttpHeaders;
 import io.metersphere.sdk.controller.handler.ResultHolder;
 import io.metersphere.sdk.dto.TestResourceDTO;
 import io.metersphere.sdk.dto.TestResourceNodeDTO;
 import io.metersphere.sdk.exception.MSException;
 import io.metersphere.sdk.util.LogUtils;
 import io.metersphere.sdk.util.Translator;
+import jakarta.annotation.Resource;
 import org.apache.commons.collections.CollectionUtils;
 import org.apache.commons.lang3.StringUtils;
 import org.apache.commons.lang3.tuple.ImmutablePair;
+import org.springframework.http.HttpEntity;
+import org.springframework.http.HttpHeaders;
+import org.springframework.http.HttpMethod;
 import org.springframework.http.ResponseEntity;
 import org.springframework.http.client.HttpComponentsClientHttpRequestFactory;
 import org.springframework.stereotype.Service;
@@ -21,10 +27,13 @@ import java.util.List;
 @Service
 @Transactional(rollbackFor = Exception.class)
 public class NodeResourcePoolService {
+    @Resource
+    private TOTPGenerator totpGenerator;
 
     private final static String nodeControllerUrl = "http://%s:%s/status";
 
     private static final RestTemplate restTemplateWithTimeOut = new RestTemplate();
+
     static {
         HttpComponentsClientHttpRequestFactory httpRequestFactory = new HttpComponentsClientHttpRequestFactory();
         httpRequestFactory.setConnectionRequestTimeout(2000);
@@ -55,7 +64,7 @@ public class NodeResourcePoolService {
             }
             isValid = validateNode(testResourceNodeDTO);
             if (!isValid) {
-                 break;
+                break;
             }
         }
         //校验节点
@@ -73,7 +82,11 @@ public class NodeResourcePoolService {
 
     private boolean validateNode(TestResourceNodeDTO node) {
         try {
-            ResponseEntity<ResultHolder> entity = restTemplateWithTimeOut.getForEntity(String.format(nodeControllerUrl, node.getIp(), node.getPort()), ResultHolder.class);
+            String token = totpGenerator.now();
+            HttpHeaders headers = new HttpHeaders();
+            headers.add(MsHttpHeaders.OTP_TOKEN, token);
+            HttpEntity<String> httpEntity = new HttpEntity<>(headers);
+            ResponseEntity<ResultHolder> entity = restTemplateWithTimeOut.exchange(String.format(nodeControllerUrl, node.getIp(), node.getPort()), HttpMethod.GET, httpEntity, ResultHolder.class);
             ResultHolder body = entity.getBody();
             if (body == null) {
                 return false;