[feature] Optimize script protocol and add Windows-script yml template (

#2367) Signed-off-by: Zhang Yuxuan <[email protected]> Co-authored-by: tomsun28 <[email protected]>
apache · Jul 24, 2024 · 348b50c · 348b50c
1 parent 8853889
commit 348b50c
Show file tree

Hide file tree

Showing 4 changed files with 383 additions and 16 deletions.
diff --git a/collector/src/main/java/org/apache/hertzbeat/collector/collect/script/ScriptCollectImpl.java b/collector/src/main/java/org/apache/hertzbeat/collector/collect/script/ScriptCollectImpl.java
@@ -46,10 +46,13 @@
 public class ScriptCollectImpl extends AbstractCollect {
     public static final String WINDOWS_SCRIPT = "windows_script";
     public static final String LINUX_SCRIPT = "linux_script";
-    private static final String CMD = "cmd.exe";
+    private static final String CMD = "cmd";
     private static final String BASH = "bash";
+    private static final String POWERSHELL = "powershell";
     private static final String CMD_C = "/c";
     private static final String BASH_C = "-c";
+    private static final String POWERSHELL_C = "-Command";
+    private static final String POWERSHELL_FILE = "-File";
     private static final String PARSE_TYPE_ONE_ROW = "oneRow";
     private static final String PARSE_TYPE_MULTI_ROW = "multiRow";
     private static final String PARSE_TYPE_NETCAT = "netcat";
@@ -62,6 +65,7 @@ public void preCheck(Metrics metrics) throws IllegalArgumentException {
         Assert.notNull(scriptProtocol, "Script collect must has Imap params");
         Assert.notNull(scriptProtocol.getCharset(), "Script charset is required");
         Assert.notNull(scriptProtocol.getParseType(), "Script parse type is required");
+        Assert.notNull(scriptProtocol.getScriptTool(), "Script tool is required");
         if (!(StringUtils.hasText(scriptProtocol.getScriptCommand()) || StringUtils.hasText(scriptProtocol.getScriptPath()))) {
             throw new IllegalArgumentException("At least one script command or script path is required.");
         }
@@ -72,35 +76,50 @@ public void collect(CollectRep.MetricsData.Builder builder, long monitorId, Stri
         ScriptProtocol scriptProtocol = metrics.getScript();
         long startTime = System.currentTimeMillis();
         ProcessBuilder processBuilder;
-        if (WINDOWS_SCRIPT.equalsIgnoreCase(app)) {
-            if (StringUtils.hasText(scriptProtocol.getScriptCommand())) {
-                processBuilder = new ProcessBuilder(CMD, CMD_C, scriptProtocol.getScriptCommand().trim());
-            } else {
-                processBuilder = new ProcessBuilder(CMD, scriptProtocol.getScriptPath().trim());
+        // use command
+        if (StringUtils.hasText(scriptProtocol.getScriptCommand())) {
+            switch (scriptProtocol.getScriptTool()) {
+                case BASH -> processBuilder = new ProcessBuilder(BASH, BASH_C, scriptProtocol.getScriptCommand().trim());
+                case CMD -> processBuilder = new ProcessBuilder(CMD, CMD_C, scriptProtocol.getScriptCommand().trim());
+                case POWERSHELL -> processBuilder = new ProcessBuilder("powershell.exe", POWERSHELL_C, scriptProtocol.getScriptCommand().trim());
+                default -> {
+                    builder.setCode(CollectRep.Code.FAIL);
+                    builder.setMsg("Not support script tool:" + scriptProtocol.getScriptTool());
+                    return;
+                }
             }
-
-        } else if (LINUX_SCRIPT.equalsIgnoreCase(app)) {
-            if (StringUtils.hasText(scriptProtocol.getScriptCommand())) {
-                processBuilder = new ProcessBuilder(BASH, BASH_C, scriptProtocol.getScriptCommand().trim());
-            } else {
-                processBuilder = new ProcessBuilder(BASH, scriptProtocol.getScriptPath().trim());
+        // use command file
+        } else if (StringUtils.hasText((scriptProtocol.getScriptPath()))) {
+            switch (scriptProtocol.getScriptTool()) {
+                case BASH -> processBuilder = new ProcessBuilder(BASH, scriptProtocol.getScriptPath().trim());
+                case CMD -> processBuilder = new ProcessBuilder(CMD,  scriptProtocol.getScriptPath().trim());
+                case POWERSHELL -> processBuilder = new ProcessBuilder(POWERSHELL, POWERSHELL_FILE, scriptProtocol.getScriptPath().trim());
+                default -> {
+                    builder.setCode(CollectRep.Code.FAIL);
+                    builder.setMsg("Not support script tool:" + scriptProtocol.getScriptTool());
+                    return;
+                }
             }
         } else {
             builder.setCode(CollectRep.Code.FAIL);
-            builder.setMsg("Not support OS:" + app);
+            builder.setMsg("At least one script command or script path is required.");
             return;
         }
+        // set work directory
         String workDirectory = scriptProtocol.getWorkDirectory();
         if (StringUtils.hasText(workDirectory)) {
             processBuilder.directory(new File(workDirectory));
         }
+        // execute command
         try {
             Process process = processBuilder.start();
             BufferedReader reader = new BufferedReader(new InputStreamReader(process.getInputStream(), Charset.forName(scriptProtocol.getCharset())));
             StringBuilder response = new StringBuilder();
             String line;
             while ((line = reader.readLine()) != null) {
-                response.append(line).append("\n");
+                if (StringUtils.hasText(line)) {
+                    response.append(line).append("\n");
+                }
             }
             process.waitFor();
             Long responseTime = System.currentTimeMillis() - startTime;
@@ -224,13 +243,13 @@ private void parseResponseDataByMulti(String result, List<String> aliasFields,
             log.error("ssh response data only has header: {}", result);
             return;
         }
-        String[] fields = lines[0].split(" ");
+        String[] fields = lines[0].split("\\s+");
         Map<String, Integer> fieldMapping = new HashMap<>(fields.length);
         for (int i = 0; i < fields.length; i++) {
             fieldMapping.put(fields[i].trim().toLowerCase(), i);
         }
         for (int i = 1; i < lines.length; i++) {
-            String[] values = lines[i].split(" ");
+            String[] values = lines[i].split("\\s+");
             CollectRep.ValueRow.Builder valueRowBuilder = CollectRep.ValueRow.newBuilder();
             for (String alias : aliasFields) {
                 if (CollectorConstants.RESPONSE_TIME.equalsIgnoreCase(alias)) {

diff --git a/common/src/main/java/org/apache/hertzbeat/common/entity/job/protocol/ScriptProtocol.java b/common/src/main/java/org/apache/hertzbeat/common/entity/job/protocol/ScriptProtocol.java
@@ -55,4 +55,9 @@ public class ScriptProtocol {
      * Response data parsing mode：oneRow, multiRow
      */
     private String parseType;
+
+    /**
+     * Script tool name, exp: bash, cmd, powershell
+     */
+    private String scriptTool;
 }
diff --git a/manager/src/main/resources/define/app-linux_script.yml b/manager/src/main/resources/define/app-linux_script.yml
@@ -74,6 +74,8 @@ metrics:
     protocol: script
     # the config content when protocol is ssh
     script:
+      # script tool
+      scriptTool: bash
       # OS charset
       charset: UTF-8
       # script working directory
@@ -139,6 +141,7 @@ metrics:
       - usage=100-idle
     protocol: script
     script:
+      scriptTool: bash
       charset: UTF-8
       scriptCommand:
         |
@@ -202,6 +205,7 @@ metrics:
       - usage=(used / total) * 100
     protocol: script
     script:
+      scriptTool: bash
       charset: UTF-8
       scriptCommand: free -m | awk 'BEGIN{print "total used free buff_cache available"} NR==2{print $2,$3,$4,$6,$7}'
       parseType: multiRow
@@ -240,6 +244,7 @@ metrics:
           en-US: Write Rate
     protocol: script
     script:
+      scriptTool: bash
       charset: UTF-8
       scriptCommand: vmstat -D | awk 'NR==1{print $1}';vmstat -D | awk 'NR==2{print $1}';vmstat 1 1 | awk 'NR==3{print $10}';vmstat 1 1 | awk 'NR==3{print $9}';vmstat 1 1 | awk 'NR==3{print $16}'
       parseType: oneRow
@@ -273,6 +278,7 @@ metrics:
       - transmit_bytes=B->MB
     protocol: script
     script:
+      scriptTool: bash
       charset: UTF-8
       scriptCommand: cat /proc/net/dev | tail -n +3 | awk 'BEGIN{ print "interface_name receive_bytes transmit_bytes"} {print $1,$2,$10}'
       parseType: multiRow
@@ -314,6 +320,7 @@ metrics:
           en-US: Mounted
     protocol: script
     script:
+      scriptTool: bash
       charset: UTF-8
       scriptCommand: df -mP | tail -n +2 | awk 'BEGIN{ print "filesystem used available usage mounted"} {print $1,$3,$4,$5,$6}'
       parseType: multiRow
@@ -349,6 +356,7 @@ metrics:
           en-US: Command
     protocol: script
     script:
+      scriptTool: bash
       charset: UTF-8
       scriptCommand: ps aux | sort -k3nr | awk 'BEGIN{ print "pid cpu_usage mem_usage command" } {printf "%s %s %s ", $2, $3, $4; for (i=11; i<=NF; i++) { printf "%s", $i; if (i < NF) printf " "; } print ""}' | head -n 11
       parseType: multiRow
@@ -384,6 +392,7 @@ metrics:
           en-US: Command
     protocol: script
     script:
+      scriptTool: bash
       charset: UTF-8
       scriptCommand: ps aux | sort -k4nr | awk 'BEGIN{ print "pid cpu_usage mem_usage command" } {printf "%s %s %s ", $2, $3, $4; for (i=11; i<=NF; i++) { printf "%s", $i; if (i < NF) printf " "; } print ""}' | head -n 11
       parseType: multiRow