archive/MemShellParty
1
0
Fork 0
mirror of https://github.com/ReaJason/MemShellParty.git synced 2026-09-21 22:50:42 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

fix: add generate exception global handler

Browse Source
This commit is contained in:
ReaJason
2025-08-13 23:53:40 +08:00
parent 5aae3c4dbe
commit 6a475983db
7 changed files with 40 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
boot/src/main/java/com/reajason/javaweb/boot/api/GlobalExceptionHandler.java
+8 -1
View File
@@ -1,5 +1,6 @@
package com.reajason.javaweb.boot.api;
import com.reajason.javaweb.GenerationException;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.HttpStatus;
import org.springframework.http.ResponseEntity;
@@ -22,10 +23,16 @@ public class GlobalExceptionHandler {
return ResponseEntity.status(HttpStatus.NOT_FOUND).build();
}
@ResponseStatus(HttpStatus.INTERNAL_SERVER_ERROR)
@ExceptionHandler(GenerationException.class)
public ErrorResponse handleGenerationException(GenerationException exception) {
return new ErrorResponse(exception.getMessage());
}
@ResponseStatus(HttpStatus.INTERNAL_SERVER_ERROR)
@ExceptionHandler(Throwable.class)
public ErrorResponse handleThrowable(Throwable throwable) {
log.error("请求出错", throwable);
log.error("Internal Exception", throwable);
return new ErrorResponse(throwable.getMessage());
}
}
generator/src/main/java/com/reajason/javaweb/memshell/MemShellGenerator.java
+3 -2
View File
@@ -1,5 +1,6 @@
package com.reajason.javaweb.memshell;
import com.reajason.javaweb.GenerationException;
import com.reajason.javaweb.memshell.config.InjectorConfig;
import com.reajason.javaweb.memshell.config.ShellConfig;
import com.reajason.javaweb.memshell.config.ShellToolConfig;
@@ -21,7 +22,7 @@ public class MemShellGenerator {
String serverName = shellConfig.getServer();
AbstractServer server = ServerFactory.getServer(serverName);
if (server == null) {
throw new IllegalArgumentException("Unsupported server: " + serverName);
throw new GenerationException("Unsupported server: " + serverName);
}
if (StringUtils.isBlank(shellToolConfig.getShellClassName())) {
@@ -39,7 +40,7 @@ public class MemShellGenerator {
} else {
Pair<Class<?>, Class<?>> shellInjectorPair = server.getShellInjectorPair(shellConfig.getShellTool(), shellConfig.getShellType());
if (shellInjectorPair == null) {
throw new UnsupportedOperationException(serverName + " unsupported shell type: " + shellConfig.getShellType() + " for tool: " + shellConfig.getShellTool());
throw new GenerationException(serverName + " unsupported shell type: " + shellConfig.getShellType() + " for tool: " + shellConfig.getShellTool());
}
Class<?> shellClass = shellInjectorPair.getLeft();
injectorClass = shellInjectorPair.getRight();
generator/src/main/java/com/reajason/javaweb/memshell/generator/CustomShellGenerator.java
+2 -1
View File
@@ -1,6 +1,7 @@
package com.reajason.javaweb.memshell.generator;
import com.reajason.javaweb.ClassBytesShrink;
import com.reajason.javaweb.GenerationException;
import com.reajason.javaweb.asm.ClassRenameUtils;
import com.reajason.javaweb.memshell.config.CustomConfig;
import com.reajason.javaweb.memshell.config.ShellConfig;
@@ -23,7 +24,7 @@ public class CustomShellGenerator extends ASMShellGenerator<CustomConfig> {
String shellClassBase64 = shellToolConfig.getShellClassBase64();
if (StringUtils.isBlank(shellClassBase64)) {
throw new IllegalArgumentException("Custom shell class is empty");
throw new GenerationException("Custom shell class is empty");
}
byte[] classBytes = Base64.getDecoder().decode(shellClassBase64);
byte[] bytes = ClassRenameUtils.renameClass(classBytes, shellToolConfig.getShellClassName());
generator/src/main/java/com/reajason/javaweb/probe/ProbeMethod.java
+6 -8
View File
@@ -5,6 +5,7 @@ import com.reajason.javaweb.probe.config.*;
import com.reajason.javaweb.probe.generator.DnsLogGenerator;
import com.reajason.javaweb.probe.generator.SleepGenerator;
import com.reajason.javaweb.probe.generator.response.ResponseBodyGenerator;
import lombok.SneakyThrows;
import java.lang.reflect.Constructor;
@@ -25,14 +26,11 @@ public enum ProbeMethod {
this.configClass = configClass;
}
@SneakyThrows
public byte[] generateBytes(ProbeConfig probeConfig, ProbeContentConfig probeContentConfig) {
try {
Constructor<? extends ShellGenerator> constructor =
generatorClass.getConstructor(ProbeConfig.class, configClass);
ShellGenerator generator = constructor.newInstance(probeConfig, configClass.cast(probeContentConfig));
return generator.getBytes();
} catch (Exception e) {
throw new RuntimeException("shell generate failed: " + e.getMessage(), e);
}
Constructor<? extends ShellGenerator> constructor =
generatorClass.getConstructor(ProbeConfig.class, configClass);
ShellGenerator generator = constructor.newInstance(probeConfig, configClass.cast(probeContentConfig));
return generator.getBytes();
}
}
generator/src/main/java/com/reajason/javaweb/probe/generator/DnsLogGenerator.java
+6 -1
View File
@@ -1,5 +1,6 @@
package com.reajason.javaweb.probe.generator;
import com.reajason.javaweb.GenerationException;
import com.reajason.javaweb.buddy.TargetJreVersionVisitorWrapper;
import com.reajason.javaweb.probe.ProbeContent;
import com.reajason.javaweb.probe.config.DnsLogConfig;
@@ -12,6 +13,7 @@ import com.reajason.javaweb.utils.CommonUtil;
import net.bytebuddy.ByteBuddy;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.dynamic.DynamicType;
import org.apache.commons.lang3.StringUtils;
import static net.bytebuddy.matcher.ElementMatchers.named;
@@ -27,6 +29,9 @@ public class DnsLogGenerator extends ByteBuddyShellGenerator<DnsLogConfig> {
@Override
protected DynamicType.Builder<?> build(ByteBuddy buddy) {
if (StringUtils.isBlank(probeContentConfig.getHost())) {
throw new GenerationException("DNSLog probeShell host must be specified");
}
ProbeContent detectContent = probeConfig.getProbeContent();
switch (detectContent) {
case Server:
@@ -42,7 +47,7 @@ public class DnsLogGenerator extends ByteBuddyShellGenerator<DnsLogConfig> {
.field(named("host")).value(probeContentConfig.getHost())
.visit(Advice.to(JdkProbe.class).on(named("getJdk")));
default:
throw new UnsupportedOperationException(detectContent + " not supported");
throw new GenerationException(detectContent + " not supported");
}
}
}
generator/src/main/java/com/reajason/javaweb/probe/generator/SleepGenerator.java
+9 -1
View File
@@ -1,5 +1,6 @@
package com.reajason.javaweb.probe.generator;
import com.reajason.javaweb.GenerationException;
import com.reajason.javaweb.buddy.TargetJreVersionVisitorWrapper;
import com.reajason.javaweb.probe.ProbeContent;
import com.reajason.javaweb.probe.config.ProbeConfig;
@@ -10,6 +11,7 @@ import com.reajason.javaweb.utils.CommonUtil;
import net.bytebuddy.ByteBuddy;
import net.bytebuddy.asm.Advice;
import net.bytebuddy.dynamic.DynamicType;
import org.apache.commons.lang3.StringUtils;
import static net.bytebuddy.matcher.ElementMatchers.named;
@@ -28,6 +30,12 @@ public class SleepGenerator extends ByteBuddyShellGenerator<SleepConfig> {
protected DynamicType.Builder<?> build(ByteBuddy buddy) {
ProbeContent detectContent = probeConfig.getProbeContent();
if (ProbeContent.Server.equals(detectContent)) {
if (probeContentConfig.getSeconds() <= 0) {
throw new GenerationException("sleepProbe seconds must be greater than 0");
}
if (StringUtils.isEmpty(probeContentConfig.getServer())) {
throw new GenerationException("sleepProbe server must be specified");
}
return buddy.redefine(SleepServer.class)
.name(CommonUtil.generateShellClassName())
.visit(TargetJreVersionVisitorWrapper.DEFAULT)
@@ -35,6 +43,6 @@ public class SleepGenerator extends ByteBuddyShellGenerator<SleepConfig> {
.field(named("seconds")).value(probeContentConfig.getSeconds())
.visit(Advice.to(ServerProbe.class).on(named("getServer")));
}
throw new UnsupportedOperationException("Sleep Probe not supported for " + detectContent);
throw new GenerationException("sleepProbe not supported for " + detectContent);
}
}
generator/src/main/java/com/reajason/javaweb/probe/generator/response/ResponseBodyGenerator.java
+6 -3
View File
@@ -1,5 +1,6 @@
package com.reajason.javaweb.probe.generator.response;
import com.reajason.javaweb.GenerationException;
import com.reajason.javaweb.Server;
import com.reajason.javaweb.buddy.MethodCallReplaceVisitorWrapper;
import com.reajason.javaweb.buddy.TargetJreVersionVisitorWrapper;
@@ -35,9 +36,11 @@ public class ResponseBodyGenerator extends ByteBuddyShellGenerator<ResponseBodyC
if (probeContentConfig.getReqParamName() != null) {
name = probeContentConfig.getReqParamName();
getDataFromReqInterceptor = getDataFromReqParamInterceptor.class;
} else {
} else if(probeContentConfig.getReqHeaderName() != null) {
name = probeContentConfig.getReqHeaderName();
getDataFromReqInterceptor = getDataFromReqHeaderInterceptor.class;
}else{
throw new GenerationException("responseBody probeShell must set headerName or paramName");
}
Class<?> writerClass = getWriterClass();
Class<?> runnerClass = getRunnerClass();
@@ -58,7 +61,7 @@ public class ResponseBodyGenerator extends ByteBuddyShellGenerator<ResponseBodyC
case Bytecode:
return ByteCodeProbe.class;
default:
throw new IllegalArgumentException("responseBody not supported for probe content: " + probeConfig.getProbeContent());
throw new GenerationException("responseBody not supported for probe content: " + probeConfig.getProbeContent());
}
}
@@ -86,7 +89,7 @@ public class ResponseBodyGenerator extends ByteBuddyShellGenerator<ResponseBodyC
case Server.Apusic:
return ApusicWriter.class;
default:
throw new IllegalArgumentException("responseBody now supported for server: " + probeContentConfig.getServer());
throw new GenerationException("responseBody not supported for server: " + probeContentConfig.getServer());
}
}