Skip to content

Commit 241df2c

Browse files
ReaJasonReaJason
committed
feat: support custom shell generator (#49)
1 parent efb5ae4 commit 241df2c

File tree

18 files changed

+313
-68
lines changed

18 files changed

+313
-68
lines changed

‎boot/src/main/java/com/reajason/javaweb/boot/controller/ConfigController.java‎

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,6 @@
1111
import org.springframework.web.bind.annotation.RestController;
1212

1313
import java.util.*;
14-
import java.util.stream.Collectors;
1514

1615
/**
1716
* @author ReaJason
@@ -40,12 +39,14 @@ public ResponseEntity<?> config() {
4039
coreMap.put(value.name(), map);
4140
}
4241
Config config = new Config();
43-
config.setServers(
44-
Arrays.stream(Server.values())
45-
.filter(s -> s.getShell() != null)
46-
.map(Server::name)
47-
.collect(Collectors.toList())
48-
);
42+
Map<String, List<String>> servers = new LinkedHashMap<>();
43+
for (Server server : Server.values()) {
44+
if (server.getShell() != null) {
45+
Set<String> supportedShellTypes = server.getShell().getShellInjectorMapping().getSupportedShellTypes();
46+
servers.put(server.name(), supportedShellTypes.stream().toList());
47+
}
48+
}
49+
config.setServers(servers);
4950
config.setCore(coreMap);
5051
config.setPackers(
5152
Arrays.stream(Packers.values())

‎boot/src/main/java/com/reajason/javaweb/boot/dto/GenerateRequest.java‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
package com.reajason.javaweb.boot.dto;
22

3-
import com.reajason.javaweb.memshell.config.*;
43
import com.reajason.javaweb.memshell.Packers;
4+
import com.reajason.javaweb.memshell.config.*;
55
import lombok.Data;
66

77
/**
@@ -25,6 +25,7 @@ static class ShellToolConfigDTO {
2525
private String antSwordPass;
2626
private String headerName;
2727
private String headerValue;
28+
private String shellClassBase64;
2829
}
2930

3031
public ShellToolConfig parseShellToolConfig() {
@@ -51,7 +52,7 @@ public ShellToolConfig parseShellToolConfig() {
5152
.headerName(shellToolConfig.getHeaderName())
5253
.headerValue(shellToolConfig.getHeaderValue())
5354
.build();
54-
case AntSword -> AntSwordConfig.builder()
55+
case AntSword -> AntSwordConfig.builder()
5556
.shellClassName(shellToolConfig.getShellClassName())
5657
.pass(shellToolConfig.getAntSwordPass())
5758
.headerName(shellToolConfig.getHeaderName())
@@ -62,6 +63,10 @@ public ShellToolConfig parseShellToolConfig() {
6263
.headerName(shellToolConfig.getHeaderName())
6364
.headerValue(shellToolConfig.getHeaderValue())
6465
.build();
66+
case Custom -> CustomConfig.builder()
67+
.shellClassBase64(shellToolConfig.getShellClassBase64())
68+
.shellClassName(shellToolConfig.getShellClassName())
69+
.build();
6570
default -> throw new UnsupportedOperationException("unknown shell tool " + shellConfig.getShellTool());
6671
};
6772
}

‎boot/src/main/java/com/reajason/javaweb/boot/entity/Config.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@
1111
*/
1212
@Data
1313
public class Config {
14-
private List<String> servers;
14+
private Map<String, List<String>> servers;
1515
private Map<String, Map<?, ?>> core;
1616
private List<String> packers;
1717
}

‎common/src/main/java/com/reajason/javaweb/ClassBytesShrink.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ public class ClassBytesShrink {
1313

1414
public static byte[] shrink(byte[] bytes, boolean full) {
1515
ClassReader cr = new ClassReader(bytes);
16-
ClassWriter cw = new ClassWriter(0);
16+
ClassWriter cw = new ClassWriter(ClassWriter.COMPUTE_MAXS);
1717
ClassVisitor cv = new ClassVisitor(Opcodes.ASM9, cw) {
1818
@Override
1919
public void visitSource(String source, String debug) {

‎generator/src/main/java/com/reajason/javaweb/memshell/MemShellGenerator.java‎

Lines changed: 19 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,9 @@ public static GenerateResult generate(ShellConfig shellConfig, InjectorConfig in
1717
Server server = shellConfig.getServer();
1818
AbstractShell shell = server.getShell();
1919
if (shell == null) {
20-
throw new IllegalArgumentException("Unsupported server");
20+
throw new IllegalArgumentException("Unsupported server: " + server);
2121
}
22+
2223
if (StringUtils.isBlank(shellToolConfig.getShellClassName())) {
2324
shellToolConfig.setShellClassName(CommonUtil.generateShellClassName(server, shellConfig.getShellType()));
2425
}
@@ -27,22 +28,25 @@ public static GenerateResult generate(ShellConfig shellConfig, InjectorConfig in
2728
injectorConfig.setInjectorClassName(CommonUtil.generateInjectorClassName());
2829
}
2930

30-
Pair<Class<?>, Class<?>> shellInjectorPair = shellConfig.getServer().getShell().getShellInjectorPair(shellConfig.getShellTool(), shellConfig.getShellType());
31-
if (shellInjectorPair == null) {
32-
throw new UnsupportedOperationException("Unknown shell type: " + shellConfig.getShellType());
33-
}
34-
Class<?> shellClass = shellInjectorPair.getLeft();
35-
Class<?> injectorClass = shellInjectorPair.getRight();
31+
Class<?> injectorClass = null;
3632

37-
shellToolConfig.setShellClass(shellClass);
33+
if (ShellTool.Custom.equals(shellConfig.getShellTool())) {
34+
injectorClass = shellConfig.getServer().getShell().getShellInjectorMapping().getInjector(shellConfig.getShellType());
35+
} else {
36+
Pair<Class<?>, Class<?>> shellInjectorPair = shellConfig.getServer().getShell().getShellInjectorPair(shellConfig.getShellTool(), shellConfig.getShellType());
37+
if (shellInjectorPair == null) {
38+
throw new UnsupportedOperationException(server + " unsupported shell type: " + shellConfig.getShellType() + " for tool: " + shellConfig.getShellTool());
39+
}
40+
Class<?> shellClass = shellInjectorPair.getLeft();
41+
injectorClass = shellInjectorPair.getRight();
42+
shellToolConfig.setShellClass(shellClass);
43+
}
3844

3945
byte[] shellBytes = generateShellBytes(shellConfig, shellToolConfig);
4046

41-
injectorConfig = injectorConfig
42-
.toBuilder()
43-
.injectorClass(injectorClass)
44-
.shellClassName(shellToolConfig.getShellClassName())
45-
.shellClassBytes(shellBytes).build();
47+
injectorConfig.setInjectorClass(injectorClass);
48+
injectorConfig.setShellClassName(shellToolConfig.getShellClassName());
49+
injectorConfig.setShellClassBytes(shellBytes);
4650

4751
byte[] injectorBytes = new InjectorGenerator(shellConfig, injectorConfig).generate();
4852

@@ -71,6 +75,8 @@ private static byte[] generateShellBytes(ShellConfig shellConfig, ShellToolConfi
7175
return new AntSwordGenerator(shellConfig, (AntSwordConfig) shellToolConfig).getBytes();
7276
case NeoreGeorg:
7377
return new NeoreGeorgGenerator(shellConfig, (NeoreGeorgConfig) shellToolConfig).getBytes();
78+
case Custom:
79+
return new CustomShellGenerator(shellConfig, (CustomConfig) shellToolConfig).getBytes();
7480
default:
7581
throw new UnsupportedOperationException("Unknown shell tool: " + shellConfig.getShellTool());
7682
}

‎generator/src/main/java/com/reajason/javaweb/memshell/ShellTool.java‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,5 +35,10 @@ public enum ShellTool {
3535
*/
3636
NeoreGeorg,
3737

38+
/**
39+
* 自定义
40+
*/
41+
Custom,
42+
3843
;
3944
}

‎generator/src/main/java/com/reajason/javaweb/memshell/config/CustomConfig.java‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
package com.reajason.javaweb.memshell.config;
2+
3+
import lombok.AllArgsConstructor;
4+
import lombok.Getter;
5+
import lombok.NoArgsConstructor;
6+
import lombok.ToString;
7+
import lombok.experimental.SuperBuilder;
8+
9+
/**
10+
* @author ReaJason
11+
* @since 2025/2/12
12+
*/
13+
@Getter
14+
@SuperBuilder
15+
@NoArgsConstructor
16+
@AllArgsConstructor
17+
@ToString
18+
public class CustomConfig extends ShellToolConfig {
19+
private String shellClassBase64;
20+
}

‎generator/src/main/java/com/reajason/javaweb/memshell/generator/CustomShellGenerator.java‎

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,53 @@
1+
package com.reajason.javaweb.memshell.generator;
2+
3+
import com.reajason.javaweb.ClassBytesShrink;
4+
import com.reajason.javaweb.memshell.config.CustomConfig;
5+
import com.reajason.javaweb.memshell.config.ShellConfig;
6+
import net.bytebuddy.jar.asm.ClassReader;
7+
import net.bytebuddy.jar.asm.ClassWriter;
8+
import net.bytebuddy.jar.asm.commons.ClassRemapper;
9+
import net.bytebuddy.jar.asm.commons.SimpleRemapper;
10+
import org.apache.commons.codec.binary.Base64;
11+
import org.apache.commons.lang3.StringUtils;
12+
13+
/**
14+
* @author ReaJason
15+
* @since 2025/3/18
16+
*/
17+
public class CustomShellGenerator {
18+
19+
private final ShellConfig shellConfig;
20+
private final CustomConfig customConfig;
21+
22+
public CustomShellGenerator(ShellConfig shellConfig, CustomConfig customConfig) {
23+
this.shellConfig = shellConfig;
24+
this.customConfig = customConfig;
25+
}
26+
27+
public byte[] getBytes() {
28+
String shellClassBase64 = customConfig.getShellClassBase64();
29+
30+
if (StringUtils.isBlank(shellClassBase64)) {
31+
throw new IllegalArgumentException("Custom shell class is empty");
32+
}
33+
34+
byte[] bytes = renameClass(Base64.decodeBase64(shellClassBase64), customConfig.getShellClassName());
35+
36+
return ClassBytesShrink.shrink(bytes, shellConfig.isShrink());
37+
}
38+
39+
private static byte[] renameClass(byte[] classBytes, String newName) {
40+
ClassReader reader = null;
41+
try {
42+
reader = new ClassReader(classBytes);
43+
} catch (Exception e) {
44+
throw new RuntimeException("invalid class bytes");
45+
}
46+
String oldClassName = reader.getClassName();
47+
String newClassName = newName.replace('.', '/');
48+
ClassWriter writer = new ClassWriter(reader, ClassWriter.COMPUTE_MAXS | ClassWriter.COMPUTE_FRAMES);
49+
ClassRemapper adapter = new ClassRemapper(writer, new SimpleRemapper(oldClassName, newClassName));
50+
reader.accept(adapter, 0);
51+
return writer.toByteArray();
52+
}
53+
}

‎generator/src/test/java/com/reajason/javaweb/memshell/generator/CustomShellGeneratorTest.java‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
package com.reajason.javaweb.memshell.generator;
2+
3+
import com.reajason.javaweb.memshell.config.CustomConfig;
4+
import com.reajason.javaweb.memshell.config.ShellConfig;
5+
import com.reajason.javaweb.memshell.utils.CommonUtil;
6+
import lombok.SneakyThrows;
7+
import net.bytebuddy.ByteBuddy;
8+
import net.bytebuddy.jar.asm.ClassReader;
9+
import org.apache.commons.codec.binary.Base64;
10+
import org.junit.jupiter.api.Test;
11+
12+
import static org.junit.jupiter.api.Assertions.assertEquals;
13+
14+
/**
15+
* @author ReaJason
16+
* @since 2025/3/19
17+
*/
18+
class CustomShellGeneratorTest {
19+
20+
@Test
21+
@SneakyThrows
22+
void test() {
23+
byte[] bytes = new ByteBuddy()
24+
.subclass(Object.class)
25+
.name(CommonUtil.generateShellClassName()).make().getBytes();
26+
String className = CommonUtil.generateShellClassName();
27+
byte[] bytes1 = new CustomShellGenerator(ShellConfig.builder().build(), CustomConfig.builder().shellClassName(className).shellClassBase64(Base64.encodeBase64String(bytes)).build()).getBytes();
28+
29+
ClassReader classReader = new ClassReader(bytes1);
30+
assertEquals(className, classReader.getClassName().replace("/", "."));
31+
}
32+
}

‎web/src/components/main-config-card.tsx‎

Lines changed: 24 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@ import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@
44
import { Switch } from "@/components/ui/switch.tsx";
55
import { Tabs, TabsList, TabsTrigger } from "@/components/ui/tabs";
66
import { FormSchema } from "@/types/schema.ts";
7-
import { JDKVersion, MainConfig, ShellToolType } from "@/types/shell.ts";
7+
import { JDKVersion, MainConfig, ServerConfig, ShellToolType } from "@/types/shell.ts";
88

99
import { JreTip } from "@/components/tips/jre-tip.tsx";
1010
import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card.tsx";
@@ -17,13 +17,15 @@ import {
1717
ShieldOffIcon,
1818
SwordIcon,
1919
WaypointsIcon,
20+
ZapIcon,
2021
} from "lucide-react";
2122
import { JSX, useState } from "react";
2223
import { FormProvider, UseFormReturn } from "react-hook-form";
2324
import { useTranslation } from "react-i18next";
2425
import { AntSwordTabContent } from "./tools/antsword-tab";
2526
import { BehinderTabContent } from "./tools/behinder-tab";
2627
import { CommandTabContent } from "./tools/command-tab";
28+
import CustomTabContent from "./tools/custom-tab";
2729
import { GodzillaTabContent } from "./tools/godzilla-tab";
2830
import { NeoRegTabContent } from "./tools/neoreg-tab";
2931
import { Suo5TabContent } from "./tools/suo5-tab";
@@ -35,17 +37,18 @@ const shellToolIcons: Record<ShellToolType, JSX.Element> = {
3537
[ShellToolType.AntSword]: <SwordIcon className="h-4 w-4" />,
3638
[ShellToolType.Suo5]: <WaypointsIcon className="h-4 w-4" />,
3739
[ShellToolType.NeoreGeorg]: <NetworkIcon className="h-4 w-4" />,
40+
[ShellToolType.Custom]: <ZapIcon className="h-4 w-4" />,
3841
};
3942

4043
export function MainConfigCard({
4144
mainConfig,
4245
form,
4346
servers,
44-
}: {
47+
}: Readonly<{
4548
mainConfig: MainConfig | undefined;
4649
form: UseFormReturn<FormSchema>;
47-
servers?: string[];
48-
}) {
50+
servers?: ServerConfig;
51+
}>) {
4952
const [shellToolMap, setShellToolMap] = useState<{
5053
[toolName: string]: string[];
5154
}>();
@@ -56,6 +59,7 @@ export function MainConfigCard({
5659
ShellToolType.Command,
5760
ShellToolType.Suo5,
5861
ShellToolType.NeoreGeorg,
62+
ShellToolType.Custom,
5963
]);
6064
const [shellTypes, setShellTypes] = useState<string[]>([]);
6165
const shellTool = form.watch("shellTool");
@@ -66,7 +70,7 @@ export function MainConfigCard({
6670
const newShellToolMap = mainConfig[value];
6771
setShellToolMap(newShellToolMap);
6872
const newShellTools = Object.keys(newShellToolMap);
69-
setShellTools(newShellTools.map((tool) => tool as ShellToolType));
73+
setShellTools([...newShellTools.map((tool) => tool as ShellToolType), ShellToolType.Custom]);
7074
if (newShellTools.length > 0) {
7175
const firstTool = newShellTools[0];
7276
setShellTypes(newShellToolMap[firstTool]);
@@ -124,8 +128,17 @@ export function MainConfigCard({
124128
form.resetField("headerValue");
125129
};
126130

131+
const resetCustom = () => {
132+
form.resetField("shellClassBase64");
133+
};
134+
127135
if (shellToolMap) {
128-
setShellTypes(shellToolMap[value]);
136+
if (value === ShellToolType.Custom) {
137+
setShellTypes(servers?.[form.getValues("server")] as string[]);
138+
} else {
139+
setShellTypes(shellToolMap[value]);
140+
}
141+
129142
form.resetField("urlPattern");
130143
form.resetField("shellType");
131144
form.resetField("shellClassName");
@@ -142,6 +155,8 @@ export function MainConfigCard({
142155
resetAntSword();
143156
} else if (value === ShellToolType.NeoreGeorg) {
144157
resetNeoreGeorg();
158+
} else if (value === ShellToolType.Custom) {
159+
resetCustom();
145160
}
146161
}
147162
form.setValue("shellTool", value);
@@ -177,7 +192,7 @@ export function MainConfigCard({
177192
</SelectTrigger>
178193
</FormControl>
179194
<SelectContent>
180-
{servers?.map((server: string) => (
195+
{Object.keys(servers ?? {}).map((server: string) => (
181196
<SelectItem key={server} value={server}>
182197
{server}
183198
</SelectItem>
@@ -291,7 +306,7 @@ export function MainConfigCard({
291306
className="flex-1 min-w-24 data-[state=active]:bg-background"
292307
>
293308
<span className="flex items-center gap-2">
294-
{shellToolIcons[shellTool as ShellToolType]}
309+
{shellToolIcons[shellTool]}
295310
{shellTool}
296311
</span>
297312
</TabsTrigger>
@@ -305,6 +320,7 @@ export function MainConfigCard({
305320
<AntSwordTabContent form={form} shellTypes={shellTypes} />
306321
<Suo5TabContent form={form} shellTypes={shellTypes} />
307322
<NeoRegTabContent form={form} shellTypes={shellTypes} />
323+
<CustomTabContent form={form} shellTypes={shellTypes} />
308324
</Tabs>
309325
</FormProvider>
310326
);

0 commit comments

Comments
 (0)