Skip to content

Commit e5aa232

Browse files
committed
feat: support command probe template
1 parent 6bbf275 commit e5aa232

File tree

10 files changed

+99
-11
lines changed

10 files changed

+99
-11
lines changed

boot/src/main/java/com/reajason/javaweb/boot/dto/ProbeShellGenerateRequest.java

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ static class ProbeContentConfigDTO {
2121
private String server;
2222
private String sleepServer;
2323
private String reqParamName;
24+
private String commandTemplate;
2425
}
2526

2627
public ProbeContentConfig parseProbeContentConfig() {
@@ -34,6 +35,7 @@ public ProbeContentConfig parseProbeContentConfig() {
3435
.build();
3536
case ResponseBody -> ResponseBodyConfig.builder()
3637
.reqParamName(probeContentConfig.reqParamName)
38+
.commandTemplate(probeContentConfig.commandTemplate)
3739
.server(probeContentConfig.server)
3840
.build();
3941
default -> throw new UnsupportedOperationException("unknown probe method: " + probeConfig.getProbeMethod());

generator/src/main/java/com/reajason/javaweb/probe/config/ResponseBodyConfig.java

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,4 +23,9 @@ public class ResponseBodyConfig extends ProbeContentConfig {
2323
* 内置执行类加载的字节码
2424
*/
2525
private String base64Bytes;
26+
27+
/**
28+
* 命令执行模板,例如 sh -c "{command}" 2>&1,使用 {command} 作为占位符
29+
*/
30+
private String commandTemplate;
2631
}

generator/src/main/java/com/reajason/javaweb/probe/generator/response/ResponseBodyGenerator.java

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -48,11 +48,14 @@ protected DynamicType.Builder<?> build(ByteBuddy buddy) {
4848
DynamicType.Builder<?> builder = buddy.redefine(writerClass)
4949
.name(probeConfig.getShellClassName())
5050
.visit(new TargetJreVersionVisitorWrapper(probeConfig.getTargetJreVersion()))
51-
.visit(Advice.to(runnerClass).on(named("run")));
51+
.visit(Advice.withCustomMapping()
52+
.bind(ValueAnnotation.class, probeContentConfig.getCommandTemplate())
53+
.to(runnerClass)
54+
.on(named("run")));
5255
if (StringUtils.isNotBlank(probeContentConfig.getReqParamName())) {
5356
builder = builder.visit(MethodCallReplaceVisitorWrapper.newInstance("getDataFromReq",
5457
probeConfig.getShellClassName(), ShellCommonUtil.class.getName()))
55-
.visit(Advice.withCustomMapping().bind(NameAnnotation.class, name)
58+
.visit(Advice.withCustomMapping().bind(ValueAnnotation.class, name)
5659
.to(getDataFromReqInterceptor).on(named("getDataFromReq")));
5760
} else if (ProbeContent.Bytecode.equals(probeConfig.getProbeContent())) {
5861
builder = builder.method(named("getDataFromReq")).intercept(FixedValue.value(probeContentConfig.getBase64Bytes()));
@@ -106,7 +109,7 @@ private Class<?> getWriterClass() {
106109
static class getDataFromReqInterceptor {
107110
@Advice.OnMethodExit
108111
public static void enter(@Advice.Argument(value = 0) Object request,
109-
@NameAnnotation String name,
112+
@ValueAnnotation String name,
110113
@Advice.Return(readOnly = false) String ret) throws Exception {
111114
try {
112115
String p = (String) ShellCommonUtil.invokeMethod(request, "getParameter", new Class[]{String.class}, new Object[]{name});
@@ -123,7 +126,7 @@ public static void enter(@Advice.Argument(value = 0) Object request,
123126
static class getDataFromReqJettyInterceptor {
124127
@Advice.OnMethodExit
125128
public static void enter(@Advice.Argument(value = 0) Object request,
126-
@NameAnnotation String name,
129+
@ValueAnnotation String name,
127130
@Advice.Return(readOnly = false) String ret) throws Exception {
128131
try {
129132
String p = (String) ShellCommonUtil.invokeMethod(request, "getParameter", new Class[]{String.class}, new Object[]{name});
@@ -144,7 +147,7 @@ public static void enter(@Advice.Argument(value = 0) Object request,
144147
}
145148

146149
@Retention(RetentionPolicy.RUNTIME)
147-
public @interface NameAnnotation {
150+
public @interface ValueAnnotation {
148151
}
149152
}
150153

generator/src/main/java/com/reajason/javaweb/probe/payload/CommandProbe.java

Lines changed: 23 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package com.reajason.javaweb.probe.payload;
22

3+
import com.reajason.javaweb.probe.generator.response.ResponseBodyGenerator;
34
import lombok.SneakyThrows;
45
import net.bytebuddy.asm.Advice;
56

@@ -17,15 +18,33 @@ public CommandProbe(String command) {
1718
}
1819

1920
@Advice.OnMethodExit
20-
public static String exit(@Advice.Argument(0) String data, @Advice.Return(readOnly = false) String ret) throws Exception {
21-
String[] cmd = System.getProperty("os.name").toLowerCase().contains("window") ? new String[]{"cmd.exe", "/c", data} : new String[]{"/bin/sh", "-c", data};
22-
Process process = new ProcessBuilder(cmd).redirectErrorStream(true).start();
21+
public static String exit(@Advice.Argument(0) String data,
22+
@Advice.Return(readOnly = false) String ret,
23+
@ResponseBodyGenerator.ValueAnnotation String template
24+
) throws Exception {
25+
String[] cmdarray = null;
26+
String t = template;
27+
if (t == null) {
28+
cmdarray = System.getProperty("os.name").toLowerCase().contains("window") ? new String[]{"cmd.exe", "/c", data} : new String[]{"/bin/sh", "-c", data};
29+
} else {
30+
if (t.contains("\"{command}\"")) {
31+
String[] split = t.split("\\s+");
32+
for (int i = 0; i < split.length; i++) {
33+
split[i] = split[i].replace("\"{command}\"", data);
34+
}
35+
cmdarray = split;
36+
} else {
37+
String cmdline = t.replace("{command}", data);
38+
cmdarray = cmdline.split("\\s+");
39+
}
40+
}
41+
Process process = new ProcessBuilder(cmdarray).redirectErrorStream(true).start();
2342
return ret = new Scanner(process.getInputStream()).useDelimiter("\\A").next();
2443
}
2544

2645
@Override
2746
@SneakyThrows
2847
public String toString() {
29-
return CommandProbe.exit(command, super.toString());
48+
return CommandProbe.exit(command, super.toString(), null);
3049
}
3150
}

web/app/components/probeshell/main-config-card.tsx

Lines changed: 55 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -156,6 +156,51 @@ export default function MainConfigCard({ form, servers }: MainConfigCardProps) {
156156
[form.control, t],
157157
);
158158

159+
const CommandTemplateField = useMemo(
160+
() => (
161+
<>
162+
<div className="space-y-2 pt-4 border-t mt-4">
163+
<FormField
164+
control={form.control}
165+
name="reqParamName"
166+
render={({ field }) => (
167+
<FormFieldItem>
168+
<FormFieldLabel>{t("common:paramName")}</FormFieldLabel>
169+
<FormControl>
170+
<Input placeholder={t("placeholders.input")} {...field} />
171+
</FormControl>
172+
<FormMessage />
173+
</FormFieldItem>
174+
)}
175+
/>
176+
</div>
177+
<div className="space-y-2">
178+
<FormField
179+
control={form.control}
180+
name="commandTemplate"
181+
render={({ field }) => (
182+
<FormFieldItem>
183+
<FormFieldLabel>
184+
{t("common:commandTemplate")} {t("common:optional")}
185+
</FormFieldLabel>
186+
<FormControl>
187+
<Input
188+
{...field}
189+
placeholder={t("common:commandTemplate.placeholder")}
190+
/>
191+
</FormControl>
192+
<p className="text-xs text-muted-foreground mt-1">
193+
{t("common:commandTemplate.description")}
194+
</p>
195+
</FormFieldItem>
196+
)}
197+
/>
198+
</div>
199+
</>
200+
),
201+
[form.control, t],
202+
);
203+
159204
const SleepFields = useMemo(
160205
() => (
161206
<div className="space-y-2 pt-4 border-t mt-4">
@@ -216,6 +261,9 @@ export default function MainConfigCard({ form, servers }: MainConfigCardProps) {
216261
const isServerContent = watchedProbeContent === "Server";
217262

218263
if (isBodyMethod && needParam) {
264+
if (watchedProbeContent === "Command") {
265+
return CommandTemplateField;
266+
}
219267
return RequestParamField;
220268
}
221269

@@ -224,7 +272,13 @@ export default function MainConfigCard({ form, servers }: MainConfigCardProps) {
224272
}
225273

226274
return null;
227-
}, [watchedProbeMethod, watchedProbeContent, RequestParamField, SleepFields]);
275+
}, [
276+
watchedProbeMethod,
277+
watchedProbeContent,
278+
RequestParamField,
279+
SleepFields,
280+
CommandTemplateField,
281+
]);
228282

229283
const DNSLogSection = useMemo(
230284
() => (

web/app/components/ui/select.tsx

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@ function SelectContent({
5252
className,
5353
children,
5454
position = "popper",
55+
align = "center",
5556
...props
5657
}: React.ComponentProps<typeof SelectPrimitive.Content>) {
5758
return (
@@ -65,6 +66,7 @@ function SelectContent({
6566
className,
6667
)}
6768
position={position}
69+
align={align}
6870
{...props}
6971
>
7072
<SelectScrollUpButton />

web/app/components/ui/tabs.tsx

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,7 @@ function TabsTrigger({
4040
<TabsPrimitive.Trigger
4141
data-slot="tabs-trigger"
4242
className={cn(
43-
"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",
43+
"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",
4444
className,
4545
)}
4646
{...props}

web/app/types/probeshell.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ export interface ProbeContentConfig {
2626
sleepServer?: string;
2727
server?: string;
2828
reqParamName?: string;
29+
commandTemplate?: string;
2930
}
3031

3132
export interface DNSLogConfig {

web/app/types/schema.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -160,6 +160,7 @@ export const probeShellFormSchema = yup.object().shape({
160160
host: yup.string().optional(),
161161
server: yup.string().optional(),
162162
reqParamName: yup.string().optional(),
163+
commandTemplate: yup.string().optional(),
163164
seconds: yup.number().optional(),
164165
sleepServer: yup.string().optional(),
165166
packingMethod: yup.string().required(),

web/app/utils/transformer.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,7 @@ export function transformToProbePostData(formValue: ProbeShellFormSchema) {
6767
sleepServer: formValue.sleepServer,
6868
server: formValue.server,
6969
reqParamName: formValue.reqParamName,
70+
commandTemplate: formValue.commandTemplate,
7071
};
7172

7273
return {

0 commit comments

Comments
 (0)