[新增][知识库] 支持对敏感知识库文件进行加密 CN-1591

9 files changed, 97 insertions, 21 deletions

diff --git a/galaxy-job-executor/pom.xml b/galaxy-job-executor/pom.xml
index e803515..86b1e39 100644
--- a/galaxy-job-executor/pom.xml
+++ b/galaxy-job-executor/pom.xml
@@ -209,7 +209,7 @@
                         <JAR_FILE>${project.build.finalName}.xjar</JAR_FILE>
                     </buildArgs>
                     <imageTags>
-                        <imageTag>2.2.0.1</imageTag>
+                        <imageTag>2.2.0.3</imageTag>
                     </imageTags>
                     <resources>
                         <resource>
diff --git a/galaxy-job-executor/src/main/java/com/mesalab/executor/core/utils/JobUtil.java b/galaxy-job-executor/src/main/java/com/mesalab/executor/core/utils/JobUtil.java
index 5b7ab13..98e9881 100644
--- a/galaxy-job-executor/src/main/java/com/mesalab/executor/core/utils/JobUtil.java
+++ b/galaxy-job-executor/src/main/java/com/mesalab/executor/core/utils/JobUtil.java
@@ -11,6 +11,7 @@ import cn.hutool.core.util.CharsetUtil;
 import cn.hutool.core.util.ObjectUtil;
 import cn.hutool.core.util.ReUtil;
 import cn.hutool.core.util.StrUtil;
+import cn.hutool.crypto.CryptoException;
 import cn.hutool.json.JSONNull;
 import cn.hutool.json.JSONUtil;
 import cn.hutool.log.Log;
@@ -18,7 +19,15 @@ import com.geedgenetworks.utils.IpLookupV2;
 import com.mesalab.executor.exception.BusinessException;
 import com.xxl.job.core.log.XxlJobLogger;
 
+import javax.crypto.BadPaddingException;
+import javax.crypto.Cipher;
+import javax.crypto.IllegalBlockSizeException;
+import javax.crypto.NoSuchPaddingException;
+import javax.crypto.spec.SecretKeySpec;
 import java.io.*;
