From 056032cae70d889522175225473dfee0dc2cc784 Mon Sep 17 00:00:00 2001
From: ReaJason <reajason1225@gmail.com>
Date: Sun, 7 Dec 2025 17:13:52 +0800
Subject: [PATCH] feat: support command probe template

---
 .../boot/dto/ProbeShellGenerateRequest.java   |  2 +
 .../probe/config/ResponseBodyConfig.java      |  5 ++
 .../response/ResponseBodyGenerator.java       | 13 +++--
 .../javaweb/probe/payload/CommandProbe.java   | 27 +++++++--
 .../probeshell/main-config-card.tsx           | 56 ++++++++++++++++++-
 web/app/components/ui/select.tsx              |  2 +
 web/app/components/ui/tabs.tsx                |  2 +-
 web/app/types/probeshell.ts                   |  1 +
 web/app/types/schema.ts                       |  1 +
 web/app/utils/transformer.ts                  |  1 +
 10 files changed, 99 insertions(+), 11 deletions(-)

diff --git a/boot/src/main/java/com/reajason/javaweb/boot/dto/ProbeShellGenerateRequest.java b/boot/src/main/java/com/reajason/javaweb/boot/dto/ProbeShellGenerateRequest.java
index 78cca343..c8c8cdf3 100644
--- a/boot/src/main/java/com/reajason/javaweb/boot/dto/ProbeShellGenerateRequest.java
+++ b/boot/src/main/java/com/reajason/javaweb/boot/dto/ProbeShellGenerateRequest.java
@@ -21,6 +21,7 @@ public class ProbeShellGenerateRequest {
         private String server;
         private String sleepServer;
         private String reqParamName;
+        private String commandTemplate;
     }
 
     public ProbeContentConfig parseProbeContentConfig() {
@@ -34,6 +35,7 @@ public class ProbeShellGenerateRequest {
                     .build();
             case ResponseBody -> ResponseBodyConfig.builder()
                     .reqParamName(probeContentConfig.reqParamName)
+                    .commandTemplate(probeContentConfig.commandTemplate)
                     .server(probeContentConfig.server)
                     .build();
             default -> throw new UnsupportedOperationException("unknown probe method: " + probeConfig.getProbeMethod());
diff --git a/generator/src/main/java/com/reajason/javaweb/probe/config/ResponseBodyConfig.java b/generator/src/main/java/com/reajason/javaweb/probe/config/ResponseBodyConfig.java
index 8f8b8a8a..1ea76754 100644
--- a/generator/src/main/java/com/reajason/javaweb/probe/config/ResponseBodyConfig.java
+++ b/generator/src/main/java/com/reajason/javaweb/probe/config/ResponseBodyConfig.java
@@ -23,4 +23,9 @@ public class ResponseBodyConfig extends ProbeContentConfig {
      * 内置执行类加载的字节码
      */
     private String base64Bytes;
+
+    /**
+     * 命令执行模板，例如 sh -c "{command}" 2>&1，使用 {command} 作为占位符
+     */
+    private String commandTemplate;
 }
diff --git a/generator/src/main/java/com/reajason/javaweb/probe/generator/response/ResponseBodyGenerator.java b/generator/src/main/java/com/reajason/javaweb/probe/generator/response/ResponseBodyGenerator.java
index f09ffa07..63aedbd1 100644
--- a/generator/src/main/java/com/reajason/javaweb/probe/generator/response/ResponseBodyGenerator.java
+++ b/generator/src/main/java/com/reajason/javaweb/probe/generator/response/ResponseBodyGenerator.java
@@ -48,11 +48,14 @@ public class ResponseBodyGenerator extends ByteBuddyShellGenerator<ResponseBodyC
         DynamicType.Builder<?> builder = buddy.redefine(writerClass)
                 .name(probeConfig.getShellClassName())
                 .visit(new TargetJreVersionVisitorWrapper(probeConfig.getTargetJreVersion()))
-                .visit(Advice.to(runnerClass).on(named("run")));
+                .visit(Advice.withCustomMapping()
+                                .bind(ValueAnnotation.class, probeContentConfig.getCommandTemplate())
+                                .to(runnerClass)
+                                .on(named("run")));
         if (StringUtils.isNotBlank(probeContentConfig.getReqParamName())) {
             builder = builder.visit(MethodCallReplaceVisitorWrapper.newInstance("getDataFromReq",
                             probeConfig.getShellClassName(), ShellCommonUtil.class.getName()))
-                    .visit(Advice.withCustomMapping().bind(NameAnnotation.class, name)
+                    .visit(Advice.withCustomMapping().bind(ValueAnnotation.class, name)
                             .to(getDataFromReqInterceptor).on(named("getDataFromReq")));
         } else if (ProbeContent.Bytecode.equals(probeConfig.getProbeContent())) {
             builder = builder.method(named("getDataFromReq")).intercept(FixedValue.value(probeContentConfig.getBase64Bytes()));
@@ -106,7 +109,7 @@ public class ResponseBodyGenerator extends ByteBuddyShellGenerator<ResponseBodyC
     static class getDataFromReqInterceptor {
         @Advice.OnMethodExit
         public static void enter(@Advice.Argument(value = 0) Object request,
-                                 @NameAnnotation String name,
+                                 @ValueAnnotation String name,
                                  @Advice.Return(readOnly = false) String ret) throws Exception {
             try {
                 String p = (String) ShellCommonUtil.invokeMethod(request, "getParameter", new Class[]{String.class}, new Object[]{name});
@@ -123,7 +126,7 @@ public class ResponseBodyGenerator extends ByteBuddyShellGenerator<ResponseBodyC
     static class getDataFromReqJettyInterceptor {
         @Advice.OnMethodExit
         public static void enter(@Advice.Argument(value = 0) Object request,
-                                 @NameAnnotation String name,
+                                 @ValueAnnotation String name,
                                  @Advice.Return(readOnly = false) String ret) throws Exception {
             try {
                 String p = (String) ShellCommonUtil.invokeMethod(request, "getParameter", new Class[]{String.class}, new Object[]{name});
@@ -144,7 +147,7 @@ public class ResponseBodyGenerator extends ByteBuddyShellGenerator<ResponseBodyC
     }
 
     @Retention(RetentionPolicy.RUNTIME)
-    public @interface NameAnnotation {
+    public @interface ValueAnnotation {
     }
 }
 
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 f6d3e7bc..9de0c033 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
@@ -1,5 +1,6 @@
 package com.reajason.javaweb.probe.payload;
 
+import com.reajason.javaweb.probe.generator.response.ResponseBodyGenerator;
 import lombok.SneakyThrows;
 import net.bytebuddy.asm.Advice;
 
@@ -17,15 +18,33 @@ 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).redirectErrorStream(true).start();
+    public static String exit(@Advice.Argument(0) String data,
+                              @Advice.Return(readOnly = false) String ret,
+                              @ResponseBodyGenerator.ValueAnnotation String template
+    ) throws Exception {
+        String[] cmdarray = null;
+        String t = template;
+        if (t == null) {
+            cmdarray = System.getProperty("os.name").toLowerCase().contains("window") ? new String[]{"cmd.exe", "/c", data} : new String[]{"/bin/sh", "-c", data};
+        } else {
+            if (t.contains("\"{command}\"")) {
+                String[] split = t.split("\\s+");
+                for (int i = 0; i < split.length; i++) {
+                    split[i] = split[i].replace("\"{command}\"", data);
+                }
+                cmdarray = split;
+            } else {
+                String cmdline = t.replace("{command}", data);
+                cmdarray = cmdline.split("\\s+");
+            }
+        }
+        Process process = new ProcessBuilder(cmdarray).redirectErrorStream(true).start();
         return ret = new Scanner(process.getInputStream()).useDelimiter("\\A").next();
     }
 
     @Override
     @SneakyThrows
     public String toString() {
-        return CommandProbe.exit(command, super.toString());
+        return CommandProbe.exit(command, super.toString(), null);
     }
 }
diff --git a/web/app/components/probeshell/main-config-card.tsx b/web/app/components/probeshell/main-config-card.tsx
index 66627a04..b621e50c 100644
--- a/web/app/components/probeshell/main-config-card.tsx
+++ b/web/app/components/probeshell/main-config-card.tsx
@@ -156,6 +156,51 @@ export default function MainConfigCard({ form, servers }: MainConfigCardProps) {
     [form.control, t],
   );
 
+  const CommandTemplateField = useMemo(
+    () => (
+      <>
+        <div className="space-y-2 pt-4 border-t mt-4">
+          <FormField
+            control={form.control}
+            name="reqParamName"
+            render={({ field }) => (
+              <FormFieldItem>
+                <FormFieldLabel>{t("common:paramName")}</FormFieldLabel>
+                <FormControl>
+                  <Input placeholder={t("placeholders.input")} {...field} />
+                </FormControl>
+                <FormMessage />
+              </FormFieldItem>
+            )}
+          />
+        </div>
+        <div className="space-y-2">
+          <FormField
+            control={form.control}
+            name="commandTemplate"
+            render={({ field }) => (
+              <FormFieldItem>
+                <FormFieldLabel>
+                  {t("common:commandTemplate")} {t("common:optional")}
+                </FormFieldLabel>
+                <FormControl>
+                  <Input
+                    {...field}
+                    placeholder={t("common:commandTemplate.placeholder")}
+                  />
+                </FormControl>
+                <p className="text-xs text-muted-foreground mt-1">
+                  {t("common:commandTemplate.description")}
+                </p>
+              </FormFieldItem>
+            )}
+          />
+        </div>
+      </>
+    ),
+    [form.control, t],
+  );
+
   const SleepFields = useMemo(
     () => (
       <div className="space-y-2 pt-4 border-t mt-4">
@@ -216,6 +261,9 @@ export default function MainConfigCard({ form, servers }: MainConfigCardProps) {
     const isServerContent = watchedProbeContent === "Server";
 
     if (isBodyMethod && needParam) {
+      if (watchedProbeContent === "Command") {
+        return CommandTemplateField;
+      }
       return RequestParamField;
     }
 
@@ -224,7 +272,13 @@ export default function MainConfigCard({ form, servers }: MainConfigCardProps) {
     }
 
     return null;
-  }, [watchedProbeMethod, watchedProbeContent, RequestParamField, SleepFields]);
+  }, [
+    watchedProbeMethod,
+    watchedProbeContent,
+    RequestParamField,
+    SleepFields,
+    CommandTemplateField,
+  ]);
 
   const DNSLogSection = useMemo(
     () => (
diff --git a/web/app/components/ui/select.tsx b/web/app/components/ui/select.tsx
index 53f1a16d..e0e6a4a6 100644
--- a/web/app/components/ui/select.tsx
+++ b/web/app/components/ui/select.tsx
@@ -52,6 +52,7 @@ function SelectContent({
   className,
   children,
   position = "popper",
+  align = "center",
   ...props
 }: React.ComponentProps<typeof SelectPrimitive.Content>) {
   return (
@@ -65,6 +66,7 @@ function SelectContent({
           className,
         )}
         position={position}
+        align={align}
         {...props}
       >
         <SelectScrollUpButton />
diff --git a/web/app/components/ui/tabs.tsx b/web/app/components/ui/tabs.tsx
index 14679113..f6ac0d38 100644
--- a/web/app/components/ui/tabs.tsx
+++ b/web/app/components/ui/tabs.tsx
@@ -40,7 +40,7 @@ function TabsTrigger({
     <TabsPrimitive.Trigger
       data-slot="tabs-trigger"
       className={cn(
-        "dark:data-[state=active]:text-foreground focus-visible:border-ring focus-visible:ring-ring/50 focus-visible:outline-ring dark:data-[state=active]:border-input dark:data-[state=active]:bg-input/30 text-foreground dark:text-muted-foreground inline-flex h-[calc(100%-1px)] flex-1 items-center justify-center gap-1.5 rounded-md border border-transparent px-2 py-1 text-sm font-medium whitespace-nowrap transition-[color,box-shadow] focus-visible:ring-[3px] focus-visible:outline-1 disabled:pointer-events-none disabled:opacity-50 data-[state=active]:shadow-sm [&_svg]:pointer-events-none [&_svg]:shrink-0 [&_svg:not([class*='size-'])]:size-4",
+        "data-[state=active]:bg-background dark:data-[state=active]:text-foreground focus-visible:border-ring focus-visible:ring-ring/50 focus-visible:outline-ring dark:data-[state=active]:border-input dark:data-[state=active]:bg-input/30 text-foreground dark:text-muted-foreground inline-flex h-[calc(100%-1px)] flex-1 items-center justify-center gap-1.5 rounded-md border border-transparent px-2 py-1 text-sm font-medium whitespace-nowrap transition-[color,box-shadow] focus-visible:ring-[3px] focus-visible:outline-1 disabled:pointer-events-none disabled:opacity-50 data-[state=active]:shadow-sm [&_svg]:pointer-events-none [&_svg]:shrink-0 [&_svg:not([class*='size-'])]:size-4",
         className,
       )}
       {...props}
diff --git a/web/app/types/probeshell.ts b/web/app/types/probeshell.ts
index f45e99b0..1b0a650e 100644
--- a/web/app/types/probeshell.ts
+++ b/web/app/types/probeshell.ts
@@ -26,6 +26,7 @@ export interface ProbeContentConfig {
   sleepServer?: string;
   server?: string;
   reqParamName?: string;
+  commandTemplate?: string;
 }
 
 export interface DNSLogConfig {
diff --git a/web/app/types/schema.ts b/web/app/types/schema.ts
index 0b155359..58534212 100644
--- a/web/app/types/schema.ts
+++ b/web/app/types/schema.ts
@@ -160,6 +160,7 @@ export const probeShellFormSchema = yup.object().shape({
   host: yup.string().optional(),
   server: yup.string().optional(),
   reqParamName: yup.string().optional(),
+  commandTemplate: yup.string().optional(),
   seconds: yup.number().optional(),
   sleepServer: yup.string().optional(),
   packingMethod: yup.string().required(),
diff --git a/web/app/utils/transformer.ts b/web/app/utils/transformer.ts
index a4d7835f..0016a603 100644
--- a/web/app/utils/transformer.ts
+++ b/web/app/utils/transformer.ts
@@ -67,6 +67,7 @@ export function transformToProbePostData(formValue: ProbeShellFormSchema) {
     sleepServer: formValue.sleepServer,
     server: formValue.server,
     reqParamName: formValue.reqParamName,
+    commandTemplate: formValue.commandTemplate,
   };
 
   return {