refactor: simplify generateBytes code

Author: ReaJason

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

@@ -1,8 +1,10 @@
 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.config.GenerateResult;
+import com.reajason.javaweb.memshell.config.InjectorConfig;
+import com.reajason.javaweb.memshell.config.ShellConfig;
+import com.reajason.javaweb.memshell.config.ShellToolConfig;
+import com.reajason.javaweb.memshell.generator.InjectorGenerator;
 import com.reajason.javaweb.memshell.server.AbstractShell;
 import com.reajason.javaweb.memshell.utils.CommonUtil;
 import org.apache.commons.lang3.StringUtils;
@@ -45,7 +47,7 @@ public static GenerateResult generate(ShellConfig shellConfig, InjectorConfig in
             shellToolConfig.setShellClass(shellClass);
         }
 
-        byte[] shellBytes = generateShellBytes(shellConfig, shellToolConfig);
+        byte[] shellBytes = shellConfig.getShellTool().generateBytes(shellConfig, shellToolConfig);
 
         injectorConfig.setInjectorClass(injectorClass);
         injectorConfig.setShellClassName(shellToolConfig.getShellClassName());
@@ -66,25 +68,4 @@ public static GenerateResult generate(ShellConfig shellConfig, InjectorConfig in
                 .injectorInnerClassBytes(innerClassBytes)
                 .build();
     }
-
-    private static byte[] generateShellBytes(ShellConfig shellConfig, ShellToolConfig shellToolConfig) {
-        switch (shellConfig.getShellTool()) {
-            case Godzilla:
-                return new GodzillaGenerator(shellConfig, (GodzillaConfig) shellToolConfig).getBytes();
-            case Command:
-                return new CommandGenerator(shellConfig, (CommandConfig) shellToolConfig).getBytes();
-            case Behinder:
-                return new BehinderGenerator(shellConfig, (BehinderConfig) shellToolConfig).getBytes();
-            case Suo5:
-                return new Suo5Generator(shellConfig, ((Suo5Config) shellToolConfig)).getBytes();
-            case AntSword:
-                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

@@ -1,44 +1,40 @@
 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 java.lang.reflect.Constructor;
+
 /**
  * @author ReaJason
  * @since 2024/11/22
  */
 public enum ShellTool {
-    /**
-     * 哥斯拉
-     */
-    Godzilla,
-
-    /**
-     * 冰蝎
-     */
-    Behinder,
-
-    /**
-     * 蚁剑
-     */
-    AntSword,
-
-    /**
-     * 命令回显
-     */
-    Command,
-
-    /**
-     * Suo5 隧道代理 <a href="https://github.com/zema1/suo5">zema1/suo5</a>
-     */
-    Suo5,
-
-    /**
-     * Neo-reGeorg <a href="https://github.com/L-codes/Neo-reGeorg">L-codes/Neo-reGeorg</a>
-     */
-    NeoreGeorg,
-
-    /**
-     * 自定义
-     */
-    Custom,
-
-    ;
+    Godzilla(GodzillaGenerator.class, GodzillaConfig.class),
+    Command(CommandGenerator.class, CommandConfig.class),
+    Behinder(BehinderGenerator.class, BehinderConfig.class),
+    Suo5(Suo5Generator.class, Suo5Config.class),
+    AntSword(AntSwordGenerator.class, AntSwordConfig.class),
+    NeoreGeorg(NeoreGeorgGenerator.class, NeoreGeorgConfig.class),
+    Custom(CustomShellGenerator.class, CustomConfig.class);
+
+    private final Class<? extends ShellGenerator> generatorClass;
+    private final Class<? extends ShellToolConfig> configClass;
+
+    ShellTool(Class<? extends ShellGenerator> generatorClass, Class<? extends ShellToolConfig> configClass) {
+        this.generatorClass = generatorClass;
+        this.configClass = configClass;
+    }
+
+    public byte[] generateBytes(ShellConfig shellConfig, ShellToolConfig shellToolConfig) {
+        try {
+            Constructor<? extends ShellGenerator> constructor =
+                    generatorClass.getConstructor(ShellConfig.class, configClass);
+            ShellGenerator generator = constructor.newInstance(shellConfig, configClass.cast(shellToolConfig));
+            return generator.getBytes();
+        } catch (Exception e) {
+            throw new RuntimeException("Failed to create generator for " + this, e);
+        }
+    }
 }