summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
authorgujinkai <[email protected]>2024-03-11 14:10:21 +0800
committergujinkai <[email protected]>2024-03-14 10:24:20 +0800
commit1f4ea178b36c81d82af82783fe664e4eac074585 (patch)
treea004db0b90b6f1fae47df54af0934a8e0ad50464
parent7115979e16f6e9577df9063f25d678e9cfbc5aac (diff)
[Improve][core] CN knowledge Handler add option: download from local
-rw-r--r--groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/AbstractKnowledgeBaseHandler.java3
-rw-r--r--groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/AbstractMultipleKnowledgeBaseHandler.java47
-rw-r--r--groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/AbstractSingleKnowledgeBaseHandler.java17
-rw-r--r--groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/AppCategoryKnowledgeBaseHandler.java7
-rw-r--r--groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/AppTagUserDefineKnowledgeBaseHandler.java13
-rw-r--r--groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/DnsServerInfoKnowledgeBaseHandler.java7
-rw-r--r--groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/DomainTagUserDefineKnowledgeBaseHandler.java16
-rw-r--r--groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/DomainVpnKnowledgeBaseHandler.java7
-rw-r--r--groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/FqdnCategoryKnowledgeBaseHandler.java7
-rw-r--r--groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/FqdnIcpKnowledgeBaseHandler.java7
-rw-r--r--groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/FqdnWhoisKnowledgeBaseHandler.java7
-rw-r--r--groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/IdcRenterKnowledgeBaseHandler.java7
-rw-r--r--groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/InternalIpKnowledgeBaseHandler.java7
-rw-r--r--groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/IocDarkwebKnowledgeBaseHandler.java9
-rw-r--r--groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/IocMalwareKnowledgeBaseHandler.java10
-rw-r--r--groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/IpTagUserDefineKnowledgeBaseHandler.java15
-rw-r--r--groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/IpVpnKnowledgeBaseHandler.java8
-rw-r--r--groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/LinkDirectionKnowledgeBaseHandler.java8
-rw-r--r--groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/RuleKnowledgeBaseHandler.java28
-rw-r--r--groot-core/src/test/java/com/geedgenetworks/core/udf/cn/LookupTestUtils.java12
20 files changed, 197 insertions, 45 deletions
diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/AbstractKnowledgeBaseHandler.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/AbstractKnowledgeBaseHandler.java
index 645cdb1..dd95198 100644
--- a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/AbstractKnowledgeBaseHandler.java
+++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/AbstractKnowledgeBaseHandler.java
@@ -4,7 +4,6 @@ import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONArray;
import com.geedgenetworks.common.config.KnowledgeBaseConfig;
import com.geedgenetworks.common.utils.FileUtils;
-import com.geedgenetworks.common.utils.HdfsUtils;
import com.geedgenetworks.core.pojo.KnowLedgeBaseFileMeta;
import com.geedgenetworks.core.utils.HttpClientPoolUtil;
import lombok.extern.slf4j.Slf4j;
@@ -16,8 +15,6 @@ import java.util.ArrayList;
import java.util.List;
import java.util.Map;
-import static com.geedgenetworks.common.utils.FileUtils.getFileFromDir;
-
@Slf4j
public abstract class AbstractKnowledgeBaseHandler {
diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/AbstractMultipleKnowledgeBaseHandler.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/AbstractMultipleKnowledgeBaseHandler.java
index 548e40a..39be307 100644
--- a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/AbstractMultipleKnowledgeBaseHandler.java
+++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/AbstractMultipleKnowledgeBaseHandler.java
@@ -2,6 +2,8 @@ package com.geedgenetworks.core.udf.knowlegdebase.handler;
import com.alibaba.fastjson2.JSON;
import com.geedgenetworks.common.config.KnowledgeBaseConfig;
+import com.geedgenetworks.common.exception.CommonErrorCode;
+import com.geedgenetworks.common.exception.GrootStreamRuntimeException;
import com.geedgenetworks.core.pojo.KnowLedgeBaseFileMeta;
import lombok.AllArgsConstructor;
import lombok.Data;
@@ -38,7 +40,9 @@ public abstract class AbstractMultipleKnowledgeBaseHandler extends AbstractKnowl
@Override
public boolean initKnowledgeBase(KnowledgeBaseConfig knowledgeBaseConfig) {
this.knowledgeBaseConfig = knowledgeBaseConfig;
- this.knowledgeMetedataCacheMap = getMetadata(knowledgeBaseConfig.getFsPath());
+ if ("http".equals(knowledgeBaseConfig.getFsType())) {
+ this.knowledgeMetedataCacheMap = getMetadata(knowledgeBaseConfig.getFsPath());
+ }
return buildKnowledgeBase();
}
@@ -51,6 +55,26 @@ public abstract class AbstractMultipleKnowledgeBaseHandler extends AbstractKnowl
protected abstract Boolean buildKnowledgeBase();
+ protected byte[] downloadFile(Long id) {
+ if ("http".equals(knowledgeBaseConfig.getFsType())) {
+ return downloadFile(knowledgeMetedataCacheMap.get(encodeId(id)).getPath(), 1);
+ }
+ if ("local".equals(knowledgeBaseConfig.getFsType())) {
+ return getFileFromLocal(knowledgeBaseConfig.getFsPath() + id);
+ }
+ throw new GrootStreamRuntimeException(CommonErrorCode.ILLEGAL_ARGUMENT, knowledgeBaseConfig.getFsType() + " is illegal");
+ }
+
+ protected List<Long> getAllFileIds() {
+ if ("http".equals(knowledgeBaseConfig.getFsType())) {
+ return knowledgeMetedataCacheMap.keySet().stream().map(AbstractMultipleKnowledgeBaseHandler::decodeId).collect(Collectors.toList());
+ }
+ if ("local".equals(knowledgeBaseConfig.getFsType())) {
+ return knowledgeBaseConfig.getFiles().stream().map(AbstractMultipleKnowledgeBaseHandler::decodeId).collect(Collectors.toList());
+ }
+ throw new GrootStreamRuntimeException(CommonErrorCode.ILLEGAL_ARGUMENT, knowledgeBaseConfig.getFsType() + " is illegal");
+ }
+
protected Boolean ifNeedUpdate() {
Map<String, KnowLedgeBaseFileMeta> knowledgeMetedataMap = getMetadata(knowledgeBaseConfig.getFsPath());
if (knowledgeMetedataMap.size() != knowledgeMetedataCacheMap.size()) {
@@ -76,7 +100,7 @@ public abstract class AbstractMultipleKnowledgeBaseHandler extends AbstractKnowl
KnowledgeResponse knowledgeResponse = JSON.parseObject(content, KnowledgeResponse.class);
List<KnowLedgeBaseFileMeta> knowledgeMetedataList = JSON.parseArray(knowledgeResponse.data, KnowLedgeBaseFileMeta.class);
return knowledgeMetedataList.stream()
- .filter(metadata -> "latest".equals(metadata.getVersion()) && metadata.getIsValid() == 1 && metadata.getSha256() != null)
+ .filter(metadata -> "latest".equals(metadata.getVersion()) && metadata.getIsValid() == 1 && metadata.getSha256() != null && checkId(metadata.getKb_id()))
.collect(Collectors.toMap(KnowLedgeBaseFileMeta::getKb_id, Function.identity(), (existing, replacement) -> existing, HashMap::new));
}
} catch (IOException e) {
@@ -85,6 +109,16 @@ public abstract class AbstractMultipleKnowledgeBaseHandler extends AbstractKnowl
return new HashMap<>();
}
+ public static boolean checkId(String id) {
+ try {
+ decodeId(id);
+ return true;
+ } catch (Exception e) {
+ logger.error("knowledge id:" + id + " is illegal", e);
+ return false;
+ }
+ }
+
@Data
private static final class KnowledgeResponse {
private int status;
@@ -99,15 +133,20 @@ public abstract class AbstractMultipleKnowledgeBaseHandler extends AbstractKnowl
* @param id
* @return
*/
- protected Long convertId(String id) {
+ public static Long decodeId(String id) {
return Long.parseLong(id);
}
+ public static String encodeId(Long id) {
+ return String.valueOf(id);
+ }
+
@Data
@AllArgsConstructor
public static final class Node {
private String tag;
private Long kbId;
-
}
+
+ public abstract void close();
}
diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/AbstractSingleKnowledgeBaseHandler.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/AbstractSingleKnowledgeBaseHandler.java
index 1691b88..cd831de 100644
--- a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/AbstractSingleKnowledgeBaseHandler.java
+++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/AbstractSingleKnowledgeBaseHandler.java
@@ -1,6 +1,8 @@
package com.geedgenetworks.core.udf.knowlegdebase.handler;
import com.alibaba.fastjson2.JSON;
import com.geedgenetworks.common.config.KnowledgeBaseConfig;
+import com.geedgenetworks.common.exception.CommonErrorCode;
+import com.geedgenetworks.common.exception.GrootStreamRuntimeException;
import com.geedgenetworks.core.pojo.KnowLedgeBaseFileMeta;
import lombok.Data;
import org.apache.http.HttpEntity;
@@ -34,7 +36,9 @@ public abstract class AbstractSingleKnowledgeBaseHandler extends AbstractKnowled
@Override
public boolean initKnowledgeBase(KnowledgeBaseConfig knowledgeBaseConfig) {
this.knowledgeBaseConfig = knowledgeBaseConfig;
- this.knowledgeMetedataCache = getMetadata(knowledgeBaseConfig.getFsType(), knowledgeBaseConfig.getFsPath(), knowledgeBaseConfig.getFiles().get(0));
+ if ("http".equals(knowledgeBaseConfig.getFsType())) {
+ this.knowledgeMetedataCache = getMetadata(knowledgeBaseConfig.getFsType(), knowledgeBaseConfig.getFsPath(), knowledgeBaseConfig.getFiles().get(0));
+ }
return buildKnowledgeBase();
}
@@ -47,6 +51,16 @@ public abstract class AbstractSingleKnowledgeBaseHandler extends AbstractKnowled
protected abstract Boolean buildKnowledgeBase();
+ public byte[] downloadFile() {
+ if ("http".equals(knowledgeBaseConfig.getFsType())) {
+ return downloadFile(knowledgeMetedataCache.getPath(), 1);
+ }
+ if ("local".equals(knowledgeBaseConfig.getFsType())) {
+ return getFileFromLocal(knowledgeBaseConfig.getFsPath() + knowledgeBaseConfig.getFiles().get(0));
+ }
+ throw new GrootStreamRuntimeException(CommonErrorCode.ILLEGAL_ARGUMENT, knowledgeBaseConfig.getFsType() + " is illegal");
+ }
+
protected Boolean ifNeedUpdate() {
KnowLedgeBaseFileMeta knowledgeMetedata = getMetadata(knowledgeBaseConfig.getFsType(), knowledgeBaseConfig.getFsPath(), knowledgeBaseConfig.getFiles().get(0));
if (knowledgeMetedata == null || knowledgeMetedata.getSha256() == null) {
@@ -89,4 +103,5 @@ public abstract class AbstractSingleKnowledgeBaseHandler extends AbstractKnowled
private String data;
}
+ public abstract void close();
}
diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/AppCategoryKnowledgeBaseHandler.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/AppCategoryKnowledgeBaseHandler.java
index cf579fd..c66a86a 100644
--- a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/AppCategoryKnowledgeBaseHandler.java
+++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/AppCategoryKnowledgeBaseHandler.java
@@ -43,7 +43,7 @@ public class AppCategoryKnowledgeBaseHandler extends AbstractSingleKnowledgeBase
needColumns.add("app_subcategory");
needColumns.add("app_company");
needColumns.add("app_company_category");
- byte[] content = downloadFile(knowledgeMetedataCache.getPath(), 1);
+ byte[] content = downloadFile();
HighCsvReader highCsvReader = new HighCsvReader(new InputStreamReader(new ByteArrayInputStream(content)), needColumns);
HashMap<String, AppCategory> newAppMap = new HashMap<>((int) (highCsvReader.getLineNumber() / 0.75F + 1.0F));
HighCsvReader.CsvIterator iterator = highCsvReader.getIterator();
@@ -88,4 +88,9 @@ public class AppCategoryKnowledgeBaseHandler extends AbstractSingleKnowledgeBase
private String appCompany;
private String appCompanyCategory;
}
+
+ public void close() {
+ appMap.clear();
+ appMap = null;
+ }
}
diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/AppTagUserDefineKnowledgeBaseHandler.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/AppTagUserDefineKnowledgeBaseHandler.java
index 0ff0183..d09c11e 100644
--- a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/AppTagUserDefineKnowledgeBaseHandler.java
+++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/AppTagUserDefineKnowledgeBaseHandler.java
@@ -1,6 +1,5 @@
package com.geedgenetworks.core.udf.knowlegdebase.handler;
-import com.geedgenetworks.core.pojo.KnowLedgeBaseFileMeta;
import com.geedgenetworks.core.utils.cn.csv.HighCsvReader;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -37,19 +36,17 @@ public class AppTagUserDefineKnowledgeBaseHandler extends AbstractMultipleKnowle
@Override
protected Boolean buildKnowledgeBase() {
Map<String, List<Node>> newAppTagRules = new HashMap<>();
- this.knowledgeMetedataCacheMap.forEach((key, value) -> {
- buildSingleKnowledgeBase(newAppTagRules, convertId(key), value);
- });
+ getAllFileIds().forEach(id -> buildSingleKnowledgeBase(newAppTagRules, id));
appTagRules = newAppTagRules;
return true;
}
- private void buildSingleKnowledgeBase(Map<String, List<Node>> appTags, Long id, KnowLedgeBaseFileMeta metadata) {
+ private void buildSingleKnowledgeBase(Map<String, List<Node>> appTags, Long id) {
try {
List<String> needColumns = new ArrayList<>();
needColumns.add("tag_value");
needColumns.add("app_name");
- byte[] content = downloadFile(metadata.getPath(), 1);
+ byte[] content = downloadFile(id);
HighCsvReader highCsvReader = new HighCsvReader(new InputStreamReader(new ByteArrayInputStream(content)), needColumns);
HighCsvReader.CsvIterator iterator = highCsvReader.getIterator();
while (iterator.hasNext()) {
@@ -79,4 +76,8 @@ public class AppTagUserDefineKnowledgeBaseHandler extends AbstractMultipleKnowle
return new ArrayList<>();
}
+ public void close() {
+ appTagRules.clear();
+ appTagRules = null;
+ }
}
diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/DnsServerInfoKnowledgeBaseHandler.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/DnsServerInfoKnowledgeBaseHandler.java
index 1f7196a..bb3a8c9 100644
--- a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/DnsServerInfoKnowledgeBaseHandler.java
+++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/DnsServerInfoKnowledgeBaseHandler.java
@@ -39,7 +39,7 @@ public class DnsServerInfoKnowledgeBaseHandler extends AbstractSingleKnowledgeBa
List<String> needColumns = new ArrayList<>();
needColumns.add("ip_addr");
needColumns.add("dns_server_role");
- byte[] content = downloadFile(knowledgeMetedataCache.getPath(), 1);
+ byte[] content = downloadFile();
HighCsvReader highCsvReader = new HighCsvReader(new InputStreamReader(new ByteArrayInputStream(content)), needColumns);
HashMap<String, List<String>> newDnsMap = new HashMap<>((int) (highCsvReader.getLineNumber() / 0.75F + 1.0F));
HighCsvReader.CsvIterator iterator = highCsvReader.getIterator();
@@ -64,4 +64,9 @@ public class DnsServerInfoKnowledgeBaseHandler extends AbstractSingleKnowledgeBa
public List<String> lookup(String ip) {
return dnsMap.get(ip);
}
+
+ public void close() {
+ dnsMap.clear();
+ dnsMap = null;
+ }
}
diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/DomainTagUserDefineKnowledgeBaseHandler.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/DomainTagUserDefineKnowledgeBaseHandler.java
index 37abe9e..85e1ad4 100644
--- a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/DomainTagUserDefineKnowledgeBaseHandler.java
+++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/DomainTagUserDefineKnowledgeBaseHandler.java
@@ -1,6 +1,5 @@
package com.geedgenetworks.core.udf.knowlegdebase.handler;
-import com.geedgenetworks.core.pojo.KnowLedgeBaseFileMeta;
import com.geedgenetworks.core.utils.cn.common.Trie;
import com.geedgenetworks.core.utils.cn.csv.HighCsvReader;
import org.apache.commons.lang3.StringUtils;
@@ -41,20 +40,18 @@ public class DomainTagUserDefineKnowledgeBaseHandler extends AbstractMultipleKno
protected Boolean buildKnowledgeBase() {
Trie<Node> newFqdnTagFuzzyMatchRules = new Trie<>();
Map<String, List<Node>> newFqdnTagFullMatchRules = new HashMap<>();
- this.knowledgeMetedataCacheMap.forEach((key, value) -> {
- buildSingleKnowledgeBase(newFqdnTagFuzzyMatchRules, newFqdnTagFullMatchRules, convertId(key), value);
- });
+ getAllFileIds().forEach(id -> buildSingleKnowledgeBase(newFqdnTagFuzzyMatchRules, newFqdnTagFullMatchRules, id));
fqdnTagFuzzyMatchRules = newFqdnTagFuzzyMatchRules;
fqdnTagFullMatchRules = newFqdnTagFullMatchRules;
return true;
}
- private void buildSingleKnowledgeBase(Trie<Node> fqdnTagsFuzzy, Map<String, List<Node>> fqdnTagsFull, Long id, KnowLedgeBaseFileMeta metadata) {
+ private void buildSingleKnowledgeBase(Trie<Node> fqdnTagsFuzzy, Map<String, List<Node>> fqdnTagsFull, Long id) {
try {
List<String> needColumns = new ArrayList<>();
needColumns.add("tag_value");
needColumns.add("domain");
- byte[] content = downloadFile(metadata.getPath(), 1);
+ byte[] content = downloadFile(id);
HighCsvReader highCsvReader = new HighCsvReader(new InputStreamReader(new ByteArrayInputStream(content)), needColumns);
HighCsvReader.CsvIterator iterator = highCsvReader.getIterator();
while (iterator.hasNext()) {
@@ -102,4 +99,11 @@ public class DomainTagUserDefineKnowledgeBaseHandler extends AbstractMultipleKno
}
return nodes;
}
+
+ @Override
+ public void close() {
+ fqdnTagFuzzyMatchRules = null;
+ fqdnTagFullMatchRules.clear();
+ fqdnTagFullMatchRules = null;
+ }
}
diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/DomainVpnKnowledgeBaseHandler.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/DomainVpnKnowledgeBaseHandler.java
index 2eada2a..a8fa70c 100644
--- a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/DomainVpnKnowledgeBaseHandler.java
+++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/DomainVpnKnowledgeBaseHandler.java
@@ -39,7 +39,7 @@ public class DomainVpnKnowledgeBaseHandler extends AbstractSingleKnowledgeBaseHa
List<String> needColumns = new ArrayList<>();
needColumns.add("domain");
needColumns.add("vpn_service_name");
- byte[] content = downloadFile(knowledgeMetedataCache.getPath(), 1);
+ byte[] content = downloadFile();
HighCsvReader highCsvReader = new HighCsvReader(new InputStreamReader(new ByteArrayInputStream(content)), needColumns);
Map<String, String> newMap = new HashMap<>((int) (highCsvReader.getLineNumber() / 0.75F + 1.0F));
HighCsvReader.CsvIterator iterator = highCsvReader.getIterator();
@@ -77,4 +77,9 @@ public class DomainVpnKnowledgeBaseHandler extends AbstractSingleKnowledgeBaseHa
}
}
}
+
+ public void close() {
+ map.clear();
+ map = null;
+ }
}
diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/FqdnCategoryKnowledgeBaseHandler.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/FqdnCategoryKnowledgeBaseHandler.java
index ba7d166..a70fdf8 100644
--- a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/FqdnCategoryKnowledgeBaseHandler.java
+++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/FqdnCategoryKnowledgeBaseHandler.java
@@ -42,7 +42,7 @@ public class FqdnCategoryKnowledgeBaseHandler extends AbstractSingleKnowledgeBas
needColumns.add("reputation_level");
needColumns.add("category_name");
needColumns.add("category_group");
- byte[] content = downloadFile(knowledgeMetedataCache.getPath(), 1);
+ byte[] content = downloadFile();
HighCsvReader highCsvReader = new HighCsvReader(new InputStreamReader(new ByteArrayInputStream(content)), needColumns);
HashMap<String, FqdnCategory> newCategoryMap = new HashMap<>((int) (highCsvReader.getLineNumber() / 0.75F + 1.0F));
HighCsvReader.CsvIterator iterator = highCsvReader.getIterator();
@@ -96,4 +96,9 @@ public class FqdnCategoryKnowledgeBaseHandler extends AbstractSingleKnowledgeBas
private String categoryGroup;
private String reputationLevel;
}
+
+ public void close() {
+ categoryMap.clear();
+ categoryMap = null;
+ }
}
diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/FqdnIcpKnowledgeBaseHandler.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/FqdnIcpKnowledgeBaseHandler.java
index b983ac0..8891a46 100644
--- a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/FqdnIcpKnowledgeBaseHandler.java
+++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/FqdnIcpKnowledgeBaseHandler.java
@@ -38,7 +38,7 @@ public class FqdnIcpKnowledgeBaseHandler extends AbstractSingleKnowledgeBaseHand
List<String> needColumns = new ArrayList<>();
needColumns.add("fqdn");
needColumns.add("icp_company_name");
- byte[] content = downloadFile(knowledgeMetedataCache.getPath(), 1);
+ byte[] content = downloadFile();
HighCsvReader highCsvReader = new HighCsvReader(new InputStreamReader(new ByteArrayInputStream(content)), needColumns);
HashMap<String, String> newIcpMap = new HashMap<>((int) (highCsvReader.getLineNumber() / 0.75F + 1.0F));
HighCsvReader.CsvIterator iterator = highCsvReader.getIterator();
@@ -76,4 +76,9 @@ public class FqdnIcpKnowledgeBaseHandler extends AbstractSingleKnowledgeBaseHand
}
}
}
+
+ public void close() {
+ icpMap.clear();
+ icpMap = null;
+ }
}
diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/FqdnWhoisKnowledgeBaseHandler.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/FqdnWhoisKnowledgeBaseHandler.java
index b9a2457..2803d10 100644
--- a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/FqdnWhoisKnowledgeBaseHandler.java
+++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/FqdnWhoisKnowledgeBaseHandler.java
@@ -39,7 +39,7 @@ public class FqdnWhoisKnowledgeBaseHandler extends AbstractSingleKnowledgeBaseHa
List<String> needColumns = new ArrayList<>();
needColumns.add("fqdn");
needColumns.add("whois_registrant_org");
- byte[] content = downloadFile(knowledgeMetedataCache.getPath(), 1);
+ byte[] content = downloadFile();
HighCsvReader highCsvReader = new HighCsvReader(new InputStreamReader(new ByteArrayInputStream(content)), needColumns);
HashMap<String, String> newWhoIsMap = new HashMap<>((int) (highCsvReader.getLineNumber() / 0.75F + 1.0F));
HighCsvReader.CsvIterator iterator = highCsvReader.getIterator();
@@ -77,4 +77,9 @@ public class FqdnWhoisKnowledgeBaseHandler extends AbstractSingleKnowledgeBaseHa
}
}
}
+
+ public void close() {
+ whoIsMap.clear();
+ whoIsMap = null;
+ }
}
diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/IdcRenterKnowledgeBaseHandler.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/IdcRenterKnowledgeBaseHandler.java
index d7c7f1f..50127bb 100644
--- a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/IdcRenterKnowledgeBaseHandler.java
+++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/IdcRenterKnowledgeBaseHandler.java
@@ -44,7 +44,7 @@ public class IdcRenterKnowledgeBaseHandler extends AbstractSingleKnowledgeBaseHa
needColumns.add("ip1");
needColumns.add("ip2");
needColumns.add("idc_renter");
- byte[] content = downloadFile(knowledgeMetedataCache.getPath(), 1);
+ byte[] content = downloadFile();
HighCsvReader highCsvReader = new HighCsvReader(new InputStreamReader(new ByteArrayInputStream(content)), needColumns);
TreeRangeMap<IPAddress, String> newTreeRangeMap = TreeRangeMap.create();
HighCsvReader.CsvIterator iterator = highCsvReader.getIterator();
@@ -103,4 +103,9 @@ public class IdcRenterKnowledgeBaseHandler extends AbstractSingleKnowledgeBaseHa
}
return null;
}
+
+ public void close() {
+ treeRangeMap.clear();
+ treeRangeMap = null;
+ }
}
diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/InternalIpKnowledgeBaseHandler.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/InternalIpKnowledgeBaseHandler.java
index 4faefc5..4d6ad3e 100644
--- a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/InternalIpKnowledgeBaseHandler.java
+++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/InternalIpKnowledgeBaseHandler.java
@@ -43,7 +43,7 @@ public class InternalIpKnowledgeBaseHandler extends AbstractSingleKnowledgeBaseH
needColumns.add("addr_format");
needColumns.add("ip1");
needColumns.add("ip2");
- byte[] content = downloadFile(knowledgeMetedataCache.getPath(), 1);
+ byte[] content = downloadFile();
HighCsvReader highCsvReader = new HighCsvReader(new InputStreamReader(new ByteArrayInputStream(content)), needColumns);
TreeRangeSet<IPAddress> newTreeRangeSet = TreeRangeSet.create();
HighCsvReader.CsvIterator iterator = highCsvReader.getIterator();
@@ -107,4 +107,9 @@ public class InternalIpKnowledgeBaseHandler extends AbstractSingleKnowledgeBaseH
}
return false;
}
+
+ public void close() {
+ treeRangeSet.clear();
+ treeRangeSet = null;
+ }
}
diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/IocDarkwebKnowledgeBaseHandler.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/IocDarkwebKnowledgeBaseHandler.java
index 859206d..0d0dec4 100644
--- a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/IocDarkwebKnowledgeBaseHandler.java
+++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/IocDarkwebKnowledgeBaseHandler.java
@@ -41,7 +41,7 @@ public class IocDarkwebKnowledgeBaseHandler extends AbstractSingleKnowledgeBaseH
needColumns.add("ioc_type");
needColumns.add("ioc_value");
needColumns.add("node_type");
- byte[] content = downloadFile(knowledgeMetedataCache.getPath(), 1);
+ byte[] content = downloadFile();
HighCsvReader highCsvReader = new HighCsvReader(new InputStreamReader(new ByteArrayInputStream(content)), needColumns);
HashMap<String, String> newIpMap = new HashMap<>((int) (highCsvReader.getLineNumber() / 0.75F + 1.0F));
HashMap<String, String> newDomainMap = new HashMap<>((int) (highCsvReader.getLineNumber() / 0.75F + 1.0F));
@@ -90,4 +90,11 @@ public class IocDarkwebKnowledgeBaseHandler extends AbstractSingleKnowledgeBaseH
}
}
}
+
+ public void close() {
+ ipMap.clear();
+ ipMap = null;
+ domainMap.clear();
+ domainMap = null;
+ }
}
diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/IocMalwareKnowledgeBaseHandler.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/IocMalwareKnowledgeBaseHandler.java
index 94dc628..8e6f83b 100644
--- a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/IocMalwareKnowledgeBaseHandler.java
+++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/IocMalwareKnowledgeBaseHandler.java
@@ -43,7 +43,7 @@ public class IocMalwareKnowledgeBaseHandler extends AbstractSingleKnowledgeBaseH
needColumns.add("ioc_type");
needColumns.add("ioc_value");
needColumns.add("malware_name");
- byte[] content = downloadFile(knowledgeMetedataCache.getPath(), 1);
+ byte[] content = downloadFile();
HighCsvReader highCsvReader = new HighCsvReader(new InputStreamReader(new ByteArrayInputStream(content)), needColumns);
HashMap<String, String> newIpMap = new HashMap<>((int) (highCsvReader.getLineNumber() / 0.75F + 1.0F));
HashMap<String, String> newDomainMap = new HashMap<>((int) (highCsvReader.getLineNumber() / 0.75F + 1.0F));
@@ -104,4 +104,12 @@ public class IocMalwareKnowledgeBaseHandler extends AbstractSingleKnowledgeBaseH
}
return list.get(0);
}
+
+ public void close() {
+ ipMap.clear();
+ ipMap = null;
+ domainMap.clear();
+ domainMap = null;
+ urlTrie = null;
+ }
}
diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/IpTagUserDefineKnowledgeBaseHandler.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/IpTagUserDefineKnowledgeBaseHandler.java
index d5a4a13..3b2f1e3 100644
--- a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/IpTagUserDefineKnowledgeBaseHandler.java
+++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/IpTagUserDefineKnowledgeBaseHandler.java
@@ -1,6 +1,5 @@
package com.geedgenetworks.core.udf.knowlegdebase.handler;
-import com.geedgenetworks.core.pojo.KnowLedgeBaseFileMeta;
import com.geedgenetworks.core.utils.cn.csv.HighCsvReader;
import inet.ipaddr.IPAddress;
import inet.ipaddr.IPAddressString;
@@ -40,21 +39,19 @@ public class IpTagUserDefineKnowledgeBaseHandler extends AbstractMultipleKnowled
@Override
protected Boolean buildKnowledgeBase() {
TreeRangeMap<IPAddress, List<Node>> newIpTagRules = TreeRangeMap.create();
- this.knowledgeMetedataCacheMap.forEach((key, value) -> {
- buildSingleKnowledgeBase(newIpTagRules, convertId(key), value);
- });
+ getAllFileIds().forEach(id -> buildSingleKnowledgeBase(newIpTagRules, id));
ipTagRules = newIpTagRules;
return true;
}
- private void buildSingleKnowledgeBase(TreeRangeMap<IPAddress, List<Node>> treeRangeMap, Long id, KnowLedgeBaseFileMeta metadata) {
+ private void buildSingleKnowledgeBase(TreeRangeMap<IPAddress, List<Node>> treeRangeMap, Long id) {
try {
List<String> needColumns = new ArrayList<>();
needColumns.add("tag_value");
needColumns.add("addr_format");
needColumns.add("ip1");
needColumns.add("ip2");
- byte[] content = downloadFile(metadata.getPath(), 1);
+ byte[] content = downloadFile(id);
HighCsvReader highCsvReader = new HighCsvReader(new InputStreamReader(new ByteArrayInputStream(content)), needColumns);
HighCsvReader.CsvIterator iterator = highCsvReader.getIterator();
while (iterator.hasNext()) {
@@ -124,4 +121,10 @@ public class IpTagUserDefineKnowledgeBaseHandler extends AbstractMultipleKnowled
}
return nodes == null ? new ArrayList<>() : nodes;
}
+
+ @Override
+ public void close() {
+ ipTagRules.clear();
+ ipTagRules = null;
+ }
}
diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/IpVpnKnowledgeBaseHandler.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/IpVpnKnowledgeBaseHandler.java
index a1fc9ff..99b3168 100644
--- a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/IpVpnKnowledgeBaseHandler.java
+++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/IpVpnKnowledgeBaseHandler.java
@@ -44,7 +44,7 @@ public class IpVpnKnowledgeBaseHandler extends AbstractSingleKnowledgeBaseHandle
needColumns.add("ip1");
needColumns.add("ip2");
needColumns.add("vpn_service_name");
- byte[] content = downloadFile(knowledgeMetedataCache.getPath(), 1);
+ byte[] content = downloadFile();
HighCsvReader highCsvReader = new HighCsvReader(new InputStreamReader(new ByteArrayInputStream(content)), needColumns);
TreeRangeMap<IPAddress, String> newTreeRangeMap = TreeRangeMap.create();
HighCsvReader.CsvIterator iterator = highCsvReader.getIterator();
@@ -103,4 +103,10 @@ public class IpVpnKnowledgeBaseHandler extends AbstractSingleKnowledgeBaseHandle
}
return null;
}
+
+ @Override
+ public void close() {
+ treeRangeMap.clear();
+ treeRangeMap = null;
+ }
}
diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/LinkDirectionKnowledgeBaseHandler.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/LinkDirectionKnowledgeBaseHandler.java
index 861eb70..ba0925c 100644
--- a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/LinkDirectionKnowledgeBaseHandler.java
+++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/LinkDirectionKnowledgeBaseHandler.java
@@ -38,7 +38,7 @@ public class LinkDirectionKnowledgeBaseHandler extends AbstractSingleKnowledgeBa
List<String> needColumns = new ArrayList<>();
needColumns.add("peer_city");
needColumns.add("link_id");
- byte[] content = downloadFile(knowledgeMetedataCache.getPath(), 1);
+ byte[] content = downloadFile();
HighCsvReader highCsvReader = new HighCsvReader(new InputStreamReader(new ByteArrayInputStream(content)), needColumns);
HashMap<Long, String> newLinkMap = new HashMap<>((int) (highCsvReader.getLineNumber() / 0.75F + 1.0F));
HighCsvReader.CsvIterator iterator = highCsvReader.getIterator();
@@ -71,4 +71,10 @@ public class LinkDirectionKnowledgeBaseHandler extends AbstractSingleKnowledgeBa
}
return null;
}
+
+ @Override
+ public void close() {
+ linkMap.clear();
+ linkMap = null;
+ }
}
diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/RuleKnowledgeBaseHandler.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/RuleKnowledgeBaseHandler.java
index a97a2a1..a3335c4 100644
--- a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/RuleKnowledgeBaseHandler.java
+++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/RuleKnowledgeBaseHandler.java
@@ -2,6 +2,8 @@ package com.geedgenetworks.core.udf.knowlegdebase.handler;
import com.alibaba.fastjson2.JSON;
import com.geedgenetworks.common.config.KnowledgeBaseConfig;
+import com.geedgenetworks.common.exception.CommonErrorCode;
+import com.geedgenetworks.common.exception.GrootStreamRuntimeException;
import lombok.Data;
import org.apache.http.HttpEntity;
import org.apache.http.client.methods.CloseableHttpResponse;
@@ -53,17 +55,29 @@ public class RuleKnowledgeBaseHandler extends AbstractKnowledgeBaseHandler {
@Override
public boolean initKnowledgeBase(KnowledgeBaseConfig knowledgeBaseConfig) {
this.knowledgeBaseConfig = knowledgeBaseConfig;
- updateCache(knowledgeBaseConfig.getFsPath());
+ updateCache();
return true;
}
@Override
public void updateKnowledgeBase() {
- updateCache(knowledgeBaseConfig.getFsPath());
+ updateCache();
}
- public void updateCache(String address) {
- String url = getUrl(address);
+ public void updateCache() {
+ if ("http".equals(knowledgeBaseConfig.getFsType())) {
+ requestApi();
+ }
+ if ("local".equals(knowledgeBaseConfig.getFsType())) {
+ byte[] localRuleContent = getFileFromLocal(knowledgeBaseConfig.getFsPath());
+ RuleResponse ruleResponse = JSON.parseObject(new String(localRuleContent), RuleResponse.class);
+ processResponse(ruleResponse);
+ }
+ throw new GrootStreamRuntimeException(CommonErrorCode.ILLEGAL_ARGUMENT, knowledgeBaseConfig.getFsType() + " is illegal");
+ }
+
+ private void requestApi() {
+ String url = getUrl(knowledgeBaseConfig.getFsPath());
final HttpGet httpGet = new HttpGet(url);
httpGet.addHeader("Accept", "application/json");
httpGet.addHeader("Cn-Authorization", knowledgeBaseConfig.getProperties().get("token"));
@@ -125,6 +139,12 @@ public class RuleKnowledgeBaseHandler extends AbstractKnowledgeBaseHandler {
return kbIdMap.get(kbId);
}
+ public void close() {
+ ruleMap.clear();
+ nameMap.clear();
+ kbIdMap.clear();
+ }
+
@Data
private static final class RuleResponse {
private Integer code;
diff --git a/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/LookupTestUtils.java b/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/LookupTestUtils.java
index 56f1776..0178375 100644
--- a/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/LookupTestUtils.java
+++ b/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/LookupTestUtils.java
@@ -5,9 +5,9 @@ import com.geedgenetworks.common.Constants;
import com.geedgenetworks.common.config.CommonConfig;
import com.geedgenetworks.common.config.KnowledgeBaseConfig;
import com.geedgenetworks.core.pojo.KnowLedgeBaseFileMeta;
+import com.geedgenetworks.core.udf.knowlegdebase.KnowledgeBaseUpdateJob;
import com.geedgenetworks.core.udf.knowlegdebase.handler.AbstractKnowledgeBaseHandler;
import com.geedgenetworks.core.udf.knowlegdebase.handler.AbstractMultipleKnowledgeBaseHandler;
-import com.geedgenetworks.core.udf.knowlegdebase.KnowledgeBaseUpdateJob;
import org.apache.flink.api.common.ExecutionConfig;
import org.apache.flink.api.common.functions.RuntimeContext;
import org.apache.flink.configuration.Configuration;
@@ -19,6 +19,8 @@ import java.util.HashMap;
import java.util.List;
import java.util.Map;
+import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.Mockito.mockStatic;
/**
@@ -29,7 +31,7 @@ import static org.mockito.Mockito.mockStatic;
public class LookupTestUtils {
private static String fsPath = "testPath";
- private static String fsType = "testType";
+ private static String fsType = "http";
private static List<String> fsFiles = Arrays.asList("testFile");
public static String kbName = "testKbName";
private static String downloadPath = "testDownloadPath";
@@ -75,11 +77,15 @@ public class LookupTestUtils {
static void mockMultipleKnowledgeBaseHandler(String downloadContent) {
checkStaticMock();
KnowLedgeBaseFileMeta KnowLedgeBaseFileMeta = new KnowLedgeBaseFileMeta();
+ KnowLedgeBaseFileMeta.setKb_id("1");
KnowLedgeBaseFileMeta.setPath(downloadPath);
Map<String, KnowLedgeBaseFileMeta> KnowLedgeBaseFileMetaMap = new HashMap<>();
KnowLedgeBaseFileMetaMap.put("1", KnowLedgeBaseFileMeta);
abstractMultipleKnowledgeBaseHandlerMockedStatic.when(() -> AbstractMultipleKnowledgeBaseHandler.getMetadata(fsPath)).thenReturn(KnowLedgeBaseFileMetaMap);
abstractKnowledgeBaseHandlerMockedStatic.when(() -> AbstractKnowledgeBaseHandler.downloadFile(downloadPath, 1)).thenReturn(downloadContent.getBytes());
+ abstractMultipleKnowledgeBaseHandlerMockedStatic.when(() -> AbstractMultipleKnowledgeBaseHandler.decodeId(anyString())).thenCallRealMethod();
+ abstractMultipleKnowledgeBaseHandlerMockedStatic.when(() -> AbstractMultipleKnowledgeBaseHandler.encodeId(anyLong())).thenCallRealMethod();
+ abstractMultipleKnowledgeBaseHandlerMockedStatic.when(() -> AbstractMultipleKnowledgeBaseHandler.checkId(anyString())).thenCallRealMethod();
}
static RuntimeContext mockRuleRuntimeContext() {
@@ -95,7 +101,7 @@ public class LookupTestUtils {
knowledgeBaseConfig.setName(kbName);
KnowledgeBaseConfig ruleKnowledgeBaseConfig = new KnowledgeBaseConfig();
ruleKnowledgeBaseConfig.setFsPath("testRulePath");
- ruleKnowledgeBaseConfig.setFsType("testRuleType");
+ ruleKnowledgeBaseConfig.setFsType(fsType);
Map<String, String> properties = new HashMap<>();
properties.put("token", "testToken");
ruleKnowledgeBaseConfig.setProperties(properties);