From e986ae8241e7390479425a874cb95438aa57f566 Mon Sep 17 00:00:00 2001
From: ReaJason <reajason1225@gmail.com>
Date: Wed, 27 Aug 2025 12:48:00 +0800
Subject: [PATCH] refactor: simplify CommandProbe

---
 generator/build.gradle.kts                    |  1 +
 .../javaweb/probe/payload/CommandProbe.java   |  9 ++-----
 .../probe/payload/CommandProbeTest.java       | 24 +++++++++++++++++++
 3 files changed, 27 insertions(+), 7 deletions(-)
 create mode 100644 generator/src/test/java/com/reajason/javaweb/probe/payload/CommandProbeTest.java

diff --git a/generator/build.gradle.kts b/generator/build.gradle.kts
index 23d9a221..3543233c 100644
--- a/generator/build.gradle.kts
+++ b/generator/build.gradle.kts
@@ -41,6 +41,7 @@ dependencies {
     implementation(libs.bcel)
     implementation(libs.jackson.databind)
     testImplementation(libs.junit.jupiter)
+    testImplementation(libs.hamcrest)
     testRuntimeOnly(libs.junit.platform.launcher)
     testImplementation(libs.bundles.mockito)
 }
\ No newline at end of file
diff --git a/generator/src/main/java/com/reajason/javaweb/probe/payload/CommandProbe.java b/generator/src/main/java/com/reajason/javaweb/probe/payload/CommandProbe.java
index 4a1d5135..f6d3e7bc 100644
--- a/generator/src/main/java/com/reajason/javaweb/probe/payload/CommandProbe.java
+++ b/generator/src/main/java/com/reajason/javaweb/probe/payload/CommandProbe.java
@@ -3,7 +3,6 @@ package com.reajason.javaweb.probe.payload;
 import lombok.SneakyThrows;
 import net.bytebuddy.asm.Advice;
 
-import java.util.NoSuchElementException;
 import java.util.Scanner;
 
 /**
@@ -20,12 +19,8 @@ public class CommandProbe {
     @Advice.OnMethodExit
     public static String exit(@Advice.Argument(0) String data, @Advice.Return(readOnly = false) String ret) throws Exception {
         String[] cmd = System.getProperty("os.name").toLowerCase().contains("window") ? new String[]{"cmd.exe", "/c", data} : new String[]{"/bin/sh", "-c", data};
-        Process process = new ProcessBuilder(cmd).start();
-        try {
-            return ret = new Scanner(process.getInputStream()).useDelimiter("\\A").next();
-        } catch (NoSuchElementException e) {
-            return ret = new Scanner(process.getErrorStream()).useDelimiter("\\A").next();
-        }
+        Process process = new ProcessBuilder(cmd).redirectErrorStream(true).start();
+        return ret = new Scanner(process.getInputStream()).useDelimiter("\\A").next();
     }
 
     @Override
diff --git a/generator/src/test/java/com/reajason/javaweb/probe/payload/CommandProbeTest.java b/generator/src/test/java/com/reajason/javaweb/probe/payload/CommandProbeTest.java
new file mode 100644
index 00000000..f846292f
--- /dev/null
+++ b/generator/src/test/java/com/reajason/javaweb/probe/payload/CommandProbeTest.java
@@ -0,0 +1,24 @@
+package com.reajason.javaweb.probe.payload;
+
+import org.junit.jupiter.api.Disabled;
+import org.junit.jupiter.api.Test;
+
+import static org.hamcrest.CoreMatchers.anyOf;
+import static org.hamcrest.CoreMatchers.containsString;
+import static org.hamcrest.MatcherAssert.assertThat;
+
+/**
+ * @author ReaJason
+ * @since 2025/8/26
+ */
+class CommandProbeTest {
+
+    @Test
+    @Disabled
+    void test() {
+        String result = new CommandProbe("hello").toString();
+        assertThat(result, anyOf(
+                containsString("not found")
+        ));
+    }
+}
\ No newline at end of file