mirror of
https://github.com/ReaJason/MemShellParty.git
synced 2026-09-23 23:41:52 +08:00
fix: custom shell not supported for listener and valve vistor
This commit is contained in:
+41
-4
@@ -1,17 +1,21 @@
|
||||
package com.reajason.javaweb.memshell.generator;
|
||||
|
||||
import com.reajason.javaweb.Server;
|
||||
import com.reajason.javaweb.asm.ClassReferenceVisitor;
|
||||
import com.reajason.javaweb.memshell.ShellType;
|
||||
import com.reajason.javaweb.memshell.config.CustomConfig;
|
||||
import com.reajason.javaweb.memshell.config.ShellConfig;
|
||||
import com.reajason.javaweb.memshell.shelltool.godzilla.GodzillaValve;
|
||||
import com.reajason.javaweb.utils.CommonUtil;
|
||||
import lombok.SneakyThrows;
|
||||
import net.bytebuddy.ByteBuddy;
|
||||
import net.bytebuddy.jar.asm.ClassReader;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.objectweb.asm.ClassReader;
|
||||
|
||||
import java.util.Base64;
|
||||
import java.util.Set;
|
||||
|
||||
import static org.junit.jupiter.api.Assertions.assertEquals;
|
||||
|
||||
import static org.junit.jupiter.api.Assertions.*;
|
||||
|
||||
|
||||
/**
|
||||
@@ -27,9 +31,42 @@ class CustomShellGeneratorTest {
|
||||
.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.getEncoder().encodeToString(bytes)).build()).getBytes();
|
||||
ShellConfig shellConfig = ShellConfig.builder()
|
||||
.shellType(ShellType.FILTER)
|
||||
.build();
|
||||
CustomConfig customConfig = CustomConfig.builder()
|
||||
.shellClassName(className)
|
||||
.shellClassBase64(Base64.getEncoder().encodeToString(bytes))
|
||||
.build();
|
||||
byte[] bytes1 = new CustomShellGenerator(shellConfig, customConfig).getBytes();
|
||||
|
||||
ClassReader classReader = new ClassReader(bytes1);
|
||||
assertEquals(className, classReader.getClassName().replace("/", "."));
|
||||
}
|
||||
|
||||
@Test
|
||||
@SneakyThrows
|
||||
void testValue() {
|
||||
byte[] bytes = new ByteBuddy()
|
||||
.redefine(GodzillaValve.class)
|
||||
.name(CommonUtil.generateShellClassName()).make().getBytes();
|
||||
String className = CommonUtil.generateShellClassName();
|
||||
ShellConfig shellConfig = ShellConfig.builder()
|
||||
.server(Server.BES)
|
||||
.shellType(ShellType.VALVE)
|
||||
.build();
|
||||
CustomConfig customConfig = CustomConfig.builder()
|
||||
.shellClassName(className)
|
||||
.shellClassBase64(Base64.getEncoder().encodeToString(bytes))
|
||||
.build();
|
||||
byte[] bytes1 = new CustomShellGenerator(shellConfig, customConfig).getBytes();
|
||||
|
||||
ClassReader classReader = new ClassReader(bytes1);
|
||||
ClassReferenceVisitor classVisitor = new ClassReferenceVisitor();
|
||||
classReader.accept(classVisitor, 0);
|
||||
Set<String> referencedClasses = classVisitor.getReferencedClasses();
|
||||
assertEquals(className, classReader.getClassName().replace("/", "."));
|
||||
assertTrue(referencedClasses.contains("com/bes/enterprise/webtier/Valve"));
|
||||
assertFalse(referencedClasses.contains("org/apache/catalina/Valve"));
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user