+import java.security.InvalidKeyException;
+import java.security.Key;
+import java.security.NoSuchAlgorithmException;
 import java.util.*;
 
 public class JobUtil {
@@ -27,6 +36,8 @@ public class JobUtil {
     public static final String DAY_TRUNC_REGEX = "\\$\\{now\\(\\)-\\d+DAY\\}";
     public static final String HOUR_TRUNC_REGEX = "\\$\\{now\\(\\)-\\d+HOUR\\}";
     public static final String NOW_REGEX = "\\$\\{now\\(\\)\\}";
+    private static final String ALGORITHM = "AES";
+    private static final String TRANSFORMATION = "AES";
 
     /**
      * 对执行器中日志字符串做处理防止特殊符号转义报错
@@ -274,4 +285,41 @@ public class JobUtil {
         }
         return true;
     }
+
+    public static void doCrypto(int cipherMode, String key, File inputFile, File outputFile) {
+
+        try (FileInputStream inputStream = new FileInputStream(inputFile);
+             BufferedInputStream bufferedInputStream = new BufferedInputStream(inputStream);
+             FileOutputStream outputStream = new FileOutputStream(outputFile);
+             BufferedOutputStream bufferedOutputStream = new BufferedOutputStream(outputStream)) {
+            Key secretKey = new SecretKeySpec(key.getBytes(), ALGORITHM);
+            Cipher cipher = Cipher.getInstance(TRANSFORMATION);
+            cipher.init(cipherMode, secretKey);
+            byte[] buffer = new byte[4096];
+            int bytesRead;
+            while ((bytesRead = bufferedInputStream.read(buffer)) != -1) {
+                byte[] outputBytes = cipher.update(buffer, 0, bytesRead);
+                if (outputBytes != null) {
+                    bufferedOutputStream.write(outputBytes);
+                }
+            }
+            byte[] finalBytes = cipher.doFinal();
+            if (finalBytes != null) {
+                bufferedOutputStream.write(finalBytes);
+            }
+        } catch (NoSuchPaddingException | NoSuchAlgorithmException
+                | InvalidKeyException | BadPaddingException
+                | IllegalBlockSizeException | IOException ex) {
+            throw new CryptoException("Error encrypting/decrypting file", ex);
+        }
+    }
+    public static void encrypt(String key, File inputFile, File outputFile)
+            throws CryptoException {
+        JobUtil.doCrypto(Cipher.ENCRYPT_MODE, key, inputFile, outputFile);
+    }
+
+    public static void decrypt(String key, File inputFile, File outputFile)
+            throws CryptoException {
+        JobUtil.doCrypto(Cipher.DECRYPT_MODE, key, inputFile, outputFile);
+    }
 }
diff --git a/galaxy-job-executor/src/main/java/com/mesalab/executor/jobhandler/knowledgeBase/KnowledgeBaseJob.java b/galaxy-job-executor/src/main/java/com/mesalab/executor/jobhandler/knowledgeBase/KnowledgeBaseJob.java
index be37582..b3a03a9 100644
--- a/galaxy-job-executor/src/main/java/com/mesalab/executor/jobhandler/knowledgeBase/KnowledgeBaseJob.java
+++ b/galaxy-job-executor/src/main/java/com/mesalab/executor/jobhandler/knowledgeBase/KnowledgeBaseJob.java
@@ -30,7 +30,7 @@ public class KnowledgeBaseJob {
 
     private static final String KNOWLEDGE_BASE_META_SQL = "SELECT a.knowledge_id knowledge_id,name,version,source,`status` FROM cn_knowledge_base_sync_version a LEFT JOIN cn_knowledge_base b ON a.knowledge_id=b.knowledge_id WHERE `is_published`=1";
     private static final String KNOWLEDGE_BASE_META_DELETE_SQL = "SELECT knowledge_id FROM cn_knowledge_base_sync_version WHERE knowledge_id NOT IN (SELECT knowledge_id FROM cn_knowledge_base)";
-    private String KNOWLEDGE_BASE_SQL = "SELECT * FROM {0} WHERE knowledge_id={1}";
+    private String KNOWLEDGE_BASE_SQL = "SELECT * FROM {0}";
 
     /**
      * [{
@@ -75,7 +75,7 @@ public class KnowledgeBaseJob {
                                 if ("mmdb".equals(sinkParam.get("format"))) {
                                     hasPush = knowledgeBaseService.mmdbBuildPush(builtParam(sink, dbKm, sinkParam));
                                 } else {
-                                    source.put("sql", MessageFormat.format(KNOWLEDGE_BASE_SQL, dbKm.get("source"), String.valueOf(dbKm.get("knowledge_id"))));
+                                    source.put("sql", MessageFormat.format(KNOWLEDGE_BASE_SQL, dbKm.get("source")));
                                     //推送到知识库接口
                                     hasPush = knowledgeBaseService.commonPushFile(source, builtParam(sink, dbKm, sinkParam));
                                 }
@@ -119,6 +119,7 @@ public class KnowledgeBaseJob {
         newParam.put("format", sinkParam.get("format"));
         newParam.put("category", ObjectUtils.defaultIfNull(sinkParam.get("category"),""));
         newParam.put("is_valid", dbKm.get("status"));
+        newParam.put("is_encryption", sinkParam.get("is_encryption"));
         sink.put("params", newParam);
         return sink;
     }
diff --git a/galaxy-job-executor/src/main/java/com/mesalab/executor/pojo/KnowledgeBase.java b/galaxy-job-executor/src/main/java/com/mesalab/executor/pojo/KnowledgeBase.java
index d2d74b3..9f3738c 100644
--- a/galaxy-job-executor/src/main/java/com/mesalab/executor/pojo/KnowledgeBase.java
+++ b/galaxy-job-executor/src/main/java/com/mesalab/executor/pojo/KnowledgeBase.java
@@ -1,15 +1,13 @@
 package com.mesalab.executor.pojo;
 
-import lombok.AllArgsConstructor;
-import lombok.Data;
-import lombok.NoArgsConstructor;
-import lombok.ToString;
+import lombok.*;
 
 import java.io.Serializable;
 
 @Data
 @NoArgsConstructor
 @AllArgsConstructor
+@Builder
 @ToString
 public class KnowledgeBase implements Serializable {
 
@@ -20,4 +18,5 @@ public class KnowledgeBase implements Serializable {
     private String sha256;
     private Integer is_valid = 1;
     private Integer builtin;
+    private Integer is_encryption = 0;
 }
\ No newline at end of file
diff --git a/galaxy-job-executor/src/main/java/com/mesalab/executor/service/DataTransformService.java b/galaxy-job-executor/src/main/java/com/mesalab/executor/service/DataTransformService.java
index 3305ef7..74c1bbe 100644
--- a/galaxy-job-executor/src/main/java/com/mesalab/executor/service/DataTransformService.java
+++ b/galaxy-job-executor/src/main/java/com/mesalab/executor/service/DataTransformService.java
@@ -3,6 +3,7 @@ package com.mesalab.executor.service;
 import cn.hutool.core.collection.ListUtil;
 import cn.hutool.core.util.ObjectUtil;
 import cn.hutool.core.util.StrUtil;
+import cn.hutool.json.JSONArray;
 import cn.hutool.json.JSONUtil;
 import cn.hutool.log.Log;
 import com.alibaba.fastjson.JSON;
@@ -222,8 +223,10 @@ public class DataTransformService {
         //处理   ioc=iov_value
         List<Map> transformResult = new ArrayList<>();
         for (Map map : json) {
-            List<Map> data = (List<Map>) map.get("data");
-            for (Map datum : data) {
+            for (Object key : map.keySet()) {
+                List<Map> maps = JSONUtil.toList((JSONArray) map.get(key), Map.class);
+                for (Map datum : maps) {
+
                 Object ioc_type = datum.get("ioc_type");
                 if (iocTypeFilter.contains(ioc_type)) {
                     Map extMap = JobUtil.extractMap(datum, FIELDS);
@@ -274,6 +277,7 @@ public class DataTransformService {
                     transformResult.add(extMap);
                 }
             }
+            }
         }
         logger.info("websketch transform data {}, spend time {} ms", transformResult.size(), System.currentTimeMillis() - now);
         XxlJobLogger.log("websketch transform data {}, spend time {} ms", transformResult.size(), System.currentTimeMillis() - now);
diff --git a/galaxy-job-executor/src/main/java/com/mesalab/executor/service/knowledgeBase/KnowledgeBaseService.java b/galaxy-job-executor/src/main/java/com/mesalab/executor/service/knowledgeBase/KnowledgeBaseService.java
index c625617..7ed971f 100644
--- a/galaxy-job-executor/src/main/java/com/mesalab/executor/service/knowledgeBase/KnowledgeBaseService.java
+++ b/galaxy-job-executor/src/main/java/com/mesalab/executor/service/knowledgeBase/KnowledgeBaseService.java
@@ -39,6 +39,8 @@ public class KnowledgeBaseService {
     @NacosValue(value = "${zookeeper.server}", autoRefreshed = true)
     private String zookeeperServer;
 
+    @NacosValue(value = "${storge.files.token}",autoRefreshed = true)
+    private String filesToken;
 
     public boolean commonPushFile(Map sourceParams, Map<String, Object> sinkParam) throws Exception {
         String url = String.valueOf(sinkParam.get("url"));
@@ -126,6 +128,12 @@ public class KnowledgeBaseService {
      */
     public boolean push(KnowledgeBase knowledgeBase, File tmpFile, String url) throws BusinessException {
         try {
+            if(knowledgeBase.getIs_encryption()== 1){
+                File tmpEncryptFile = new File(StringUtil.join(tmpFile.getParentFile(), File.separator, knowledgeBase.getName(), "_enc.", knowledgeBase.getFormat()));
+                JobUtil.encrypt(filesToken,tmpFile,tmpEncryptFile);
+                tmpFile = tmpEncryptFile;
+                JobUtil.infoLog("knowledgeBase encrypt" ,knowledgeBase.getName());
+            }
             String resultStr = HttpClientUtils.httpPostFile(url, tmpFile, BeanUtil.beanToMap(knowledgeBase));
             JobUtil.infoLog("knowledgeBase info {}, response {}", knowledgeBase.toString(), resultStr);
             if (!"-1".equals(resultStr)) {
@@ -160,7 +168,7 @@ public class KnowledgeBaseService {
                 script = "asn";
                 ipLocation = false;
             }
-            int executeResult = buildMMDB(knowledgeBase, script);
+            int executeResult = buildMMDB(script);
             String filename = StringUtil.join(script, "_builtin.", knowledgeBase.getFormat());
 
             if (executeResult == 0) {
@@ -183,14 +191,14 @@ public class KnowledgeBaseService {
         return pushStatus;
     }
 
-    private int buildMMDB(KnowledgeBase knowledgeBase, String script) throws Exception {
+    private int buildMMDB(String script) throws Exception {
         final String workDirCmd = "docker inspect --format '{{range.Mounts}}  {{.Source}} {{end}}' galaxy-job-executor |xargs -n 1|grep schedule-service";
         final String imagesNameCmd = "docker images --format '{{.Repository}}:{{.Tag}}' | grep '^perl:' | sort -r | head -n 1";
 
         String workDir = JobUtil.exeCmd(workDirCmd);
         String imagesName = JobUtil.exeCmd(imagesNameCmd);
 
-        final String cmd = "docker run --rm --network=host -v " + workDir + "/mmdb:"+ KNOWLEDGE_WORK_DIR +" --name perl-mmdb "+ imagesName +" sh -c 'cd "+ KNOWLEDGE_WORK_DIR +" && perl " + script + ".pl " + knowledgeBase.getKb_id() + "'";
+        final String cmd = "docker run --rm --network=host -v " + workDir + "/mmdb:"+ KNOWLEDGE_WORK_DIR +" --name perl-mmdb "+ imagesName +" sh -c 'cd "+ KNOWLEDGE_WORK_DIR +" && perl " + script + ".pl'";
         JobUtil.infoLog("ip locate run perl cmd {}", cmd);
 
         int executeResult = JobUtil.executeCommand(cmd);
diff --git a/galaxy-job-executor/src/main/resources/plugins/mmdb/cn/asn.pl b/galaxy-job-executor/src/main/resources/plugins/mmdb/cn/asn.pl
index 30b79c9..ee8f1e1 100644
--- a/galaxy-job-executor/src/main/resources/plugins/mmdb/cn/asn.pl
+++ b/galaxy-job-executor/src/main/resources/plugins/mmdb/cn/asn.pl
@@ -25,9 +25,7 @@ my $database = "cn_api";        # 数据库
 my $username = "root";          # 数据库用户名
 my $password = "111111";        # 数据库密码
 
-my @args = @ARGV;
-my $where= @args ? " WHERE knowledge_id=".$args[0]:"";
-my $sql = "SELECT ip1,ip2,asn,organization FROM `cn_ip_asn_built_in` $where";
+my $sql = "SELECT ip1,ip2,asn,organization FROM `cn_ip_asn_built_in`";
 # 驱动程序对象的句柄
 my $dsn = "DBI:mysql:database=$database:$host";
 my $outfile = 'asn_builtin.mmdb';
@@ -42,7 +40,7 @@ my $sth = $dbh->prepare("$sql");
 $sth->execute();
 
 my $record_size = 24;
-my $db_type = "ASNLOCATOR";
+my $db_type = "IPLOCATOR";
 my $ip_version = 6;
 
 GetOptions(
@@ -51,7 +49,7 @@ GetOptions(
 	'record-size=i' => \$record_size,
 ) or die "Usage: $0 --db-type TYPE --ipv VERSION --record-size SIZE\n";
 
-my $db_name = "ASNLOCATOR";
+my $db_name = "IPLOCATOR";
 my $db_description = "IP Geolocation Library";
 my %types = (
 	ASN				      => 'utf8_string',
diff --git a/galaxy-job-executor/src/main/resources/plugins/mmdb/cn/ip.pl b/galaxy-job-executor/src/main/resources/plugins/mmdb/cn/ip.pl
index 1aaf7c5..50b2955 100644
--- a/galaxy-job-executor/src/main/resources/plugins/mmdb/cn/ip.pl
+++ b/galaxy-job-executor/src/main/resources/plugins/mmdb/cn/ip.pl
@@ -16,9 +16,7 @@ my $database = "cn_api";       # 数据库
 my $username = "root";
 my $password = "111111";
 
-my @args = @ARGV;
-my $where= @args && $args[0] ? " WHERE knowledge_id=".$args[0]:"";
-my $sql = "SELECT start_ip,end_ip,latwgs,lngwgs,isp,continent,country,province,city FROM cn_ip_location_built_in $where";
+my $sql = "SELECT start_ip,end_ip,latwgs,lngwgs,isp,continent,country,province,city FROM cn_ip_location_built_in";
 my $dsn = "DBI:mysql:database=$database:$host";
 ###################数据库连接信息###################
 
diff --git a/galaxy-job-executor/src/test/java/com/mesalab/executor/test/EncryptConfigUtil.java b/galaxy-job-executor/src/test/java/com/mesalab/executor/test/EncryptConfigUtil.java
index 3153c43..039826b 100644
--- a/galaxy-job-executor/src/test/java/com/mesalab/executor/test/EncryptConfigUtil.java
+++ b/galaxy-job-executor/src/test/java/com/mesalab/executor/test/EncryptConfigUtil.java
@@ -1,11 +1,16 @@
 package com.mesalab.executor.test;
 
+import cn.hutool.core.io.FileUtil;
+import com.mesalab.executor.core.config.StorgeConfig;
+import com.mesalab.executor.core.utils.JobUtil;
 import org.jasypt.util.text.BasicTextEncryptor;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.springframework.boot.test.context.SpringBootTest;
 import org.springframework.test.context.junit4.SpringRunner;
 
+import javax.annotation.Resource;
+
 @RunWith(SpringRunner.class)
 @SpringBootTest(webEnvironment = SpringBootTest.WebEnvironment.RANDOM_PORT)
 public class EncryptConfigUtil {
@@ -21,6 +26,21 @@ public class EncryptConfigUtil {
 //        String pin = textEncryptor.decrypt("galaxy2019");
 //        System.err.println("user: "+user);
 //        System.err.println("pin: "+pin);
-        System.err.println("pin: "+textEncryptor.decrypt("vrvDJu5crIzkrALX+O8v3lSuBhwzCoVNCHkh94VuubzhXW/7z1crClpUT9AZY1jL"));
+        System.err.println("pin: "+textEncryptor.decrypt("LDEb2OekU7iZWiFw6pUYBSozVKP27r1y"));
+    }
+
+    @Resource
+    StorgeConfig storgeConfig;
+
+    @Test
+    public  void encryptFileTest() {
+        JobUtil.encrypt(storgeConfig.getFilesToken(), FileUtil.file("C:\\test\\h2.zip"),FileUtil.file("C:\\test\\test.csv.enc"));
     }
+    @Test
+    public  void decryptFileTest() {
+        JobUtil.decrypt(storgeConfig.getFilesToken(), FileUtil.file("C:\\test\\9_latest"),FileUtil.file("C:\\test\\test2.csv"));
+    }
+
+
+
 }
\ No newline at end of file