From 3ec0a33d3add1e0e44351b2a829eee949f04bcc5 Mon Sep 17 00:00:00 2001 From: zhangshuai Date: Mon, 9 Sep 2024 14:47:29 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20ASW-62=20Environment=20session=20?= =?UTF-8?q?=E5=81=9C=E6=AD=A2=E6=8D=95=E5=8C=85=E6=8E=A5=E5=8F=A3=E5=BC=80?= =?UTF-8?q?=E5=8F=91?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../net/geedge/asw/common/util/Constants.java | 2 + .../controller/EnvironmentController.java | 64 +++++++++++++++++++ .../environment/util/EnvironmentUtil.java | 53 ++++++++++++++- src/main/resources/application.yml | 6 +- 4 files changed, 120 insertions(+), 5 deletions(-) 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 ba2e5d5..f609987 100644 --- a/src/main/java/net/geedge/asw/common/util/Constants.java +++ b/src/main/java/net/geedge/asw/common/util/Constants.java @@ -73,4 +73,6 @@ public class Constants { * env api websocket path */ public static final String ENV_API_WEBSOCKET_PATH = "/api/v1/env/websocket"; + + public static final String ENV_API_TCPDUMP_PATH = "/api/v1/env/pcap"; } 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 51306c9..826c2fa 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,5 +1,7 @@ package net.geedge.asw.module.environment.controller; +import cn.dev33.satoken.stp.StpUtil; +import cn.hutool.core.util.ObjectUtil; import cn.hutool.http.HttpRequest; import cn.hutool.http.HttpResponse; import cn.hutool.json.JSONObject; @@ -15,12 +17,20 @@ import net.geedge.asw.module.environment.entity.EnvironmentSessionEntity; import net.geedge.asw.module.environment.service.IEnvironmentService; import net.geedge.asw.module.environment.service.IEnvironmentSessionService; import net.geedge.asw.module.environment.util.EnvironmentUtil; +import net.geedge.asw.module.runner.entity.PcapEntity; +import net.geedge.asw.module.runner.service.IPcapService; +import net.geedge.asw.module.runner.util.RunnerConstant; import net.geedge.asw.module.sys.service.ISysUserService; +import net.geedge.asw.module.workspace.entity.WorkspaceEntity; import net.geedge.asw.module.workspace.service.IWorkspaceService; import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.http.*; import org.springframework.web.bind.annotation.*; +import java.io.File; +import java.io.FileOutputStream; import java.io.IOException; +import java.text.SimpleDateFormat; import java.util.*; @RestController @@ -42,6 +52,9 @@ public class EnvironmentController { @Autowired private IWorkspaceService workspaceService; + @Autowired + private IPcapService pcapService; + @GetMapping("/{id}") public R detail(@PathVariable("id") String id) { EnvironmentEntity entity = environmentService.queryInfo(id); @@ -148,4 +161,55 @@ public class EnvironmentController { environmentSessionService.updateById(session); return R.ok(); } + + + @DeleteMapping("/{envId}/session/{sessionId}/pcap/{pcapId}") + public void stopTcpdump(@PathVariable("envId") String envId, + @PathVariable("sessionId") String sessionId, + @PathVariable("pcapId") String pcapId, + @RequestParam Map param, + HttpServletResponse response) 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); + } + EnvironmentEntity environment = environmentService.getById(envId); + if (T.ObjectUtil.isNull(environment)) { + throw new ASWException(RCode.ENVIRONMENT_NOT_EXIST); + } + // build query param + Map params = T.MapUtil.builder().put("id", pcapId).put("returnFile", true).build(); + ResponseEntity responseEntity = EnvironmentUtil.stopTcpdump(environment, params); + if (ObjectUtil.isNotNull(param.get("savePcap"))){ + // save pcap to workspace + WorkspaceEntity workspace = workspaceService.getById(session.getWorkspaceId()); + String pcapName = T.StrUtil.emptyToDefault(param.get("pcapName").toString(), pcapId); + File destination = T.FileUtil.file(T.WebPathUtil.getRootPath(), workspace.getId(), T.StrUtil.concat(true,pcapName, ".pcap")); + if (destination.exists()){ + 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(); + log.info("save pcap to path:{}", destination.getAbsolutePath()); + // save entity + PcapEntity entity = new PcapEntity(); + entity.setId(pcapId); + entity.setName(destination.getName()); + entity.setSize(destination.length()); + entity.setStatus(RunnerConstant.PcapStatus.UPLOADED.getValue()); + entity.setCreateTimestamp(System.currentTimeMillis()); + entity.setCreateUserId(StpUtil.getLoginIdAsString()); + entity.setWorkspaceId(workspace.getId()); + entity.setPath(destination.getPath()); + entity.setMd5(T.DigestUtil.md5Hex(destination)); + + pcapService.save(entity); + response.getWriter().write(T.JSONUtil.toJsonStr(R.ok())); + }else { + EnvironmentUtil.writeResponseWithHeaders(response, responseEntity); + } + } } \ No newline at end of file diff --git a/src/main/java/net/geedge/asw/module/environment/util/EnvironmentUtil.java b/src/main/java/net/geedge/asw/module/environment/util/EnvironmentUtil.java index 08f4b15..a725d92 100644 --- a/src/main/java/net/geedge/asw/module/environment/util/EnvironmentUtil.java +++ b/src/main/java/net/geedge/asw/module/environment/util/EnvironmentUtil.java @@ -38,9 +38,54 @@ public class EnvironmentUtil { private static Log log = Log.get(); private static RestTemplate restTemplate; + /** + * agent stop tcpdump + * @param environment + * @param params + * @return + * @throws IOException + * @throws ServletException + */ + public static ResponseEntity stopTcpdump(EnvironmentEntity environment, Map params) throws IOException, ServletException { + JSONObject jsonObject = environment.getParamJSONObject(); + String url = jsonObject.getStr("url"); + String token = jsonObject.getStr("token"); + + String urlStr = UrlBuilder.of(url) + .setPath(UrlPath.of(Constants.ENV_API_TCPDUMP_PATH, Charset.forName("UTF-8"))) + .setQuery(UrlQuery.of(params)) + .setCharset(StandardCharsets.UTF_8).toString(); + // token + HttpHeaders headers = new HttpHeaders(); + headers.add(HttpHeaders.AUTHORIZATION, token); + HttpEntity httpEntity = new HttpEntity(headers); + ResponseEntity responseEntity = null; + try { + responseEntity = restTemplate.exchange(new URI(urlStr), HttpMethod.DELETE, httpEntity, byte[].class); + } catch (Exception e) { + log.error(e, "stop tcpdump request error. url:{}", urlStr); + String message = e.getMessage(); + if (ObjectUtil.isNotNull(e.getCause())) { + message = e.getCause().getMessage(); + } + throw new ASWException(message, HttpStatus.INTERNAL_SERVER_ERROR.value()); + } + int statusCode = responseEntity.getStatusCodeValue(); + log.info("stop tcpdump request url:{}, responseStatus:{}", urlStr, statusCode); + return responseEntity; + } + + /** + * env api agent + * @param device + * @param request + * @param response + * @param sessionId + * @throws IOException + * @throws ServletException + */ public static void getForObject(EnvironmentEntity device, HttpServletRequest request, HttpServletResponse response, String sessionId) throws IOException, ServletException { // path - String pathProfix = T.StrUtil.concat(true, Constants.ENV_API_PREFIX, "/", device.getId()); String[] paths = request.getServletPath().split(sessionId); String path = Arrays.asList(paths).getLast(); path = path.startsWith("/") ? (String.format("%s%s", Constants.ENV_API_PREFIX, path)) @@ -113,11 +158,13 @@ public class EnvironmentUtil { } throw new ASWException(message, HttpStatus.INTERNAL_SERVER_ERROR.value()); } + log.info("env request url:{}, responseStatus:{}", urlStr, responseEntity.getStatusCode()); + writeResponseWithHeaders(response, responseEntity); + } + public static void writeResponseWithHeaders(HttpServletResponse response, ResponseEntity responseEntity) throws IOException { HttpHeaders httpHeaders = responseEntity.getHeaders(); int statusCode = responseEntity.getStatusCodeValue(); - log.info("env request url:{}, responseStatus:{}", urlStr, statusCode); - byte[] responseBody = responseEntity.getBody(); response.reset(); response.setStatus(statusCode); diff --git a/src/main/resources/application.yml b/src/main/resources/application.yml index 7d4597b..d7a9634 100644 --- a/src/main/resources/application.yml +++ b/src/main/resources/application.yml @@ -10,6 +10,8 @@ spring: no-cache: true # max-age: 30d # cache-public: true + jackson: + default-property-inclusion: non_empty profiles: active: prod include: magic-api @@ -26,7 +28,7 @@ spring: baseline-on-migrate: true # 连接数据库中存在表时设置为true locations: classpath:db/migration # 脚本路径 clean-disabled: false # flyway 的 clean 命令会删除指定 schema 下的所有 table, 生产务必禁掉。这个默认值是 false 理论上作为默认配置是不科学的 - validate-on-migrate: true # 执行迁移时是否自动调用验证 当你的 版本不符合逻辑 比如 你先执行了 DML 而没有 对应的DDL 会抛出异常 + validate-on-migrate: false # 执行迁移时是否自动调用验证 当你的 版本不符合逻辑 比如 你先执行了 DML 而没有 对应的DDL 会抛出异常 placeholder-replacement: false # 不做取值替换 默认替换为 ${} ，初始化sql中有sql语句存在freemarker替换，所以禁用此项 servlet: multipart: @@ -71,4 +73,4 @@ mybatis-plus: banner: false logging: - config: ./config/logback-spring.xml \ No newline at end of file + config: D:\IdeaProjects\asw-controller\src\main\resources\config\logback-spring.xml \ No newline at end of file