Skip to content

Commit de9de88

Browse files
ReaJasonReaJason
committed
refactor: extract method
1 parent 3602716 commit de9de88

File tree

1 file changed

+90
-67
lines changed

1 file changed

+90
-67
lines changed

‎generator/src/main/java/com/reajason/javaweb/memshell/packer/jar/AgentJarPacker.java‎

Lines changed: 90 additions & 67 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package com.reajason.javaweb.memshell.packer.jar;
22

33
import com.reajason.javaweb.asm.ClassRenameUtils;
4+
import com.reajason.javaweb.memshell.ShellType;
45
import com.reajason.javaweb.memshell.config.GenerateResult;
56
import com.reajason.javaweb.memshell.utils.CommonUtil;
67
import lombok.SneakyThrows;
@@ -17,74 +18,94 @@
1718
import java.nio.file.Path;
1819
import java.util.Enumeration;
1920
import java.util.Map;
20-
import java.util.jar.JarEntry;
21-
import java.util.jar.JarFile;
22-
import java.util.jar.JarOutputStream;
23-
import java.util.jar.Manifest;
21+
import java.util.jar.*;
2422

2523
/**
2624
* @author ReaJason
2725
* @since 2025/1/1
2826
*/
2927
public class AgentJarPacker implements JarPacker {
30-
31-
static Path tempBootPath;
28+
private static Path tempBootPath;
3229

3330
@Override
3431
@SneakyThrows
3532
public byte[] packBytes(GenerateResult generateResult) {
36-
String mainClass = generateResult.getInjectorClassName();
37-
String advisorClass = generateResult.getShellClassName();
33+
Manifest manifest = createManifest(generateResult.getInjectorClassName());
34+
ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
35+
36+
String relocatePrefix = CommonUtil.getRandomPackageName().replace(".", "/") + "/";
37+
boolean isAsm = generateResult.getShellConfig().getShellType().endsWith(ShellType.ASM);
38+
39+
try (JarOutputStream targetJar = new JarOutputStream(outputStream, manifest)) {
40+
addDependencies(targetJar, relocatePrefix, isAsm);
41+
addClassesToJar(targetJar, generateResult, relocatePrefix, isAsm);
42+
}
3843

44+
return outputStream.toByteArray();
45+
}
46+
47+
private Manifest createManifest(String mainClass) {
3948
Manifest manifest = new Manifest();
40-
manifest.getMainAttributes().putValue("Manifest-Version", "1.0");
41-
manifest.getMainAttributes().putValue("Agent-Class", mainClass);
42-
manifest.getMainAttributes().putValue("Premain-Class", mainClass);
43-
manifest.getMainAttributes().putValue("Can-Redefine-Classes", "true");
44-
manifest.getMainAttributes().putValue("Can-Retransform-Classes", "true");
45-
ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
46-
47-
String RELOCATE_PREFIX = CommonUtil.getRandomPackageName().replace(".", "/") + "/";
48-
boolean RELOCATE_ENABLED = true;
49-
50-
try (JarOutputStream targetJar = new JarOutputStream(byteArrayOutputStream, manifest)) {
51-
String dependencyPackageName = null;
52-
if (generateResult.getShellConfig().getShellType().endsWith("ASM")) {
53-
dependencyPackageName = Opcodes.class.getPackage().getName();
54-
addDependency(targetJar, Opcodes.class, true, RELOCATE_PREFIX);
55-
} else {
56-
RELOCATE_ENABLED = false;
57-
dependencyPackageName = ByteBuddy.class.getPackage().getName();
58-
addDependency(targetJar, ByteBuddy.class, false, RELOCATE_PREFIX);
59-
}
49+
Attributes attributes = manifest.getMainAttributes();
50+
attributes.putValue("Manifest-Version", "1.0");
51+
attributes.putValue("Agent-Class", mainClass);
52+
attributes.putValue("Premain-Class", mainClass);
53+
attributes.putValue("Can-Redefine-Classes", "true");
54+
attributes.putValue("Can-Retransform-Classes", "true");
55+
return manifest;
56+
}
6057

61-
byte[] injectorBytes = generateResult.getInjectorBytes();
62-
if (RELOCATE_ENABLED) {
63-
injectorBytes = ClassRenameUtils.relocateClass(injectorBytes, dependencyPackageName, RELOCATE_PREFIX + dependencyPackageName);
64-
}
65-
targetJar.putNextEntry(new JarEntry(mainClass.replace('.', '/') + ".class"));
66-
targetJar.write(injectorBytes);
67-
targetJar.closeEntry();
58+
@SneakyThrows
59+
private void addDependencies(JarOutputStream targetJar, String relocatePrefix, boolean isAsm) {
60+
if (isAsm) {
61+
addDependency(targetJar, Opcodes.class, true, relocatePrefix);
62+
} else {
63+
addDependency(targetJar, ByteBuddy.class, false, relocatePrefix);
64+
}
65+
}
6866

69-
byte[] shellBytes = generateResult.getShellBytes();
70-
if (RELOCATE_ENABLED) {
71-
shellBytes = ClassRenameUtils.relocateClass(shellBytes, dependencyPackageName, RELOCATE_PREFIX + dependencyPackageName);
72-
}
73-
targetJar.putNextEntry(new JarEntry(advisorClass.replace('.', '/') + ".class"));
74-
targetJar.write(shellBytes);
75-
targetJar.closeEntry();
76-
77-
for (Map.Entry<String, byte[]> entry : generateResult.getInjectorInnerClassBytes().entrySet()) {
78-
targetJar.putNextEntry(new JarEntry(entry.getKey().replace('.', '/') + ".class"));
79-
byte[] innerClassBytes = entry.getValue();
80-
if (RELOCATE_ENABLED) {
81-
innerClassBytes = ClassRenameUtils.relocateClass(innerClassBytes, dependencyPackageName, RELOCATE_PREFIX + dependencyPackageName);
82-
}
83-
targetJar.write(innerClassBytes);
84-
targetJar.closeEntry();
85-
}
67+
@SneakyThrows
68+
private void addClassesToJar(JarOutputStream targetJar, GenerateResult generateResult,
69+
String relocatePrefix, boolean isRelocateEnabled) {
70+
String dependencyPackage = isRelocateEnabled ?
71+
Opcodes.class.getPackage().getName() : ByteBuddy.class.getPackage().getName();
72+
73+
// Add injector class
74+
addClassEntry(targetJar,
75+
generateResult.getInjectorClassName(),
76+
generateResult.getInjectorBytes(),
77+
dependencyPackage,
78+
relocatePrefix,
79+
isRelocateEnabled);
80+
81+
// Add shell class
82+
addClassEntry(targetJar,
83+
generateResult.getShellClassName(),
84+
generateResult.getShellBytes(),
85+
dependencyPackage,
86+
relocatePrefix,
87+
isRelocateEnabled);
88+
89+
// Add inner classes
90+
for (Map.Entry<String, byte[]> entry : generateResult.getInjectorInnerClassBytes().entrySet()) {
91+
addClassEntry(targetJar,
92+
entry.getKey(),
93+
entry.getValue(),
94+
dependencyPackage,
95+
relocatePrefix,
96+
isRelocateEnabled);
8697
}
87-
return byteArrayOutputStream.toByteArray();
98+
}
99+
100+
@SneakyThrows
101+
private void addClassEntry(JarOutputStream targetJar, String className, byte[] classBytes,
102+
String dependencyPackage, String relocatePrefix, boolean isRelocateEnabled) {
103+
targetJar.putNextEntry(new JarEntry(className.replace('.', '/') + ".class"));
104+
byte[] processedBytes = isRelocateEnabled ?
105+
ClassRenameUtils.relocateClass(classBytes, dependencyPackage, relocatePrefix + dependencyPackage) :
106+
classBytes;
107+
targetJar.write(processedBytes);
108+
targetJar.closeEntry();
88109
}
89110

90111
@SneakyThrows
@@ -128,27 +149,29 @@ public static void addDependency(JarOutputStream targetJar, Class<?> baseClass,
128149
sourceJar.close();
129150
}
130151

152+
/**
153+
* Extracts a JAR file to a temporary directory
154+
*
155+
* @param jarPath Path to the source JAR file
156+
* @param tempPath Path to the temporary directory
157+
*/
131158
@SneakyThrows
132159
public static void unzip(String jarPath, String tempPath) {
133160
try (JarFile jarFile = new JarFile(jarPath)) {
134161
Enumeration<JarEntry> entries = jarFile.entries();
135162
while (entries.hasMoreElements()) {
136163
JarEntry jarEntry = entries.nextElement();
137-
String entryName = jarEntry.getName();
138-
File file = new File(tempPath, entryName);
164+
File targetFile = new File(tempPath, jarEntry.getName());
165+
139166
if (jarEntry.isDirectory()) {
140-
file.mkdir();
141-
} else {
142-
InputStream inputStream = null;
143-
FileOutputStream outputStream = null;
144-
try {
145-
inputStream = jarFile.getInputStream(jarEntry);
146-
outputStream = new FileOutputStream(file);
147-
IOUtils.copy(inputStream, outputStream);
148-
} finally {
149-
IOUtils.closeQuietly(inputStream);
150-
IOUtils.closeQuietly(outputStream);
151-
}
167+
targetFile.mkdirs();
168+
continue;
169+
}
170+
171+
targetFile.getParentFile().mkdirs();
172+
try (InputStream inputStream = jarFile.getInputStream(jarEntry);
173+
FileOutputStream outputStream = new FileOutputStream(targetFile)) {
174+
IOUtils.copy(inputStream, outputStream);
152175
}
153176
}
154177
}

0 commit comments

Comments
 (0)