Merge pull request #337 from Nizernizer/main

lostsnow · web-flow · commit 4f04ad88b6cc · 2022-08-01T15:37:05.000+08:00
agent state
diff --git a/dongtai-agent/src/main/java/io/dongtai/iast/agent/Agent.java b/dongtai-agent/src/main/java/io/dongtai/iast/agent/Agent.java
@@ -2,7 +2,6 @@
 
 import java.io.*;
 import java.util.Arrays;
-import java.util.concurrent.TimeUnit;
 
 import io.dongtai.iast.agent.util.FileUtils;
 import org.apache.commons.cli.*;
@@ -152,7 +151,6 @@ private static void extractJattach() throws IOException {
      * @param args
      */
     public static void main(String[] args) throws InterruptedException {
-        TimeUnit.SECONDS.sleep(10);
         String[] agentArgs = new String[0];
         try {
             agentArgs = parseAgentArgs(args);
diff --git a/dongtai-agent/src/main/java/io/dongtai/iast/agent/Constant.java b/dongtai-agent/src/main/java/io/dongtai/iast/agent/Constant.java
@@ -9,7 +9,9 @@ public class Constant {
 
     public final static String API_AGENT_REGISTER = "/api/v1/agent/register";
     public final static String API_REPORT_UPLOAD = "/api/v1/report/upload";
-    public final static String API_ENGINE_ACTION = "/api/v1/engine/action";
+    public final static String API_AGENT_STATUS = "/api/v1/actual_action";
+//    public final static String API_ENGINE_ACTION = "/api/v1/engine/action";
+    public final static String API_ENGINE_ACTION = "/api/v1/except_action";
 
     public static final String KEY_NETWORK = "network";
     public static final String KEY_AGENT_VERSION = "version";
@@ -39,6 +41,8 @@ public class Constant {
     public static final String KEY_PERFORMANCE = "performance";
     public static final String KEY_CORE_INSTALLED = "isCoreInstalled";
     public static final String KEY_CORE_RUNNING = "isCoreRunning";
+    public static final String KEY_CORE_STATUS = "actualRunningStatus";
+    public static final String KEY_CORE_STATE = "stateStatus";
     public static final String KEY_RETURN_QUEUE = "returnQueue";
 
     // 线程异常上报
diff --git a/dongtai-agent/src/main/java/io/dongtai/iast/agent/manager/EngineManager.java b/dongtai-agent/src/main/java/io/dongtai/iast/agent/manager/EngineManager.java
@@ -34,6 +34,7 @@ public class EngineManager {
     private static final String FALLBACK_MANAGER_CLASS = "io.dongtai.iast.core.bytecode.enhance.plugin.fallback.FallbackManager";
     private static final String REMOTE_CONFIG_UTILS_CLASS = "io.dongtai.iast.core.utils.config.RemoteConfigUtils";
     private static final String ENGINE_MANAGER_CLASS = "io.dongtai.iast.core.EngineManager";
+    private static final String ENGINE_FALLBACK_CLASS = "io.dongtai.iast.core.bytecode.enhance.plugin.fallback.FallbackSwitch";
     private static final String INJECT_PACKAGE_REMOTE_URI = "/api/v1/engine/download?engineName=dongtai-spy";
     private static final String INJECT_PACKAGE_REMOTE_URI_JDK6 = "/api/v1/engine/download?engineName=dongtai-spy-jdk6";
     private static final String ENGINE_PACKAGE_REMOTE_URI = "/api/v1/engine/download?engineName=dongtai-core";
@@ -124,6 +125,22 @@ public static boolean checkCoreIsRunning() {
         }
     }
 
+    public static boolean checkCoreIsFallback() {
+        if (IAST_CLASS_LOADER == null) {
+            return false;
+        }
+        try {
+            final Class<?> engineManagerClass = IAST_CLASS_LOADER.loadClass(ENGINE_FALLBACK_CLASS);
+            if (engineManagerClass == null) {
+                return false;
+            }
+            return (Boolean) engineManagerClass.getMethod("isEngineFallback").invoke(null);
+        } catch (Throwable e) {
+            DongTaiLog.info("checkCoreIsRunning failed, msg:{}, cause:{}", e.getMessage(), e.getCause());
+            return false;
+        }
+    }
+
     /**
      * 在核心包中加载并获取远端配置工具类
      */
diff --git a/dongtai-agent/src/main/java/io/dongtai/iast/agent/monitor/impl/EngineMonitor.java b/dongtai-agent/src/main/java/io/dongtai/iast/agent/monitor/impl/EngineMonitor.java
@@ -5,6 +5,7 @@
 import io.dongtai.iast.agent.monitor.MonitorDaemonThread;
 import io.dongtai.iast.agent.monitor.ServerCommandEnum;
 import io.dongtai.iast.agent.report.AgentRegisterReport;
+import io.dongtai.iast.agent.report.HeartBeatReport;
 import io.dongtai.iast.agent.util.ThreadUtils;
 import io.dongtai.iast.agent.util.http.HttpClientUtils;
 import io.dongtai.iast.agent.Constant;
@@ -20,6 +21,7 @@ public class EngineMonitor implements IMonitor {
     public static Boolean isCoreRegisterStart = false;
     private static final String NAME = "EngineMonitor";
     private static Boolean isUninstallHeart = true;
+    private String coreStatus = "1";
 
     public EngineMonitor(EngineManager engineManager) {
         this.engineManager = engineManager;
@@ -34,7 +36,21 @@ public String getName() {
     @Override
     public void check() throws Exception {
         String status = checkForStatus();
-        ServerCommandEnum serviceCmdEnum = ServerCommandEnum.getEnum(status);
+        if ("1".equals(status) && !"1".equals(coreStatus)) {
+            coreStatus = "1";
+            DongTaiLog.info("engine start");
+            if (EngineManager.checkCoreIsInstalled()) {
+                engineManager.start();
+            } else {
+                startEngine();
+            }
+        } else if ("2".equals(status) && !"2".equals(coreStatus)) {
+            coreStatus = "2";
+            DongTaiLog.info("engine stop");
+            engineManager.stop();
+        }
+        HttpClientUtils.sendPost(Constant.API_AGENT_STATUS, HeartBeatReport.generateAgentStatusMsg());
+/*        ServerCommandEnum serviceCmdEnum = ServerCommandEnum.getEnum(status);
         if (serviceCmdEnum == null || serviceCmdEnum == ServerCommandEnum.NO_CMD) {
             return;
         }
@@ -70,7 +86,7 @@ public void check() throws Exception {
                 forceSwitchPerformanceBreaker(false);
                 break;
             default:
-        }
+        }*/
     }
 
     private void forceSwitchPerformanceBreaker(boolean turnOn) {
@@ -91,7 +107,8 @@ private String checkForStatus() {
             String respRaw = String.valueOf(HttpClientUtils.sendGet(Constant.API_ENGINE_ACTION, "agentId", String.valueOf(AgentRegisterReport.getAgentFlag())));
             if (respRaw != null && !respRaw.isEmpty()) {
                 JSONObject resp = new JSONObject(respRaw);
-                return resp.get("data").toString();
+                JSONObject data = (JSONObject) resp.get("data");
+                return data.get("exceptRunningStatus").toString();
             }
         } catch (Exception e) {
             return "other";
@@ -101,12 +118,12 @@ private String checkForStatus() {
 
     public void startEngine() {
         boolean status = true;
-        if(couldInstallEngine()){
+        if (couldInstallEngine()) {
             // jdk8以上
             status = status && engineManager.extractPackage();
             status = status && engineManager.install();
             status = status && engineManager.start();
-        }else {
+        } else {
             // jdk6-7
             status = status && engineManager.extractPackage();
             status = status && engineManager.install();
diff --git a/dongtai-agent/src/main/java/io/dongtai/iast/agent/report/HeartBeatReport.java b/dongtai-agent/src/main/java/io/dongtai/iast/agent/report/HeartBeatReport.java
@@ -39,6 +39,20 @@ public static String generateHeartBeatMsg() {
         return report.toString();
     }
 
+    public static String generateAgentStatusMsg() {
+        JSONObject detail = new JSONObject();
+        detail.put(Constant.KEY_AGENT_ID, AgentRegisterReport.getAgentFlag());
+        Integer status = 0;
+        if (EngineManager.checkCoreIsRunning()){
+            status = 1;
+        }else {
+            status = 2;
+        }
+        detail.put(Constant.KEY_CORE_STATUS, status);
+        detail.put(Constant.KEY_CORE_STATE, EngineManager.checkCoreIsFallback() ? 3 : status);
+        return detail.toString();
+    }
+
     /**
      * Query CPU usage
      *
diff --git a/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/models/IastReplayModel.java b/dongtai-core/src/main/java/io/dongtai/iast/core/handler/hookpoint/models/IastReplayModel.java
@@ -73,7 +73,11 @@ public String getFullUrl() {
 
         StringBuilder url = new StringBuilder();
         // fixme 根据协议，判断使用http/https
-        url.append("http://");
+        if (EngineManager.SERVER.getProtocol().toLowerCase().contains("https")){
+            url.append("https://");
+        }else {
+            url.append("http://");
+        }
         url.append(host);
         if (getRequestQueryString().isEmpty()) {
             url.append(requestUrl);
