feat: support custom shell generator (#49)

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<?>, 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<?>, 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());
         }

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

@@ -35,5 +35,10 @@ public enum ShellTool {
      */
     NeoreGeorg,
 
+    /**
+     * 自定义
+     */
+    Custom,
+
     ;
 }

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

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

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("/", "."));
+    }
+}