feat: support tomcat agent shell (resolved #12)

2026-09-22 07:00:43 +08:00 · 2025-01-03 00:25:56 +08:00
parent 8e741b6647
commit 2da2e180cb
84 changed files with 1058 additions and 285 deletions
@@ -1,4 +1,14 @@
 import { UrlPatternTip } from "@/components/tips/url-pattern-tip.tsx";
+import {
+  AlertDialog,
+  AlertDialogAction,
+  AlertDialogCancel,
+  AlertDialogContent,
+  AlertDialogDescription,
+  AlertDialogFooter,
+  AlertDialogHeader,
+  AlertDialogTitle,
+} from "@/components/ui/alert-dialog";
 import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card.tsx";
 import { FormControl, FormDescription, FormField, FormItem, FormLabel } from "@/components/ui/form.tsx";
 import { Input } from "@/components/ui/input.tsx";
@@ -9,6 +19,7 @@ import { Switch } from "@/components/ui/switch.tsx";
 import { FormSchema } from "@/types/schema.ts";
 import { MainConfig } from "@/types/shell.ts";
 import { ServerIcon } from "lucide-react";
+
 import { useState } from "react";
 import { FormProvider, UseFormReturn } from "react-hook-form";

@@ -33,6 +44,8 @@ export function MainConfigCard({
  const [shellToolMap, setShellToolMap] = useState<{ [toolName: string]: string[] }>();
  const [shellTools, setShellTools] = useState<string[]>([]);
  const [shellTypes, setShellTypes] = useState<string[]>([]);
+  const [openAgentConfirm, setOpenAgentConfirm] = useState(false);
+  const [currentShellType, setCurrentShellType] = useState<string>("");

  const handleServerChange = (value: string) => {
    if (mainConfig) {
@@ -121,7 +134,17 @@ export function MainConfigCard({
              render={({ field }) => (
                <FormItem className="space-y-1">
                  <FormLabel>JRE(可选)</FormLabel>
-                  <Select onValueChange={field.onChange} value={field.value}>
+                  <Select
+                    onValueChange={(v) => {
+                      if (Number.parseInt(v) >= 53) {
+                        form.setValue("bypassJavaModule", true);
+                      } else {
+                        form.setValue("bypassJavaModule", false);
+                      }
+                      field.onChange(v);
+                    }}
+                    value={field.value}
+                  >
                    <FormControl>
                      <SelectTrigger className="h-8">
                        <SelectValue placeholder="请选择" />
@@ -224,7 +247,39 @@ export function MainConfigCard({
              render={({ field }) => (
                <FormItem className="space-y-1">
                  <FormLabel>内存马挂载类型</FormLabel>
-                  <Select onValueChange={field.onChange} value={field.value}>
+                  <AlertDialog open={openAgentConfirm} onOpenChange={setOpenAgentConfirm}>
+                    <AlertDialogContent>
+                      <AlertDialogHeader>
+                        <AlertDialogTitle>Agent 注入当前仅支持 Java8 以上，是否仍要选择？</AlertDialogTitle>
+                        <AlertDialogDescription>确认后会将 JRE 版本改为 Java8</AlertDialogDescription>
+                      </AlertDialogHeader>
+                      <AlertDialogFooter>
+                        <AlertDialogCancel>Cancel</AlertDialogCancel>
+                        <AlertDialogAction
+                          onClick={() => {
+                            form.setValue("targetJdkVersion", "52");
+                            form.setValue("shellType", currentShellType);
+                          }}
+                        >
+                          Confirm
+                        </AlertDialogAction>
+                      </AlertDialogFooter>
+                    </AlertDialogContent>
+                  </AlertDialog>
+                  <Select
+                    onValueChange={(v) => {
+                      if (
+                        v.startsWith("Agent") &&
+                        Number.parseInt(form.getValues("targetJdkVersion") as string) === 50
+                      ) {
+                        setOpenAgentConfirm(true);
+                        setCurrentShellType(v);
+                      } else {
+                        field.onChange(v);
+                      }
+                    }}
+                    value={field.value}
+                  >
                    <FormControl>
                      <SelectTrigger className="h-8">
                        <SelectValue placeholder="请选择" />
@@ -4,17 +4,64 @@ import { Alert, AlertDescription, AlertTitle } from "@/components/ui/alert.tsx";
 import { Button } from "@/components/ui/button.tsx";
 import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card.tsx";
 import { Label } from "@/components/ui/label.tsx";
+import { Separator } from "@/components/ui/separator.tsx";
 import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs.tsx";
-import { downloadJavaClass } from "@/lib/utils.ts";
+import { downloadBytes } from "@/lib/utils.ts";
 import { GenerateResult } from "@/types/shell.ts";
 import { TicketsIcon, TriangleAlertIcon } from "lucide-react";

+function AgentResult({ packResult, generateResult }: { packResult: string; generateResult?: GenerateResult }) {
+  return (
+    <section>
+      <ol className="list-decimal list-inside space-y-4">
+        <li className="flex items-center justify-between">
+          <span>下载 MemShellAgent.jar</span>
+          <Button
+            size="sm"
+            variant="outline"
+            className="w-28"
+            type="button"
+            onClick={() =>
+              downloadBytes(
+                packResult,
+                undefined,
+                `${generateResult?.shellConfig.server}${generateResult?.shellConfig.shellTool}MemShellAgent`,
+              )
+            }
+          >
+            下载 Jar
+          </Button>
+        </li>
+        <li className="flex items-center justify-between">
+          <span>下载 Jattach 工具（后期考虑直接封装在 Jar 中）</span>
+          <Button
+            size="sm"
+            variant="outline"
+            className="w-28"
+            type="button"
+            onClick={() => window.open("https://github.com/jattach/jattach/releases")}
+          >
+            下载 Jattach
+          </Button>
+        </li>
+        <Separator />
+        <h2 className="text-2xl font-bold mb-4">使用方法：</h2>
+        <li>将 MemShellAgent.jar 和 jattach 移动到容器中（如果测试环境使用容器部署）</li>
+        <li>获取目标 jvm 的进程 pid （使用 jps 或 ps）</li>
+        <li>执行命令进行注入：/path/to/jattach pid load instrument false /path/to/agent.jar</li>
+        <li>尝试连接测试</li>
+      </ol>
+    </section>
+  );
+}
+
 export function ShellResult({
  packResult,
  packMethod,
  generateResult,
 }: { packResult: string; packMethod: string; generateResult?: GenerateResult }) {
  const showCode = packMethod === "JSP";
+  const isAgent = packMethod.startsWith("Agent");
  return (
    <Card className="h-full">
      <CardHeader className="pb-2">
@@ -31,7 +78,7 @@ export function ShellResult({
            <TabsTrigger value="injector">注入器类</TabsTrigger>
          </TabsList>
          <TabsContent value="packResult" className="mt-4">
-            {generateResult && (
+            {generateResult && !isAgent && (
              <div className="gap-4 my-2">
                <CopyableField
                  label="注入器类名"
@@ -45,12 +92,15 @@ export function ShellResult({
                />
              </div>
            )}
-            <CodeViewer
-              code={packResult}
-              wrapLongLines={!showCode}
-              showLineNumbers={showCode}
-              language={showCode ? "java" : "text"}
-            />
+            {!isAgent && (
+              <CodeViewer
+                code={packResult}
+                wrapLongLines={!showCode}
+                showLineNumbers={showCode}
+                language={showCode ? "java" : "text"}
+              />
+            )}
+            {isAgent && <AgentResult packResult={packResult} generateResult={generateResult} />}
          </TabsContent>
          <TabsContent value="shell" className="mt-4">
            <Alert>
@@ -71,14 +121,17 @@ export function ShellResult({
                  </div>
                </div>
              )}
-              <Button
-                size="sm"
-                className="h-8 gap-1"
-                type="button"
-                onClick={() => downloadJavaClass(generateResult?.shellBytesBase64Str, generateResult?.shellClassName)}
-              >
-                下载 Class
-              </Button>
+              {generateResult && (
+                <Button
+                  size="sm"
+                  variant="outline"
+                  className="w-28"
+                  type="button"
+                  onClick={() => downloadBytes(generateResult?.shellBytesBase64Str, generateResult?.shellClassName)}
+                >
+                  下载 Class
+                </Button>
+              )}
            </div>
            <CodeViewer
              showLineNumbers={false}
@@ -106,16 +159,19 @@ export function ShellResult({
                  </div>
                </div>
              )}
-              <Button
-                size="sm"
-                className="h-8 gap-1"
-                type="button"
-                onClick={() =>
-                  downloadJavaClass(generateResult?.injectorBytesBase64Str, generateResult?.injectorClassName)
-                }
-              >
-                下载 Class
-              </Button>
+              {generateResult && (
+                <Button
+                  size="sm"
+                  className="w-28"
+                  variant="outline"
+                  type="button"
+                  onClick={() =>
+                    downloadBytes(generateResult?.injectorBytesBase64Str, generateResult?.injectorClassName)
+                  }
+                >
+                  下载 Class
+                </Button>
+              )}
            </div>
            <CodeViewer
              showLineNumbers={false}
@@ -11,6 +11,8 @@ export function UrlPatternTip() {
        <TooltipContent>
          <p>当使用 Servlet 内存马时必须写具体的 urlPattern，不能使用 /*，不然无法使用</p>
          <p>当使用 SpringMVC ControllerHandler 内存马时必须写具体的 urlPattern，不能使用 /*，不然无法使用</p>
+          <p>当使用 SpringWebFlux HandlerMethod 内存马时必须写具体的 urlPattern，不能使用 /*，不然无法使用</p>
+          <p>当使用 SpringWebFlux HandlerFunction 内存马时必须写具体的 urlPattern，不能使用 /*，不然无法使用</p>
        </TooltipContent>
      </Tooltip>
    </TooltipProvider>
@@ -0,0 +1,104 @@
+import * as React from "react";
+import * as AlertDialogPrimitive from "@radix-ui/react-alert-dialog";
+
+import { cn } from "@/lib/utils";
+import { buttonVariants } from "@/components/ui/button";
+
+const AlertDialog = AlertDialogPrimitive.Root;
+
+const AlertDialogTrigger = AlertDialogPrimitive.Trigger;
+
+const AlertDialogPortal = AlertDialogPrimitive.Portal;
+
+const AlertDialogOverlay = React.forwardRef<
+  React.ElementRef<typeof AlertDialogPrimitive.Overlay>,
+  React.ComponentPropsWithoutRef<typeof AlertDialogPrimitive.Overlay>
+>(({ className, ...props }, ref) => (
+  <AlertDialogPrimitive.Overlay
+    className={cn(
+      "fixed inset-0 z-50 bg-black/80 data-[state=open]:animate-in data-[state=closed]:animate-out data-[state=closed]:fade-out-0 data-[state=open]:fade-in-0",
+      className,
+    )}
+    {...props}
+    ref={ref}
+  />
+));
+AlertDialogOverlay.displayName = AlertDialogPrimitive.Overlay.displayName;
+
+const AlertDialogContent = React.forwardRef<
+  React.ElementRef<typeof AlertDialogPrimitive.Content>,
+  React.ComponentPropsWithoutRef<typeof AlertDialogPrimitive.Content>
+>(({ className, ...props }, ref) => (
+  <AlertDialogPortal>
+    <AlertDialogOverlay />
+    <AlertDialogPrimitive.Content
+      ref={ref}
+      className={cn(
+        "fixed left-[50%] top-[50%] z-50 grid w-full max-w-lg translate-x-[-50%] translate-y-[-50%] gap-4 border bg-background p-6 shadow-lg duration-200 data-[state=open]:animate-in data-[state=closed]:animate-out data-[state=closed]:fade-out-0 data-[state=open]:fade-in-0 data-[state=closed]:zoom-out-95 data-[state=open]:zoom-in-95 data-[state=closed]:slide-out-to-left-1/2 data-[state=closed]:slide-out-to-top-[48%] data-[state=open]:slide-in-from-left-1/2 data-[state=open]:slide-in-from-top-[48%] sm:rounded-lg",
+        className,
+      )}
+      {...props}
+    />
+  </AlertDialogPortal>
+));
+AlertDialogContent.displayName = AlertDialogPrimitive.Content.displayName;
+
+const AlertDialogHeader = ({ className, ...props }: React.HTMLAttributes<HTMLDivElement>) => (
+  <div className={cn("flex flex-col space-y-2 text-center sm:text-left", className)} {...props} />
+);
+AlertDialogHeader.displayName = "AlertDialogHeader";
+
+const AlertDialogFooter = ({ className, ...props }: React.HTMLAttributes<HTMLDivElement>) => (
+  <div className={cn("flex flex-col-reverse sm:flex-row sm:justify-end sm:space-x-2", className)} {...props} />
+);
+AlertDialogFooter.displayName = "AlertDialogFooter";
+
+const AlertDialogTitle = React.forwardRef<
+  React.ElementRef<typeof AlertDialogPrimitive.Title>,
+  React.ComponentPropsWithoutRef<typeof AlertDialogPrimitive.Title>
+>(({ className, ...props }, ref) => (
+  <AlertDialogPrimitive.Title ref={ref} className={cn("text-lg font-semibold", className)} {...props} />
+));
+AlertDialogTitle.displayName = AlertDialogPrimitive.Title.displayName;
+
+const AlertDialogDescription = React.forwardRef<
+  React.ElementRef<typeof AlertDialogPrimitive.Description>,
+  React.ComponentPropsWithoutRef<typeof AlertDialogPrimitive.Description>
+>(({ className, ...props }, ref) => (
+  <AlertDialogPrimitive.Description ref={ref} className={cn("text-sm text-muted-foreground", className)} {...props} />
+));
+AlertDialogDescription.displayName = AlertDialogPrimitive.Description.displayName;
+
+const AlertDialogAction = React.forwardRef<
+  React.ElementRef<typeof AlertDialogPrimitive.Action>,
+  React.ComponentPropsWithoutRef<typeof AlertDialogPrimitive.Action>
+>(({ className, ...props }, ref) => (
+  <AlertDialogPrimitive.Action ref={ref} className={cn(buttonVariants(), className)} {...props} />
+));
+AlertDialogAction.displayName = AlertDialogPrimitive.Action.displayName;
+
+const AlertDialogCancel = React.forwardRef<
+  React.ElementRef<typeof AlertDialogPrimitive.Cancel>,
+  React.ComponentPropsWithoutRef<typeof AlertDialogPrimitive.Cancel>
+>(({ className, ...props }, ref) => (
+  <AlertDialogPrimitive.Cancel
+    ref={ref}
+    className={cn(buttonVariants({ variant: "outline" }), "mt-2 sm:mt-0", className)}
+    {...props}
+  />
+));
+AlertDialogCancel.displayName = AlertDialogPrimitive.Cancel.displayName;
+
+export {
+  AlertDialog,
+  AlertDialogPortal,
+  AlertDialogOverlay,
+  AlertDialogTrigger,
+  AlertDialogContent,
+  AlertDialogHeader,
+  AlertDialogFooter,
+  AlertDialogTitle,
+  AlertDialogDescription,
+  AlertDialogAction,
+  AlertDialogCancel,
+};
@@ -6,9 +6,9 @@ export function cn(...inputs: ClassValue[]) {
  return twMerge(clsx(inputs));
 }

-export function downloadJavaClass(base64String?: string, className?: string) {
-  if (!base64String || !className) {
-    toast.warning("内存马字节码为空，无法下载, 请先生成内存马");
+export function downloadBytes(base64String?: string, className?: string, jarName?: string) {
+  if (!base64String) {
+    toast.warning("字节码为空，无法下载, 请先生成内存马");
    return;
  }
  const byteCharacters = atob(base64String);
@@ -19,12 +19,12 @@ export function downloadJavaClass(base64String?: string, className?: string) {
  const byteArray = new Uint8Array(byteNumbers);

  // Create a Blob from the byte array
-  const blob = new Blob([byteArray], { type: "application/java-vm" });
+  const blob = new Blob([byteArray], { type: className ? "application/java-vm" : "application/java-archive" });

  // Create a download link
  const link = document.createElement("a");
  link.href = window.URL.createObjectURL(blob);
-  link.download = `${className.substring(className.lastIndexOf("."))}.class`;
+  link.download = className ? `${className.substring(className.lastIndexOf("."))}.class` : `${jarName}.jar`;

  document.body.appendChild(link);
  link.click();
@@ -57,15 +57,46 @@ function IndexComponent() {
  const [packMethod, setPackMethod] = useState<string>("");
  const [isActionPending, startTransition] = useTransition();

+  function customValidation(values: FormSchema) {
+    if (values.shellType.endsWith("Servlet") && (values.urlPattern === "/*" || !values.urlPattern)) {
+      toast.warning("Servlet 类型的需要填写具体的 URL Pattern，例如 /hello_servlet");
+      return false;
+    }
+
+    if (values.shellType.endsWith("ControllerHandler") && (values.urlPattern === "/*" || !values.urlPattern)) {
+      toast.warning("ControllerHandler 类型的需要填写具体的 URL Pattern，例如 /hello_controller");
+      return false;
+    }
+
+    if (
+      (values.shellType === "HandlerMethod" || values.shellType === "HandlerFunction") &&
+      (values.urlPattern === "/*" || !values.urlPattern)
+    ) {
+      toast.warning("HandlerMethod/HandlerFunction 类型的需要填写具体的 URL Pattern，例如 /hello_handler");
+      return false;
+    }
+
+    if (values.shellType.startsWith("Agent") && values.packingMethod !== "AgentJar") {
+      toast.warning("Agent 注入方式当前仅支持 AgentJar 打包方式");
+      return false;
+    }
+
+    if (values.shellType.startsWith("Agent") && values.targetJdkVersion === "50") {
+      toast.warning("Agent 注入方式当前仅支持 Java8 以上");
+      return false;
+    }
+
+    if (!values.shellType.startsWith("Agent") && values.packingMethod === "AgentJar") {
+      toast.warning("Agent 注入方式当前仅支持 Tomcat，只有 Agent 注入方式才可使用 AgentJar 打包方式");
+      return false;
+    }
+
+    return true;
+  }
+
  async function onSubmit(values: FormSchema) {
    startTransition(async () => {
-      if (values.shellType.endsWith("Servlet") && (values.urlPattern === "/*" || !values.urlPattern)) {
-        toast.warning("Servlet 类型的需要填写具体的 URL Pattern，例如 /hello_servlet");
-        return;
-      }
-
-      if (values.shellType.endsWith("ControllerHandler") && (values.urlPattern === "/*" || !values.urlPattern)) {
-        toast.warning("ControllerHandler 类型的需要填写具体的 URL Pattern，例如 /hello_controller");
+      if (!customValidation(values)) {
        return;
      }