diff --git a/src/main/java/net/geedge/asw/common/util/Constants.java b/src/main/java/net/geedge/asw/common/util/Constants.java index f609987..42c2707 100644 --- a/src/main/java/net/geedge/asw/common/util/Constants.java +++ b/src/main/java/net/geedge/asw/common/util/Constants.java @@ -75,4 +75,6 @@ public class Constants { public static final String ENV_API_WEBSOCKET_PATH = "/api/v1/env/websocket"; public static final String ENV_API_TCPDUMP_PATH = "/api/v1/env/pcap"; + + public static final String EMPTY_FILE_MD5 = "d41d8cd98f00b204e9800998ecf8427e"; } diff --git a/src/main/java/net/geedge/asw/module/environment/controller/EnvironmentController.java b/src/main/java/net/geedge/asw/module/environment/controller/EnvironmentController.java index e8025f7..f088327 100644 --- a/src/main/java/net/geedge/asw/module/environment/controller/EnvironmentController.java +++ b/src/main/java/net/geedge/asw/module/environment/controller/EnvironmentController.java @@ -1,7 +1,8 @@ package net.geedge.asw.module.environment.controller; import cn.dev33.satoken.stp.StpUtil; -import cn.hutool.core.util.ObjectUtil; +import cn.hutool.core.io.FileUtil; +import cn.hutool.core.util.ArrayUtil; import cn.hutool.http.HttpRequest; import cn.hutool.http.HttpResponse; import cn.hutool.json.JSONObject; @@ -164,11 +165,10 @@ public class EnvironmentController { @DeleteMapping("/{envId}/session/{sessionId}/pcap/{pcapId}") - public void stopTcpdump(@PathVariable("envId") String envId, + public R stopTcpdump(@PathVariable("envId") String envId, @PathVariable("sessionId") String sessionId, @PathVariable("pcapId") String pcapId, - @RequestParam Map param, - HttpServletResponse response) throws IOException, ServletException { + @RequestParam Map param) throws IOException, ServletException { EnvironmentSessionEntity session = environmentSessionService.getOne(new LambdaQueryWrapper().eq(EnvironmentSessionEntity::getId, sessionId).eq(EnvironmentSessionEntity::getStatus, 1)); if (T.ObjectUtil.isNull(session)){ throw new ASWException(RCode.ENVIRONMENT_SESSION_NOT_EXIST); @@ -178,7 +178,7 @@ public class EnvironmentController { throw new ASWException(RCode.ENVIRONMENT_NOT_EXIST); } // build query param - Map params = T.MapUtil.builder().put("id", pcapId).put("returnFile", true).build(); + Map params = T.MapUtil.builder().put("id", pcapId).put("returnFile", T.MapUtil.getBool(param, "savePcap")).build(); ResponseEntity responseEntity = EnvironmentUtil.stopTcpdump(environment, params); if (T.MapUtil.getBool(param, "savePcap")){ // save pcap to workspace @@ -189,10 +189,12 @@ public class EnvironmentController { String formatTime = new SimpleDateFormat("yyyyMMddHHmmss").format(new Date()); destination = T.FileUtil.file(T.WebPathUtil.getRootPath(), workspace.getId(), T.StrUtil.concat(true, pcapName, "-", formatTime, ".pcap")); } - FileOutputStream fos = new FileOutputStream(destination); - fos.write(responseEntity.getBody()); - fos.flush(); - fos.close(); + // create empty file + destination = FileUtil.touch(destination); + if (ArrayUtil.isNotEmpty(responseEntity.getBody())){ + FileOutputStream fos = new FileOutputStream(destination); + T.IoUtil.write(fos,true, responseEntity.getBody()); + } log.info("save pcap to path:{}", destination.getAbsolutePath()); // save entity PcapEntity entity = new PcapEntity(); @@ -204,12 +206,9 @@ public class EnvironmentController { entity.setCreateUserId(StpUtil.getLoginIdAsString()); entity.setWorkspaceId(workspace.getId()); entity.setPath(destination.getPath()); - entity.setMd5(T.DigestUtil.md5Hex(destination)); - + entity.setMd5(destination.length() == 0 ? Constants.EMPTY_FILE_MD5 : T.DigestUtil.md5Hex(destination)); pcapService.save(entity); - response.getWriter().write(T.JSONUtil.toJsonStr(R.ok())); - }else { - EnvironmentUtil.writeResponseWithHeaders(response, responseEntity); } + return R.ok(); } } \ No newline at end of file