refactor: extract method

This commit is contained in:
ReaJason
2025-03-30 15:33:31 +08:00
parent e8124c02cc
commit 17fe4b4746
@@ -1,6 +1,7 @@
package com.reajason.javaweb.memshell.packer.jar; package com.reajason.javaweb.memshell.packer.jar;
import com.reajason.javaweb.asm.ClassRenameUtils; import com.reajason.javaweb.asm.ClassRenameUtils;
import com.reajason.javaweb.memshell.ShellType;
import com.reajason.javaweb.memshell.config.GenerateResult; import com.reajason.javaweb.memshell.config.GenerateResult;
import com.reajason.javaweb.memshell.utils.CommonUtil; import com.reajason.javaweb.memshell.utils.CommonUtil;
import lombok.SneakyThrows; import lombok.SneakyThrows;
@@ -17,75 +18,95 @@ import java.nio.file.Files;
import java.nio.file.Path; import java.nio.file.Path;
import java.util.Enumeration; import java.util.Enumeration;
import java.util.Map; import java.util.Map;
import java.util.jar.JarEntry; import java.util.jar.*;
import java.util.jar.JarFile;
import java.util.jar.JarOutputStream;
import java.util.jar.Manifest;
/** /**
* @author ReaJason * @author ReaJason
* @since 2025/1/1 * @since 2025/1/1
*/ */
public class AgentJarPacker implements JarPacker { public class AgentJarPacker implements JarPacker {
private static Path tempBootPath;
static Path tempBootPath;
@Override @Override
@SneakyThrows @SneakyThrows
public byte[] packBytes(GenerateResult generateResult) { public byte[] packBytes(GenerateResult generateResult) {
String mainClass = generateResult.getInjectorClassName(); Manifest manifest = createManifest(generateResult.getInjectorClassName());
String advisorClass = generateResult.getShellClassName(); ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
String relocatePrefix = CommonUtil.getRandomPackageName().replace(".", "/") + "/";
boolean isAsm = generateResult.getShellConfig().getShellType().endsWith(ShellType.ASM);
try (JarOutputStream targetJar = new JarOutputStream(outputStream, manifest)) {
addDependencies(targetJar, relocatePrefix, isAsm);
addClassesToJar(targetJar, generateResult, relocatePrefix, isAsm);
}
return outputStream.toByteArray();
}
private Manifest createManifest(String mainClass) {
Manifest manifest = new Manifest(); Manifest manifest = new Manifest();
manifest.getMainAttributes().putValue("Manifest-Version", "1.0"); Attributes attributes = manifest.getMainAttributes();
manifest.getMainAttributes().putValue("Agent-Class", mainClass); attributes.putValue("Manifest-Version", "1.0");
manifest.getMainAttributes().putValue("Premain-Class", mainClass); attributes.putValue("Agent-Class", mainClass);
manifest.getMainAttributes().putValue("Can-Redefine-Classes", "true"); attributes.putValue("Premain-Class", mainClass);
manifest.getMainAttributes().putValue("Can-Retransform-Classes", "true"); attributes.putValue("Can-Redefine-Classes", "true");
ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream(); attributes.putValue("Can-Retransform-Classes", "true");
return manifest;
}
String RELOCATE_PREFIX = CommonUtil.getRandomPackageName().replace(".", "/") + "/"; @SneakyThrows
boolean RELOCATE_ENABLED = true; private void addDependencies(JarOutputStream targetJar, String relocatePrefix, boolean isAsm) {
if (isAsm) {
try (JarOutputStream targetJar = new JarOutputStream(byteArrayOutputStream, manifest)) { addDependency(targetJar, Opcodes.class, true, relocatePrefix);
String dependencyPackageName = null;
if (generateResult.getShellConfig().getShellType().endsWith("ASM")) {
dependencyPackageName = Opcodes.class.getPackage().getName();
addDependency(targetJar, Opcodes.class, true, RELOCATE_PREFIX);
} else { } else {
RELOCATE_ENABLED = false; addDependency(targetJar, ByteBuddy.class, false, relocatePrefix);
dependencyPackageName = ByteBuddy.class.getPackage().getName(); }
addDependency(targetJar, ByteBuddy.class, false, RELOCATE_PREFIX);
} }
byte[] injectorBytes = generateResult.getInjectorBytes(); @SneakyThrows
if (RELOCATE_ENABLED) { private void addClassesToJar(JarOutputStream targetJar, GenerateResult generateResult,
injectorBytes = ClassRenameUtils.relocateClass(injectorBytes, dependencyPackageName, RELOCATE_PREFIX + dependencyPackageName); String relocatePrefix, boolean isRelocateEnabled) {
} String dependencyPackage = isRelocateEnabled ?
targetJar.putNextEntry(new JarEntry(mainClass.replace('.', '/') + ".class")); Opcodes.class.getPackage().getName() : ByteBuddy.class.getPackage().getName();
targetJar.write(injectorBytes);
targetJar.closeEntry();
byte[] shellBytes = generateResult.getShellBytes(); // Add injector class
if (RELOCATE_ENABLED) { addClassEntry(targetJar,
shellBytes = ClassRenameUtils.relocateClass(shellBytes, dependencyPackageName, RELOCATE_PREFIX + dependencyPackageName); generateResult.getInjectorClassName(),
} generateResult.getInjectorBytes(),
targetJar.putNextEntry(new JarEntry(advisorClass.replace('.', '/') + ".class")); dependencyPackage,
targetJar.write(shellBytes); relocatePrefix,
targetJar.closeEntry(); isRelocateEnabled);
// Add shell class
addClassEntry(targetJar,
generateResult.getShellClassName(),
generateResult.getShellBytes(),
dependencyPackage,
relocatePrefix,
isRelocateEnabled);
// Add inner classes
for (Map.Entry<String, byte[]> entry : generateResult.getInjectorInnerClassBytes().entrySet()) { for (Map.Entry<String, byte[]> entry : generateResult.getInjectorInnerClassBytes().entrySet()) {
targetJar.putNextEntry(new JarEntry(entry.getKey().replace('.', '/') + ".class")); addClassEntry(targetJar,
byte[] innerClassBytes = entry.getValue(); entry.getKey(),
if (RELOCATE_ENABLED) { entry.getValue(),
innerClassBytes = ClassRenameUtils.relocateClass(innerClassBytes, dependencyPackageName, RELOCATE_PREFIX + dependencyPackageName); dependencyPackage,
relocatePrefix,
isRelocateEnabled);
} }
targetJar.write(innerClassBytes); }
@SneakyThrows
private void addClassEntry(JarOutputStream targetJar, String className, byte[] classBytes,
String dependencyPackage, String relocatePrefix, boolean isRelocateEnabled) {
targetJar.putNextEntry(new JarEntry(className.replace('.', '/') + ".class"));
byte[] processedBytes = isRelocateEnabled ?
ClassRenameUtils.relocateClass(classBytes, dependencyPackage, relocatePrefix + dependencyPackage) :
classBytes;
targetJar.write(processedBytes);
targetJar.closeEntry(); targetJar.closeEntry();
} }
}
return byteArrayOutputStream.toByteArray();
}
@SneakyThrows @SneakyThrows
public static void addDependency(JarOutputStream targetJar, Class<?> baseClass, boolean relocate, String relocatePrefix) { public static void addDependency(JarOutputStream targetJar, Class<?> baseClass, boolean relocate, String relocatePrefix) {
@@ -128,27 +149,29 @@ public class AgentJarPacker implements JarPacker {
sourceJar.close(); sourceJar.close();
} }
/**
* Extracts a JAR file to a temporary directory
*
* @param jarPath Path to the source JAR file
* @param tempPath Path to the temporary directory
*/
@SneakyThrows @SneakyThrows
public static void unzip(String jarPath, String tempPath) { public static void unzip(String jarPath, String tempPath) {
try (JarFile jarFile = new JarFile(jarPath)) { try (JarFile jarFile = new JarFile(jarPath)) {
Enumeration<JarEntry> entries = jarFile.entries(); Enumeration<JarEntry> entries = jarFile.entries();
while (entries.hasMoreElements()) { while (entries.hasMoreElements()) {
JarEntry jarEntry = entries.nextElement(); JarEntry jarEntry = entries.nextElement();
String entryName = jarEntry.getName(); File targetFile = new File(tempPath, jarEntry.getName());
File file = new File(tempPath, entryName);
if (jarEntry.isDirectory()) { if (jarEntry.isDirectory()) {
file.mkdir(); targetFile.mkdirs();
} else { continue;
InputStream inputStream = null;
FileOutputStream outputStream = null;
try {
inputStream = jarFile.getInputStream(jarEntry);
outputStream = new FileOutputStream(file);
IOUtils.copy(inputStream, outputStream);
} finally {
IOUtils.closeQuietly(inputStream);
IOUtils.closeQuietly(outputStream);
} }
targetFile.getParentFile().mkdirs();
try (InputStream inputStream = jarFile.getInputStream(jarEntry);
FileOutputStream outputStream = new FileOutputStream(targetFile)) {
IOUtils.copy(inputStream, outputStream);
} }
} }
} }