feat: support custom shell generator (#49)

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, JSX.Element> = {
   [ShellToolType.AntSword]: <SwordIcon className="h-4 w-4" />,
   [ShellToolType.Suo5]: <WaypointsIcon className="h-4 w-4" />,
   [ShellToolType.NeoreGeorg]: <NetworkIcon className="h-4 w-4" />,
+  [ShellToolType.Custom]: <ZapIcon className="h-4 w-4" />,
 };
 
 export function MainConfigCard({
   mainConfig,
   form,
   servers,
-}: {
+}: Readonly<{
   mainConfig: MainConfig | undefined;
   form: UseFormReturn<FormSchema>;
-  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<string[]>([]);
   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({
                       </SelectTrigger>
                     </FormControl>
                     <SelectContent>
-                      {servers?.map((server: string) => (
+                      {Object.keys(servers ?? {}).map((server: string) => (
                         <SelectItem key={server} value={server}>
                           {server}
                         </SelectItem>
@@ -291,7 +306,7 @@ export function MainConfigCard({
                 className="flex-1 min-w-24 data-[state=active]:bg-background"
               >
                 <span className="flex items-center gap-2">
-                  {shellToolIcons[shellTool as ShellToolType]}
+                  {shellToolIcons[shellTool]}
                   {shellTool}
                 </span>
               </TabsTrigger>
@@ -305,6 +320,7 @@ export function MainConfigCard({
         <AntSwordTabContent form={form} shellTypes={shellTypes} />
         <Suo5TabContent form={form} shellTypes={shellTypes} />
         <NeoRegTabContent form={form} shellTypes={shellTypes} />
+        <CustomTabContent form={form} shellTypes={shellTypes} />
       </Tabs>
     </FormProvider>
   );

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 (
     <Card>
@@ -40,7 +39,7 @@ export function BasicInfo({ generateResult }: { generateResult?: GenerateResult
             value={generateResult?.injectorConfig.urlPattern}
           />
         </div>
-        <Separator className="my-2" />
+        {generateResult?.shellConfig.shellTool !== ShellToolType.Custom && <Separator className="my-2" />}
         <div className="grid grid-cols-1 md:grid-cols-2 gap-2">
           {generateResult?.shellConfig.shellTool === ShellToolType.Behinder && (
             <>
@@ -62,7 +61,7 @@ export function BasicInfo({ generateResult }: { generateResult?: GenerateResult
             </>
           )}
           {generateResult?.shellConfig.shellTool === ShellToolType.Godzilla && (
-            <Fragment>
+            <>
               <CopyableField
                 label={t("shellToolConfig.pass")}
                 text={(generateResult?.shellToolConfig as GodzillaShellToolConfig).pass}
@@ -80,28 +79,24 @@ export function BasicInfo({ generateResult }: { generateResult?: GenerateResult
                 text={`${(generateResult?.shellToolConfig as GodzillaShellToolConfig).headerName}: ${(generateResult?.shellToolConfig as GodzillaShellToolConfig).headerValue}`}
                 value={`${(generateResult?.shellToolConfig as GodzillaShellToolConfig).headerName}: ${(generateResult?.shellToolConfig as GodzillaShellToolConfig).headerValue}`}
               />
-            </Fragment>
+            </>
           )}
           {generateResult?.shellConfig.shellTool === ShellToolType.Command && (
-            <Fragment>
-              <CopyableField
-                label={t("shellToolConfig.paramName")}
-                text={(generateResult?.shellToolConfig as CommandShellToolConfig).paramName}
-                value={(generateResult?.shellToolConfig as CommandShellToolConfig).paramName}
-              />
-            </Fragment>
+            <CopyableField
+              label={t("shellToolConfig.paramName")}
+              text={(generateResult?.shellToolConfig as CommandShellToolConfig).paramName}
+              value={(generateResult?.shellToolConfig as CommandShellToolConfig).paramName}
+            />
           )}
           {generateResult?.shellConfig.shellTool === ShellToolType.Suo5 && (
-            <Fragment>
-              <CopyableField
-                label={t("shellToolConfig.suo5Header")}
-                text={`${(generateResult?.shellToolConfig as Suo5ShellToolConfig).headerName}: ${(generateResult?.shellToolConfig as Suo5ShellToolConfig).headerValue}`}
-                value={`${(generateResult?.shellToolConfig as Suo5ShellToolConfig).headerName}: ${(generateResult?.shellToolConfig as Suo5ShellToolConfig).headerValue}`}
-              />
-            </Fragment>
+            <CopyableField
+              label={t("shellToolConfig.suo5Header")}
+              text={`${(generateResult?.shellToolConfig as Suo5ShellToolConfig).headerName}: ${(generateResult?.shellToolConfig as Suo5ShellToolConfig).headerValue}`}
+              value={`${(generateResult?.shellToolConfig as Suo5ShellToolConfig).headerName}: ${(generateResult?.shellToolConfig as Suo5ShellToolConfig).headerValue}`}
+            />
           )}
           {generateResult?.shellConfig.shellTool === ShellToolType.AntSword && (
-            <Fragment>
+            <>
               <CopyableField
                 label={t("shellToolConfig.antSwordPass")}
                 text={(generateResult?.shellToolConfig as AntSwordShellToolConfig).pass}
@@ -112,17 +107,17 @@ export function BasicInfo({ generateResult }: { generateResult?: GenerateResult
                 text={`${(generateResult?.shellToolConfig as AntSwordShellToolConfig).headerName}: ${(generateResult?.shellToolConfig as AntSwordShellToolConfig).headerValue}`}
                 value={`${(generateResult?.shellToolConfig as AntSwordShellToolConfig).headerName}: ${(generateResult?.shellToolConfig as AntSwordShellToolConfig).headerValue}`}
               />
-            </Fragment>
+            </>
           )}
           {generateResult?.shellConfig.shellTool === ShellToolType.NeoreGeorg && (
-            <Fragment>
+            <>
               <CopyableField label={t("shellToolConfig.neoreGeorgKey")} text="key" value="key" />
               <CopyableField
                 label={t("shellToolConfig.neoreGeorgHeader")}
                 text={`${(generateResult?.shellToolConfig as NeoreGeorgShellToolConfig).headerName}: ${(generateResult?.shellToolConfig as NeoreGeorgShellToolConfig).headerValue}`}
                 value={`${(generateResult?.shellToolConfig as NeoreGeorgShellToolConfig).headerName}: ${(generateResult?.shellToolConfig as NeoreGeorgShellToolConfig).headerValue}`}
               />
-            </Fragment>
+            </>
           )}
         </div>
         <Separator className="my-2" />

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<FormSchema>; shellTypes: Array<string> }>) {
+  const [isFile, setIsFile] = useState(false);
+
+  return (
+    <FormProvider {...form}>
+      <TabsContent value="Custom">
+        <Card>
+          <CardContent className="space-y-2 mt-4">
+            <div className="grid grid-cols-2 gap-2">
+              <ShellTypeFormField form={form} shellTypes={shellTypes} />
+              <UrlPatternFormField form={form} />
+            </div>
+            <FormField
+              control={form.control}
+              name="shellClassBase64"
+              render={({ field }) => (
+                <FormItem className="space-y-1">
+                  <FormLabel className="h-6 flex items-center gap-1">{t("shellToolConfig.base64String")}</FormLabel>
+                  <RadioGroup
+                    value={isFile ? "file" : "base64"}
+                    onValueChange={(value) => {
+                      field.onChange("");
+                      setIsFile(value === "file");
+                    }}
+                    className="flex items-center space-x-2"
+                  >
+                    <div className="flex items-center space-x-2">
+                      <RadioGroupItem value="base64" id="option-one" />
+                      <Label htmlFor="option-one">Base64</Label>
+                    </div>
+                    <div className="flex items-center space-x-2">
+                      <RadioGroupItem value="file" id="option-two" />
+                      <Label htmlFor="option-two">File</Label>
+                    </div>
+                  </RadioGroup>
+                  <FormControl className="pt-2 mt-2">
+                    {isFile ? (
+                      <Input
+                        onChange={(e) => {
+                          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"
+                      />
+                    ) : (
+                      <Textarea {...field} placeholder={t("placeholders.input")} className="h-24" />
+                    )}
+                  </FormControl>
+                </FormItem>
+              )}
+            />
+            <OptionalClassFormField form={form} />
+          </CardContent>
+        </Card>
+      </TabsContent>
+    </FormProvider>
+  );
+}

web/src/i18n/en.json

@@ -84,7 +84,7 @@
   "shellNotWork": {
     "step1": "1. Try to enable the debug mode, regenerate the memory shell and inject it, check the console or log",
     "step2": "2. If an exception stack trace is displayed, or no exception is seen, please take a screenshot of the current generation interface and the exception log, and describe the target environment as much as possible for feedback",
-    "title": "Shell Not Work ？"
+    "title": "Shell Not Work ?"
   },
   "shellToolConfig": {
     "antSwordPass": "Shell pwd",
@@ -107,7 +107,8 @@
     "neoreGeorgKey": "Connection Key",
     "paramName": "Param Name",
     "pass": "Pass",
-    "suo5Header": "AdvanceConfiguration -> Request Header"
+    "suo5Header": "AdvanceConfiguration -> Request Header",
+    "base64String": "Shell Class"
   },
   "success": {
     "generated": "Generation successful"
@@ -128,10 +129,13 @@
     "targetServerNotFound": "Target server not found?",
     "targetServerRequest": "Request",
     "try-to-use-shell": "Try to use the memory shell",
-    "waitingForGeneration": "// Waiting for generation..."
+    "waitingForGeneration": "// Waiting for generation...",
+    "customShellClass": "Custom shell class is required, base64 or classfile"
   },
   "version": {
     "updateAvailable": "Update Available",
     "updateAvailableTooltip": "Click to Open Github Release Page ( v{{currentVersion}} -> v{{latestVersion}})"
-  }
+  },
+  "generator": "Generator",
+  "about": "About"
 }

web/src/i18n/zh-CN.json

@@ -84,7 +84,7 @@
   "shellNotWork": {
     "step1": "1. 尝试开启调试模式，重新生成内存马并注入，查看控制台或日志",
     "step2": "2. 如果出现异常堆栈信息，或未见异常，请截图当前生成界面以及异常日志，并尽可能描述目标环境进行反馈",
-    "title": "内存马利用失败 ？"
+    "title": "内存马利用失败？"
   },
   "shellToolConfig": {
     "antSwordPass": "连接密码",
@@ -107,7 +107,8 @@
     "neoreGeorgKey": "连接密钥",
     "paramName": "请求参数",
     "pass": "密码",
-    "suo5Header": "高级配置 -> 请求头"
+    "suo5Header": "高级配置 -> 请求头",
+    "base64String": "内存马类"
   },
   "success": {
     "generated": "生成成功"
@@ -117,21 +118,24 @@
     "decompileTip": "反编译还在开发中，因此当前仅能看到 base64 编码格式",
     "download-jattach": "下载 Jattach 工具（后期考虑直接封装在 Jar 中）",
     "execute-command": "执行命令进行注入：/path/to/jattach pid load instrument false /path/to/agent.jar",
-    "get-pid": "获取目标 jvm 的进程 pid （使用 jps 或 ps）",
+    "get-pid": "获取目标 jvm 的进程 pid（使用 jps 或 ps）",
     "handlerUrlPattern": "HandlerMethod/HandlerFunction 类型的需要填写具体的 URL Pattern，例如 /hello_handler",
     "jreTip": "目标 JRE 版本，一般而言为了最大的兼容性，默认 Java 6 即可，Java 高版本能加载低版本的字节码。",
     "jreTip2": "特定情况下，例如 JDK8 才能使用 lambda 表达式，JDK9 以上存在模块限制时才需要选择特定的版本。",
     "move-to-container": "将 MemShellAgent.jar 和 jattach 移动到容器中（如果测试环境使用容器部署）",
     "servletUrlPattern": "Servlet 类型的需要填写具体的 URL Pattern，例如 /hello_servlet",
-    "shellBytesEmpty": "内存马字节码为空，无法下载, 请先生成内存马",
+    "shellBytesEmpty": "内存马字节码为空，无法下载，请先生成内存马",
     "shellToolNotSelected": "请先选择内存马工具类型",
-    "targetServerNotFound": "找不到目标服务 ？",
+    "targetServerNotFound": "找不到目标服务？",
     "targetServerRequest": "请求适配",
     "try-to-use-shell": "尝试利用内存马",
-    "waitingForGeneration": "// 等待填写参数生成中..."
+    "waitingForGeneration": "// 等待填写参数生成中...",
+    "customShellClass": "请输入自定义内存马类，base64 或类文件"
   },
   "version": {
     "updateAvailable": "有可用升级",
     "updateAvailableTooltip": "点击前往 GitHub Release ( v{{currentVersion}} -> v{{latestVersion}})"
-  }
+  },
+  "about": "关于",
+  "generator": "生成器"
 }

web/src/routes/index.tsx

@@ -5,7 +5,13 @@ import { Button } from "@/components/ui/button";
 import { Form } from "@/components/ui/form.tsx";
 import { env } from "@/config.ts";
 import { FormSchema, formSchema } from "@/types/schema.ts";
-import { APIErrorResponse, ConfigResponseType, GenerateResponse, GenerateResult } from "@/types/shell.ts";
+import {
+  APIErrorResponse,
+  ConfigResponseType,
+  GenerateResponse,
+  GenerateResult,
+  ShellToolType,
+} from "@/types/shell.ts";
 import { transformToPostData } from "@/utils/transformer.ts";
 import { zodResolver } from "@hookform/resolvers/zod";
 import { useQuery } from "@tanstack/react-query";
@@ -76,6 +82,11 @@ function IndexComponent() {
       toast.warning(t("tips.handlerUrlPattern"));
       return false;
     }
+
+    if (values.shellTool === ShellToolType.Custom && !values.shellClassBase64) {
+      toast.warning(t("tips.customShellClass"));
+      return false;
+    }
     return true;
   }
 

web/src/types/schema.ts

@@ -19,6 +19,7 @@ export const formSchema = z.object({
   injectorClassName: z.optional(z.string()),
   packingMethod: z.string().min(1),
   shrink: z.optional(z.boolean()),
+  shellClassBase64: z.optional(z.string()),
 });
 
 export type FormSchema = z.infer<typeof formSchema>;

web/src/types/shell.ts

@@ -18,6 +18,7 @@ export interface ShellToolConfig {
   antSwordPass?: string;
   headerName?: string;
   headerValue?: string;
+  shellClassBase64?: string;
 }
 
 export interface CommandShellToolConfig {
@@ -66,11 +67,15 @@ export interface InjectorConfig {
 }
 
 export interface ConfigResponseType {
-  servers: string[];
+  servers: ServerConfig;
   core: MainConfig;
   packers: PackerConfig;
 }
 
+export interface ServerConfig {
+  [serverName: string]: Array<string>;
+}
+
 export interface MainConfig {
   [serverName: string]: {
     [toolName: string]: string[];
@@ -117,4 +122,5 @@ export enum ShellToolType {
   AntSword = "AntSword",
   Suo5 = "Suo5",
   NeoreGeorg = "NeoreGeorg",
+  Custom = "Custom",
 }

web/src/utils/transformer.ts

@@ -20,6 +20,7 @@ export function transformToPostData(formValue: FormSchema) {
     antSwordPass: formValue.antSwordPass,
     headerName: formValue.headerName,
     headerValue: formValue.headerValue,
+    shellClassBase64: formValue.shellClassBase64,
   };
 
   const injectorConfig: InjectorConfig = {