From 5abaa5bbc1d08e46ed2a4da6d2e96a1da03b512a Mon Sep 17 00:00:00 2001 From: ReaJason Date: Mon, 24 Mar 2025 16:20:11 +0800 Subject: [PATCH] feat: support custom shell generator (#49) --- .../boot/controller/ConfigController.java | 15 ++-- .../javaweb/boot/dto/GenerateRequest.java | 9 +- .../reajason/javaweb/boot/entity/Config.java | 2 +- .../reajason/javaweb/ClassBytesShrink.java | 2 +- .../javaweb/memshell/MemShellGenerator.java | 32 ++++--- .../reajason/javaweb/memshell/ShellTool.java | 5 ++ .../javaweb/memshell/config/CustomConfig.java | 20 +++++ .../generator/CustomShellGenerator.java | 53 ++++++++++++ .../generator/CustomShellGeneratorTest.java | 32 +++++++ web/src/components/main-config-card.tsx | 32 +++++-- web/src/components/results/basic-info.tsx | 41 ++++----- web/src/components/tools/custom-tab.tsx | 85 +++++++++++++++++++ web/src/i18n/en.json | 12 ++- web/src/i18n/zh-CN.json | 18 ++-- web/src/routes/index.tsx | 13 ++- web/src/types/schema.ts | 1 + web/src/types/shell.ts | 8 +- web/src/utils/transformer.ts | 1 + 18 files changed, 313 insertions(+), 68 deletions(-) create mode 100644 generator/src/main/java/com/reajason/javaweb/memshell/config/CustomConfig.java create mode 100644 generator/src/main/java/com/reajason/javaweb/memshell/generator/CustomShellGenerator.java create mode 100644 generator/src/test/java/com/reajason/javaweb/memshell/generator/CustomShellGeneratorTest.java create mode 100644 web/src/components/tools/custom-tab.tsx diff --git a/boot/src/main/java/com/reajason/javaweb/boot/controller/ConfigController.java b/boot/src/main/java/com/reajason/javaweb/boot/controller/ConfigController.java index 389298fb..f6194a06 100644 --- a/boot/src/main/java/com/reajason/javaweb/boot/controller/ConfigController.java +++ b/boot/src/main/java/com/reajason/javaweb/boot/controller/ConfigController.java @@ -11,7 +11,6 @@ import org.springframework.web.bind.annotation.RequestMapping; import org.springframework.web.bind.annotation.RestController; import java.util.*; -import java.util.stream.Collectors; /** * @author ReaJason @@ -40,12 +39,14 @@ public class ConfigController { coreMap.put(value.name(), map); } Config config = new Config(); - config.setServers( - Arrays.stream(Server.values()) - .filter(s -> s.getShell() != null) - .map(Server::name) - .collect(Collectors.toList()) - ); + Map> servers = new LinkedHashMap<>(); + for (Server server : Server.values()) { + if (server.getShell() != null) { + Set supportedShellTypes = server.getShell().getShellInjectorMapping().getSupportedShellTypes(); + servers.put(server.name(), supportedShellTypes.stream().toList()); + } + } + config.setServers(servers); config.setCore(coreMap); config.setPackers( Arrays.stream(Packers.values()) diff --git a/boot/src/main/java/com/reajason/javaweb/boot/dto/GenerateRequest.java b/boot/src/main/java/com/reajason/javaweb/boot/dto/GenerateRequest.java index 887a930b..5684cbcb 100644 --- a/boot/src/main/java/com/reajason/javaweb/boot/dto/GenerateRequest.java +++ b/boot/src/main/java/com/reajason/javaweb/boot/dto/GenerateRequest.java @@ -1,7 +1,7 @@ package com.reajason.javaweb.boot.dto; -import com.reajason.javaweb.memshell.config.*; import com.reajason.javaweb.memshell.Packers; +import com.reajason.javaweb.memshell.config.*; import lombok.Data; /** @@ -25,6 +25,7 @@ public class GenerateRequest { private String antSwordPass; private String headerName; private String headerValue; + private String shellClassBase64; } public ShellToolConfig parseShellToolConfig() { @@ -51,7 +52,7 @@ public class GenerateRequest { .headerName(shellToolConfig.getHeaderName()) .headerValue(shellToolConfig.getHeaderValue()) .build(); - case AntSword -> AntSwordConfig.builder() + case AntSword -> AntSwordConfig.builder() .shellClassName(shellToolConfig.getShellClassName()) .pass(shellToolConfig.getAntSwordPass()) .headerName(shellToolConfig.getHeaderName()) @@ -62,6 +63,10 @@ public class GenerateRequest { .headerName(shellToolConfig.getHeaderName()) .headerValue(shellToolConfig.getHeaderValue()) .build(); + case Custom -> CustomConfig.builder() + .shellClassBase64(shellToolConfig.getShellClassBase64()) + .shellClassName(shellToolConfig.getShellClassName()) + .build(); default -> throw new UnsupportedOperationException("unknown shell tool " + shellConfig.getShellTool()); }; } diff --git a/boot/src/main/java/com/reajason/javaweb/boot/entity/Config.java b/boot/src/main/java/com/reajason/javaweb/boot/entity/Config.java index 6f91466a..9e9eb059 100644 --- a/boot/src/main/java/com/reajason/javaweb/boot/entity/Config.java +++ b/boot/src/main/java/com/reajason/javaweb/boot/entity/Config.java @@ -11,7 +11,7 @@ import java.util.Map; */ @Data public class Config { - private List servers; + private Map> servers; private Map> core; private List packers; } \ No newline at end of file diff --git a/common/src/main/java/com/reajason/javaweb/ClassBytesShrink.java b/common/src/main/java/com/reajason/javaweb/ClassBytesShrink.java index 8f394090..c17545f9 100644 --- a/common/src/main/java/com/reajason/javaweb/ClassBytesShrink.java +++ b/common/src/main/java/com/reajason/javaweb/ClassBytesShrink.java @@ -13,7 +13,7 @@ public class ClassBytesShrink { public static byte[] shrink(byte[] bytes, boolean full) { ClassReader cr = new ClassReader(bytes); - ClassWriter cw = new ClassWriter(0); + ClassWriter cw = new ClassWriter(ClassWriter.COMPUTE_MAXS); ClassVisitor cv = new ClassVisitor(Opcodes.ASM9, cw) { @Override public void visitSource(String source, String debug) { diff --git a/generator/src/main/java/com/reajason/javaweb/memshell/MemShellGenerator.java b/generator/src/main/java/com/reajason/javaweb/memshell/MemShellGenerator.java index 519c2a94..d5d2e56e 100644 --- a/generator/src/main/java/com/reajason/javaweb/memshell/MemShellGenerator.java +++ b/generator/src/main/java/com/reajason/javaweb/memshell/MemShellGenerator.java @@ -17,8 +17,9 @@ public class MemShellGenerator { Server server = shellConfig.getServer(); AbstractShell shell = server.getShell(); if (shell == null) { - throw new IllegalArgumentException("Unsupported server"); + throw new IllegalArgumentException("Unsupported server: " + server); } + if (StringUtils.isBlank(shellToolConfig.getShellClassName())) { shellToolConfig.setShellClassName(CommonUtil.generateShellClassName(server, shellConfig.getShellType())); } @@ -27,22 +28,25 @@ public class MemShellGenerator { injectorConfig.setInjectorClassName(CommonUtil.generateInjectorClassName()); } - Pair, Class> shellInjectorPair = shellConfig.getServer().getShell().getShellInjectorPair(shellConfig.getShellTool(), shellConfig.getShellType()); - if (shellInjectorPair == null) { - throw new UnsupportedOperationException("Unknown shell type: " + shellConfig.getShellType()); - } - Class shellClass = shellInjectorPair.getLeft(); - Class injectorClass = shellInjectorPair.getRight(); + Class injectorClass = null; - shellToolConfig.setShellClass(shellClass); + if (ShellTool.Custom.equals(shellConfig.getShellTool())) { + injectorClass = shellConfig.getServer().getShell().getShellInjectorMapping().getInjector(shellConfig.getShellType()); + } else { + Pair, Class> shellInjectorPair = shellConfig.getServer().getShell().getShellInjectorPair(shellConfig.getShellTool(), shellConfig.getShellType()); + if (shellInjectorPair == null) { + throw new UnsupportedOperationException(server + " unsupported shell type: " + shellConfig.getShellType() + " for tool: " + shellConfig.getShellTool()); + } + Class shellClass = shellInjectorPair.getLeft(); + injectorClass = shellInjectorPair.getRight(); + shellToolConfig.setShellClass(shellClass); + } byte[] shellBytes = generateShellBytes(shellConfig, shellToolConfig); - injectorConfig = injectorConfig - .toBuilder() - .injectorClass(injectorClass) - .shellClassName(shellToolConfig.getShellClassName()) - .shellClassBytes(shellBytes).build(); + injectorConfig.setInjectorClass(injectorClass); + injectorConfig.setShellClassName(shellToolConfig.getShellClassName()); + injectorConfig.setShellClassBytes(shellBytes); byte[] injectorBytes = new InjectorGenerator(shellConfig, injectorConfig).generate(); @@ -71,6 +75,8 @@ public class MemShellGenerator { return new AntSwordGenerator(shellConfig, (AntSwordConfig) shellToolConfig).getBytes(); case NeoreGeorg: return new NeoreGeorgGenerator(shellConfig, (NeoreGeorgConfig) shellToolConfig).getBytes(); + case Custom: + return new CustomShellGenerator(shellConfig, (CustomConfig) shellToolConfig).getBytes(); default: throw new UnsupportedOperationException("Unknown shell tool: " + shellConfig.getShellTool()); } diff --git a/generator/src/main/java/com/reajason/javaweb/memshell/ShellTool.java b/generator/src/main/java/com/reajason/javaweb/memshell/ShellTool.java index 7de18049..fc84d3b1 100644 --- a/generator/src/main/java/com/reajason/javaweb/memshell/ShellTool.java +++ b/generator/src/main/java/com/reajason/javaweb/memshell/ShellTool.java @@ -35,5 +35,10 @@ public enum ShellTool { */ NeoreGeorg, + /** + * 自定义 + */ + Custom, + ; } diff --git a/generator/src/main/java/com/reajason/javaweb/memshell/config/CustomConfig.java b/generator/src/main/java/com/reajason/javaweb/memshell/config/CustomConfig.java new file mode 100644 index 00000000..916644ae --- /dev/null +++ b/generator/src/main/java/com/reajason/javaweb/memshell/config/CustomConfig.java @@ -0,0 +1,20 @@ +package com.reajason.javaweb.memshell.config; + +import lombok.AllArgsConstructor; +import lombok.Getter; +import lombok.NoArgsConstructor; +import lombok.ToString; +import lombok.experimental.SuperBuilder; + +/** + * @author ReaJason + * @since 2025/2/12 + */ +@Getter +@SuperBuilder +@NoArgsConstructor +@AllArgsConstructor +@ToString +public class CustomConfig extends ShellToolConfig { + private String shellClassBase64; +} \ No newline at end of file diff --git a/generator/src/main/java/com/reajason/javaweb/memshell/generator/CustomShellGenerator.java b/generator/src/main/java/com/reajason/javaweb/memshell/generator/CustomShellGenerator.java new file mode 100644 index 00000000..74be7ff3 --- /dev/null +++ b/generator/src/main/java/com/reajason/javaweb/memshell/generator/CustomShellGenerator.java @@ -0,0 +1,53 @@ +package com.reajason.javaweb.memshell.generator; + +import com.reajason.javaweb.ClassBytesShrink; +import com.reajason.javaweb.memshell.config.CustomConfig; +import com.reajason.javaweb.memshell.config.ShellConfig; +import net.bytebuddy.jar.asm.ClassReader; +import net.bytebuddy.jar.asm.ClassWriter; +import net.bytebuddy.jar.asm.commons.ClassRemapper; +import net.bytebuddy.jar.asm.commons.SimpleRemapper; +import org.apache.commons.codec.binary.Base64; +import org.apache.commons.lang3.StringUtils; + +/** + * @author ReaJason + * @since 2025/3/18 + */ +public class CustomShellGenerator { + + private final ShellConfig shellConfig; + private final CustomConfig customConfig; + + public CustomShellGenerator(ShellConfig shellConfig, CustomConfig customConfig) { + this.shellConfig = shellConfig; + this.customConfig = customConfig; + } + + public byte[] getBytes() { + String shellClassBase64 = customConfig.getShellClassBase64(); + + if (StringUtils.isBlank(shellClassBase64)) { + throw new IllegalArgumentException("Custom shell class is empty"); + } + + byte[] bytes = renameClass(Base64.decodeBase64(shellClassBase64), customConfig.getShellClassName()); + + return ClassBytesShrink.shrink(bytes, shellConfig.isShrink()); + } + + private static byte[] renameClass(byte[] classBytes, String newName) { + ClassReader reader = null; + try { + reader = new ClassReader(classBytes); + } catch (Exception e) { + throw new RuntimeException("invalid class bytes"); + } + String oldClassName = reader.getClassName(); + String newClassName = newName.replace('.', '/'); + ClassWriter writer = new ClassWriter(reader, ClassWriter.COMPUTE_MAXS | ClassWriter.COMPUTE_FRAMES); + ClassRemapper adapter = new ClassRemapper(writer, new SimpleRemapper(oldClassName, newClassName)); + reader.accept(adapter, 0); + return writer.toByteArray(); + } +} diff --git a/generator/src/test/java/com/reajason/javaweb/memshell/generator/CustomShellGeneratorTest.java b/generator/src/test/java/com/reajason/javaweb/memshell/generator/CustomShellGeneratorTest.java new file mode 100644 index 00000000..570441f3 --- /dev/null +++ b/generator/src/test/java/com/reajason/javaweb/memshell/generator/CustomShellGeneratorTest.java @@ -0,0 +1,32 @@ +package com.reajason.javaweb.memshell.generator; + +import com.reajason.javaweb.memshell.config.CustomConfig; +import com.reajason.javaweb.memshell.config.ShellConfig; +import com.reajason.javaweb.memshell.utils.CommonUtil; +import lombok.SneakyThrows; +import net.bytebuddy.ByteBuddy; +import net.bytebuddy.jar.asm.ClassReader; +import org.apache.commons.codec.binary.Base64; +import org.junit.jupiter.api.Test; + +import static org.junit.jupiter.api.Assertions.assertEquals; + +/** + * @author ReaJason + * @since 2025/3/19 + */ +class CustomShellGeneratorTest { + + @Test + @SneakyThrows + void test() { + byte[] bytes = new ByteBuddy() + .subclass(Object.class) + .name(CommonUtil.generateShellClassName()).make().getBytes(); + String className = CommonUtil.generateShellClassName(); + byte[] bytes1 = new CustomShellGenerator(ShellConfig.builder().build(), CustomConfig.builder().shellClassName(className).shellClassBase64(Base64.encodeBase64String(bytes)).build()).getBytes(); + + ClassReader classReader = new ClassReader(bytes1); + assertEquals(className, classReader.getClassName().replace("/", ".")); + } +} \ No newline at end of file diff --git a/web/src/components/main-config-card.tsx b/web/src/components/main-config-card.tsx index 890b75e0..9f5580c6 100644 --- a/web/src/components/main-config-card.tsx +++ b/web/src/components/main-config-card.tsx @@ -4,7 +4,7 @@ import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@ import { Switch } from "@/components/ui/switch.tsx"; import { Tabs, TabsList, TabsTrigger } from "@/components/ui/tabs"; import { FormSchema } from "@/types/schema.ts"; -import { JDKVersion, MainConfig, ShellToolType } from "@/types/shell.ts"; +import { JDKVersion, MainConfig, ServerConfig, ShellToolType } from "@/types/shell.ts"; import { JreTip } from "@/components/tips/jre-tip.tsx"; import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card.tsx"; @@ -17,6 +17,7 @@ import { ShieldOffIcon, SwordIcon, WaypointsIcon, + ZapIcon, } from "lucide-react"; import { JSX, useState } from "react"; import { FormProvider, UseFormReturn } from "react-hook-form"; @@ -24,6 +25,7 @@ import { useTranslation } from "react-i18next"; import { AntSwordTabContent } from "./tools/antsword-tab"; import { BehinderTabContent } from "./tools/behinder-tab"; import { CommandTabContent } from "./tools/command-tab"; +import CustomTabContent from "./tools/custom-tab"; import { GodzillaTabContent } from "./tools/godzilla-tab"; import { NeoRegTabContent } from "./tools/neoreg-tab"; import { Suo5TabContent } from "./tools/suo5-tab"; @@ -35,17 +37,18 @@ const shellToolIcons: Record = { [ShellToolType.AntSword]: , [ShellToolType.Suo5]: , [ShellToolType.NeoreGeorg]: , + [ShellToolType.Custom]: , }; export function MainConfigCard({ mainConfig, form, servers, -}: { +}: Readonly<{ mainConfig: MainConfig | undefined; form: UseFormReturn; - servers?: string[]; -}) { + servers?: ServerConfig; +}>) { const [shellToolMap, setShellToolMap] = useState<{ [toolName: string]: string[]; }>(); @@ -56,6 +59,7 @@ export function MainConfigCard({ ShellToolType.Command, ShellToolType.Suo5, ShellToolType.NeoreGeorg, + ShellToolType.Custom, ]); const [shellTypes, setShellTypes] = useState([]); const shellTool = form.watch("shellTool"); @@ -66,7 +70,7 @@ export function MainConfigCard({ const newShellToolMap = mainConfig[value]; setShellToolMap(newShellToolMap); const newShellTools = Object.keys(newShellToolMap); - setShellTools(newShellTools.map((tool) => tool as ShellToolType)); + setShellTools([...newShellTools.map((tool) => tool as ShellToolType), ShellToolType.Custom]); if (newShellTools.length > 0) { const firstTool = newShellTools[0]; setShellTypes(newShellToolMap[firstTool]); @@ -124,8 +128,17 @@ export function MainConfigCard({ form.resetField("headerValue"); }; + const resetCustom = () => { + form.resetField("shellClassBase64"); + }; + if (shellToolMap) { - setShellTypes(shellToolMap[value]); + if (value === ShellToolType.Custom) { + setShellTypes(servers?.[form.getValues("server")] as string[]); + } else { + setShellTypes(shellToolMap[value]); + } + form.resetField("urlPattern"); form.resetField("shellType"); form.resetField("shellClassName"); @@ -142,6 +155,8 @@ export function MainConfigCard({ resetAntSword(); } else if (value === ShellToolType.NeoreGeorg) { resetNeoreGeorg(); + } else if (value === ShellToolType.Custom) { + resetCustom(); } } form.setValue("shellTool", value); @@ -177,7 +192,7 @@ export function MainConfigCard({ - {servers?.map((server: string) => ( + {Object.keys(servers ?? {}).map((server: string) => ( {server} @@ -291,7 +306,7 @@ export function MainConfigCard({ className="flex-1 min-w-24 data-[state=active]:bg-background" > - {shellToolIcons[shellTool as ShellToolType]} + {shellToolIcons[shellTool]} {shellTool} @@ -305,6 +320,7 @@ export function MainConfigCard({ + ); diff --git a/web/src/components/results/basic-info.tsx b/web/src/components/results/basic-info.tsx index f2c79fd6..7748e00b 100644 --- a/web/src/components/results/basic-info.tsx +++ b/web/src/components/results/basic-info.tsx @@ -10,13 +10,12 @@ import { } from "@/types/shell"; import { FileTextIcon } from "lucide-react"; import { useTranslation } from "react-i18next"; -import { Fragment } from "react/jsx-runtime"; import { CopyableField } from "../copyable-field"; import { Card, CardContent, CardHeader, CardTitle } from "../ui/card"; import { Separator } from "../ui/separator"; import { FeedbackAlert } from "./feedback-alert"; -export function BasicInfo({ generateResult }: { generateResult?: GenerateResult }) { +export function BasicInfo({ generateResult }: Readonly<{ generateResult?: GenerateResult }>) { const { t } = useTranslation(); return ( @@ -40,7 +39,7 @@ export function BasicInfo({ generateResult }: { generateResult?: GenerateResult value={generateResult?.injectorConfig.urlPattern} /> - + {generateResult?.shellConfig.shellTool !== ShellToolType.Custom && }
{generateResult?.shellConfig.shellTool === ShellToolType.Behinder && ( <> @@ -62,7 +61,7 @@ export function BasicInfo({ generateResult }: { generateResult?: GenerateResult )} {generateResult?.shellConfig.shellTool === ShellToolType.Godzilla && ( - + <> - + )} {generateResult?.shellConfig.shellTool === ShellToolType.Command && ( - - - + )} {generateResult?.shellConfig.shellTool === ShellToolType.Suo5 && ( - - - + )} {generateResult?.shellConfig.shellTool === ShellToolType.AntSword && ( - + <> - + )} {generateResult?.shellConfig.shellTool === ShellToolType.NeoreGeorg && ( - + <> - + )}
diff --git a/web/src/components/tools/custom-tab.tsx b/web/src/components/tools/custom-tab.tsx new file mode 100644 index 00000000..9c6c076f --- /dev/null +++ b/web/src/components/tools/custom-tab.tsx @@ -0,0 +1,85 @@ +import { Card, CardContent } from "@/components/ui/card"; +import { TabsContent } from "@/components/ui/tabs"; +import { FormSchema } from "@/types/schema"; +import { t } from "i18next"; +import { useState } from "react"; +import { FormProvider, UseFormReturn } from "react-hook-form"; +import { FormControl, FormField, FormItem, FormLabel } from "../ui/form"; +import { Input } from "../ui/input"; +import { Label } from "../ui/label"; +import { RadioGroup, RadioGroupItem } from "../ui/radio-group"; +import { Textarea } from "../ui/textarea"; +import { OptionalClassFormField } from "./classname-field"; +import { ShellTypeFormField } from "./shelltype-field"; +import { UrlPatternFormField } from "./urlpattern-field"; + +export default function CustomTabContent({ + form, + shellTypes, +}: Readonly<{ form: UseFormReturn; shellTypes: Array }>) { + const [isFile, setIsFile] = useState(false); + + return ( + + + + +
+ + +
+ ( + + {t("shellToolConfig.base64String")} + { + field.onChange(""); + setIsFile(value === "file"); + }} + className="flex items-center space-x-2" + > +
+ + +
+
+ + +
+
+ + {isFile ? ( + { + const file = e.target.files?.[0]; + if (file) { + const reader = new FileReader(); + reader.onload = (event) => { + const base64String = (event.target?.result as string)?.split(",")[1] || ""; + field.onChange(base64String); + }; + reader.readAsDataURL(file); + } + }} + accept=".class" + placeholder={t("placeholders.input")} + type="file" + /> + ) : ( +