feat: pcap 文件增加 comment 信息

5 files changed, 87 insertions, 11 deletions

diff --git a/src/main/java/net/geedge/asw/module/feign/FeignClientConfiguration.java b/src/main/java/net/geedge/asw/module/feign/FeignClientConfiguration.java
index 5e18154..929cc00 100644
--- a/src/main/java/net/geedge/asw/module/feign/FeignClientConfiguration.java
+++ b/src/main/java/net/geedge/asw/module/feign/FeignClientConfiguration.java
@@ -4,10 +4,7 @@ import cn.hutool.core.net.url.UrlBuilder;
 import cn.hutool.log.Log;
 import feign.Feign;
 import feign.form.FormEncoder;
-import net.geedge.asw.module.feign.client.GeoipClient;
-import net.geedge.asw.module.feign.client.KibanaClient;
-import net.geedge.asw.module.feign.client.WebSharkClient;
-import net.geedge.asw.module.feign.client.ZeekClient;
+import net.geedge.asw.module.feign.client.*;
 import net.geedge.asw.module.feign.support.Fastjson2Decoder;
 import net.geedge.asw.module.feign.support.Fastjson2Encoder;
 import net.geedge.asw.module.feign.support.Http2Client;
@@ -32,6 +29,9 @@ public class FeignClientConfiguration {
     @Value("${webShark.url:127.0.0.1:8085}")
     private String websharkurl;
 
+    @Value("${pcapComment.url:127.0.0.1:5000}")
+    private String pcapCommentUrl;
+
     @Bean("zeekClient")
     public ZeekClient zeekClient() {
         String url = UrlBuilder.ofHttp(zeekUrl).toString();
@@ -76,4 +76,14 @@ public class FeignClientConfiguration {
                 .target(WebSharkClient.class, url);
     }
 
+    @Bean("pcapCommentClient")
+    public PcapCommentClient pcapCommentClient() {
+        String url = UrlBuilder.ofHttp(pcapCommentUrl).toString();
+        log.info("[pcapCommentClient] [url: {}]", url);
+        return Feign.builder()
+                .encoder(new FormEncoder())
+                .client(new Http2Client())
+                .target(PcapCommentClient.class, url);
+    }
+
 }
\ No newline at end of file
diff --git a/src/main/java/net/geedge/asw/module/feign/client/PcapCommentClient.java b/src/main/java/net/geedge/asw/module/feign/client/PcapCommentClient.java
new file mode 100644
index 0000000..b435504
--- /dev/null
+++ b/src/main/java/net/geedge/asw/module/feign/client/PcapCommentClient.java
@@ -0,0 +1,18 @@
+package net.geedge.asw.module.feign.client;
+
+import feign.Headers;
+import feign.Param;
+import feign.RequestLine;
+import feign.Response;
+import org.springframework.cloud.openfeign.FeignClient;
+
+import java.io.File;
+
+@FeignClient(name = "pcapCommentClient")
+public interface PcapCommentClient {
+
+    @RequestLine("POST /api/v1/pcap/comment")
+    @Headers("Content-Type: multipart/form-data")
+    Response addCommon(@Param("file") File file);
+
+}
\ No newline at end of file
diff --git a/src/main/java/net/geedge/asw/module/runner/service/impl/PcapServiceImpl.java b/src/main/java/net/geedge/asw/module/runner/service/impl/PcapServiceImpl.java
index 52c5c07..de85170 100644
--- a/src/main/java/net/geedge/asw/module/runner/service/impl/PcapServiceImpl.java
+++ b/src/main/java/net/geedge/asw/module/runner/service/impl/PcapServiceImpl.java
@@ -9,6 +9,7 @@ import com.alibaba.fastjson2.JSONArray;
 import com.alibaba.fastjson2.JSONObject;
 import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
 import com.baomidou.mybatisplus.core.conditions.update.LambdaUpdateWrapper;
+import com.baomidou.mybatisplus.core.toolkit.IdWorker;
 import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
 import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
 import net.geedge.asw.common.config.SpringContextUtils;
@@ -140,7 +141,12 @@ public class PcapServiceImpl extends ServiceImpl<PcapDao, PcapEntity> implements
 
         PcapEntity entity = new PcapEntity();
         try {
-            entity.setName(fileResource.getFilename());
+            String pcapId = IdWorker.get32UUID();
+            entity.setId(pcapId);
+
+            String fileExtName = T.StrUtil.emptyToDefault(T.FileUtil.extName(fileResource.getFilename()), "pcap");
+            String saveFileName = pcapId + "." + fileExtName;
+            entity.setName(saveFileName);
             entity.setDescription(description);
 
             byte[] bytes = fileResource.getInputStream().readAllBytes();
@@ -152,7 +158,7 @@ public class PcapServiceImpl extends ServiceImpl<PcapDao, PcapEntity> implements
             entity.setWorkspaceId(workspaceId);
 
             // path
-            File destination = T.FileUtil.file(T.WebPathUtil.getRootPath(), workspaceId, fileResource.getFilename());
+            File destination = T.FileUtil.file(T.WebPathUtil.getRootPath(), workspaceId, saveFileName);
             FileUtils.copyInputStreamToFile(fileResource.getInputStream(), destination);
             entity.setPath(destination.getPath());
 
diff --git a/src/main/java/net/geedge/asw/module/runner/util/PcapParserThread.java b/src/main/java/net/geedge/asw/module/runner/util/PcapParserThread.java
index 2386225..7fd22ca 100644
--- a/src/main/java/net/geedge/asw/module/runner/util/PcapParserThread.java
+++ b/src/main/java/net/geedge/asw/module/runner/util/PcapParserThread.java
@@ -4,13 +4,18 @@ import cn.hutool.log.Log;
 import com.alibaba.fastjson2.JSONArray;
 import com.alibaba.fastjson2.JSONObject;
 import com.baomidou.mybatisplus.core.conditions.update.LambdaUpdateWrapper;
+import feign.Response;
 import lombok.Data;
 import net.geedge.asw.common.config.SpringContextUtils;
+import net.geedge.asw.common.util.Constants;
+import net.geedge.asw.common.util.RCode;
 import net.geedge.asw.common.util.T;
 import net.geedge.asw.module.feign.client.GeoipClient;
+import net.geedge.asw.module.feign.client.PcapCommentClient;
 import net.geedge.asw.module.feign.client.ZeekClient;
 import net.geedge.asw.module.runner.entity.PcapEntity;
 import net.geedge.asw.module.runner.service.IPcapService;
+import org.apache.commons.io.FileUtils;
 import org.apache.commons.lang3.time.StopWatch;
 import org.opensearch.client.opensearch.OpenSearchClient;
 import org.opensearch.client.opensearch.core.BulkRequest;
@@ -21,6 +26,7 @@ import org.opensearch.client.opensearch.indices.DeleteIndexRequest;
 import org.opensearch.client.opensearch.indices.ExistsRequest;
 import org.opensearch.client.opensearch.indices.IndexSettings;
 
+import java.io.File;
 import java.io.IOException;
 import java.util.*;
 import java.util.stream.Collectors;
@@ -40,12 +46,14 @@ public class PcapParserThread implements Runnable {
     private ZeekClient zeekClient;
     private GeoipClient geoipClient;
     private OpenSearchClient openSearchClient;
+    private PcapCommentClient pcapCommentClient;
 
     private void init() {
         pcapService = SpringContextUtils.getBean(IPcapService.class);
         zeekClient = (ZeekClient) SpringContextUtils.getBean("zeekClient");
         geoipClient = (GeoipClient) SpringContextUtils.getBean("geoipClient");
         openSearchClient = (OpenSearchClient) SpringContextUtils.getBean("openSearchClient");
+        pcapCommentClient = (PcapCommentClient) SpringContextUtils.getBean("pcapCommentClient");
     }
 
     @Override
@@ -62,6 +70,8 @@ public class PcapParserThread implements Runnable {
             this.init();
             // parsing
             this.updateStatus(PcapStatus.PARSING.getValue());
+            // add common to pcap
+            this.addCommonToPcap();
             // parser
             this.parser();
             // indexed
@@ -77,6 +87,40 @@ public class PcapParserThread implements Runnable {
     }
 
     /**
+     * add common to pacp
+     */
+    private void addCommonToPcap() throws IOException {
+        // pcap common
+        Response response = pcapCommentClient.addCommon(T.FileUtil.file(pcapEntity.getPath()));
+
+        if (log.isDebugEnabled()) {
+            log.debug("[addCommonToPcap] [response: {}]", null != response ? response.status() : RCode.ERROR.getCode());
+        }
+
+        if (null == response || 200 != response.status()) {
+            throw new RuntimeException("Failed to add common to pcap");
+        }
+
+        // temp file
+        File file = T.FileUtil.file(Constants.TEMP_PATH, pcapEntity.getId() + ".pcapng");
+        FileUtils.copyInputStreamToFile(response.body().asInputStream(), file);
+
+        // override source file
+        T.FileUtil.del(T.FileUtil.file(pcapEntity.getPath()));
+        File updatePcapFile = T.FileUtil.file(T.WebPathUtil.getRootPath(), pcapEntity.getWorkspaceId(), file.getName());
+        T.FileUtil.move(file, updatePcapFile, true);
+
+        // update name,path,md5,size
+        String md5Hex = T.DigestUtil.md5Hex(updatePcapFile);
+        long filesize = T.FileUtil.size(updatePcapFile);
+        pcapEntity.setName(updatePcapFile.getName());
+        pcapEntity.setPath(updatePcapFile.getAbsolutePath());
+        pcapEntity.setMd5(md5Hex);
+        pcapEntity.setSize(filesize);
+        pcapService.updateById(pcapEntity);
+    }
+
+    /**
      * parser
      */
     private void parser() {
@@ -96,7 +140,7 @@ public class PcapParserThread implements Runnable {
 
         // add custom field
         String pcapId = pcapEntity.getId();
-        String pcapName = T.FileUtil.getName(pcapEntity.getPath());
+        String pcapName = pcapEntity.getName();
         Long tcpStream = 0L, udpStream = 0L;
 
         String sharkdApiHostAddr = properties.getProperty("sharkdApiHostAddr", "127.0.0.1");
@@ -200,9 +244,7 @@ public class PcapParserThread implements Runnable {
      * @param jsonArray
      */
     private void uploadToOpenSearch(JSONArray jsonArray) {
-        String pcapPath = pcapEntity.getPath();
-        String md5Hex = T.DigestUtil.md5Hex(T.FileUtil.file(pcapPath));
-
+        String md5Hex = pcapEntity.getMd5();
         String workspaceName = pcapEntity.getWorkspace().getName();
         String indexName = String.format("workspace-%s-%s", workspaceName, md5Hex);
 
diff --git a/src/main/resources/db/migration/V1.0.01__INIT_TABLES.sql b/src/main/resources/db/migration/V1.0.01__INIT_TABLES.sql
index e5ba4f8..c62ee31 100644
--- a/src/main/resources/db/migration/V1.0.01__INIT_TABLES.sql
+++ b/src/main/resources/db/migration/V1.0.01__INIT_TABLES.sql
@@ -229,7 +229,7 @@ CREATE TABLE `pcap`  (
     `id` varchar(64) NOT NULL COMMENT '主键',
     `name` varchar(256) NOT NULL DEFAULT '' COMMENT '文件名称',
     `description` text NOT NULL DEFAULT '' COMMENT '描述信息',
-    `path` varchar(64) NOT NULL DEFAULT '' COMMENT 'PCAP文件路径',
+    `path` varchar(256) NOT NULL DEFAULT '' COMMENT 'PCAP文件路径',
     `size` bigint(20) NOT NULL DEFAULT 0 COMMENT '文件大小',
     `md5` varchar(64) NOT NULL DEFAULT '' COMMENT '摘要值,根据文件md5值判断是否已上存在，存在则响应当前id',
     `status` varchar(64) NOT NULL DEFAULT '' COMMENT '状态,可选值 Uploaded，Analyzing，Completed',