From b121928ca74ccdd0f0c0790032346f750e238b8f Mon Sep 17 00:00:00 2001 From: wangkuan Date: Fri, 8 Mar 2024 18:26:25 +0800 Subject: [improve][core]asn与geoip知识库支持读取本地文件，添加同步变量用于函数关闭时判断是否需要清除内存 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- groot-common/pom.xml | 16 ++++- .../com/geedgenetworks/common/utils/FileUtils.java | 17 ++++++ .../projection/ProjectionProcessFunction.java | 7 ++- .../com/geedgenetworks/core/udf/AsnLookup.java | 9 ++- .../udf/knowlegdebase/KnowledgeBaseScheduler.java | 11 +++- .../udf/knowlegdebase/KnowledgeBaseUpdateJob.java | 4 +- .../handler/AbstractKnowledgeBaseHandler.java | 70 ++++++++++------------ .../handler/AsnKnowledgeBaseHandler.java | 52 ++++++++++++++-- .../handler/GeoIpKnowledgeBaseHandler.java | 37 +++++++++++- 9 files changed, 168 insertions(+), 55 deletions(-) diff --git a/groot-common/pom.xml b/groot-common/pom.xml index 8ff4910..4a728de 100644 --- a/groot-common/pom.xml +++ b/groot-common/pom.xml @@ -81,7 +81,21 @@ provided - + + org.apache.hadoop + hadoop-client + 2.7.1 + + + slf4j-log4j12 + org.slf4j + + + log4j-over-slf4j + org.slf4j + + + diff --git a/groot-common/src/main/java/com/geedgenetworks/common/utils/FileUtils.java b/groot-common/src/main/java/com/geedgenetworks/common/utils/FileUtils.java index 2d59477..6796514 100644 --- a/groot-common/src/main/java/com/geedgenetworks/common/utils/FileUtils.java +++ b/groot-common/src/main/java/com/geedgenetworks/common/utils/FileUtils.java @@ -190,4 +190,21 @@ public class FileUtils { public static String getFileName(@NonNull String filePath) { return filePath.substring(filePath.lastIndexOf(File.separatorChar) + 1); } + + public static byte[] getFileFromDir(String filePath) { + try (BufferedInputStream bis = new BufferedInputStream(Files.newInputStream(Paths.get(filePath))); + ByteArrayOutputStream baos = new ByteArrayOutputStream()) { + + byte[] buffer = new byte[1024]; + int bytesRead; + while ((bytesRead = bis.read(buffer)) != -1) { + baos.write(buffer, 0, bytesRead); + } + return baos.toByteArray(); + } catch (Exception e) { + String errorMsg = String.format("get local file [%s] failed", filePath); + throw new GrootStreamRuntimeException( + CommonErrorCode.UNSUPPORTED_OPERATION, errorMsg, e); + } + } } diff --git a/groot-core/src/main/java/com/geedgenetworks/core/processor/projection/ProjectionProcessFunction.java b/groot-core/src/main/java/com/geedgenetworks/core/processor/projection/ProjectionProcessFunction.java index 71454dc..cf82ab7 100644 --- a/groot-core/src/main/java/com/geedgenetworks/core/processor/projection/ProjectionProcessFunction.java +++ b/groot-core/src/main/java/com/geedgenetworks/core/processor/projection/ProjectionProcessFunction.java @@ -85,6 +85,7 @@ public class ProjectionProcessFunction extends ProcessFunction { } catch (Exception e) { throw new GrootStreamRuntimeException(CommonErrorCode.UNSUPPORTED_OPERATION, "Initialization UDF failed!", e); } + KnowledgeBaseScheduler.increment(); } @Override @@ -172,10 +173,12 @@ public class ProjectionProcessFunction extends ProcessFunction { public void close() throws Exception { for (UdfEntity udfEntity : functions) { - udfEntity.getUdfFunction().close(); } - KnowledgeBaseScheduler.stopSchedulerForKnowledgeBase(); + KnowledgeBaseScheduler.decrement(); + if(KnowledgeBaseScheduler.getCount()<=0) { + KnowledgeBaseScheduler.stopSchedulerForKnowledgeBase(); + } super.close(); } diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/AsnLookup.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/AsnLookup.java index 77e2c9a..422251f 100644 --- a/groot-core/src/main/java/com/geedgenetworks/core/udf/AsnLookup.java +++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/AsnLookup.java @@ -48,6 +48,8 @@ public class AsnLookup implements UDF { } this.lookupFieldName = udfContext.getLookup_fields().get(0); this.outputFieldName = udfContext.getOutput_fields().get(0); + AsnKnowledgeBaseHandler.increment(); + log.warn("AsnKnowledgeBaseHandlerCountIn"+AsnKnowledgeBaseHandler.getCount()); } @@ -71,8 +73,11 @@ public class AsnLookup implements UDF { @Override public void close() { - - KnowledgeBaseUpdateJob.removeKnowledgeBase(kbName); + AsnKnowledgeBaseHandler.decrement(); + log.warn("+++++++++++++++++++++++++++++++++++AsnKnowledgeBaseHandlerCountDe"+AsnKnowledgeBaseHandler.getCount()); + if (AsnKnowledgeBaseHandler.getCount()<=0) { + KnowledgeBaseUpdateJob.removeKnowledgeBase(kbName); + } } private void checkUdfContext(UDFContext udfContext) { diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/KnowledgeBaseScheduler.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/KnowledgeBaseScheduler.java index 2047bb7..f006044 100644 --- a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/KnowledgeBaseScheduler.java +++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/KnowledgeBaseScheduler.java @@ -11,8 +11,17 @@ public class KnowledgeBaseScheduler { private final static String JOB_GROUP = "knowledgeBase"; private final static String TRIGGER_NAME = "knowledgeBaseUpdateJob"; private static boolean methodCalled = false; + private static int count = 0; - + public static synchronized void increment() { + count++; + } + public static synchronized void decrement() { + count--; + } + public static synchronized int getCount() { + return count; + } public synchronized static void startSchedulerForKnowledgeBase(int intervalInMinutes) throws SchedulerException { if(!methodCalled){ diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/KnowledgeBaseUpdateJob.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/KnowledgeBaseUpdateJob.java index 84a7be4..c8c47de 100644 --- a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/KnowledgeBaseUpdateJob.java +++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/KnowledgeBaseUpdateJob.java @@ -63,7 +63,9 @@ public class KnowledgeBaseUpdateJob implements Job { if (!KnowLedgeBaseEntityMap.isEmpty()) { for (Map.Entry entry : KnowLedgeBaseEntityMap.entrySet()) { try { - updateKnowledgeBase(entry.getValue()); + if("http".equals(entry.getValue().getKnowledgeConfig().getFsType())) { + updateKnowledgeBase(entry.getValue()); + } } catch (Exception e) { log.error(e.getMessage()); } 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 96eeb11..645cdb1 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 @@ -3,6 +3,8 @@ package com.geedgenetworks.core.udf.knowlegdebase.handler; 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; @@ -13,6 +15,9 @@ import java.security.NoSuchAlgorithmException; 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 { @@ -58,30 +63,28 @@ public abstract class AbstractKnowledgeBaseHandler { return inputStream; } + + public static byte[] getFileFromLocal(String filepath) { + return FileUtils.getFileFromDir(filepath); + } + + + public static List getMetadata(String type, String path) { List knowLedgeBaseFileMetaList = new ArrayList<>(); try { - switch (type) { - case "http": - - HttpClientPoolUtil httpClientPoolUtil = HttpClientPoolUtil.getInstance(); - String metadate = httpClientPoolUtil.httpGet(URI.create(path)); - Map metaObject = JSON.parseObject(metadate, Map.class); - JSONArray jsonArray = (JSONArray) metaObject.get("data"); - List resultList = jsonArray.toJavaList(KnowLedgeBaseFileMeta.class); - knowLedgeBaseFileMetaList.addAll(resultList); - knowLedgeBaseFileMetaList.removeIf(knowLedgeBaseFileMeta -> !"latest".equals(knowLedgeBaseFileMeta.getVersion())); - break; - case "local": - break; - case "hdfs": - break; - default: - log.error("get file Metadata error: fs_type not suppoort !"); - break; - } + + HttpClientPoolUtil httpClientPoolUtil = HttpClientPoolUtil.getInstance(); + String metadate = httpClientPoolUtil.httpGet(URI.create(path)); + Map metaObject = JSON.parseObject(metadate, Map.class); + JSONArray jsonArray = (JSONArray) metaObject.get("data"); + List resultList = jsonArray.toJavaList(KnowLedgeBaseFileMeta.class); + knowLedgeBaseFileMetaList.addAll(resultList); + knowLedgeBaseFileMetaList.removeIf(knowLedgeBaseFileMeta -> !"latest".equals(knowLedgeBaseFileMeta.getVersion())); + + } catch (Exception e) { log.error("get file Metadata error! path = " + path + " " + e.getMessage()); } @@ -95,26 +98,15 @@ public abstract class AbstractKnowledgeBaseHandler { List knowLedgeBaseFileMetaList = new ArrayList<>(); try { - switch (type) { - case "http": - - HttpClientPoolUtil httpClientPoolUtil = HttpClientPoolUtil.getInstance(); - String metadate = httpClientPoolUtil.httpGet(URI.create(path + "?kb_id=" + kbId)); - Map metaObject = JSON.parseObject(metadate, Map.class); - JSONArray jsonArray = (JSONArray) metaObject.get("data"); - knowLedgeBaseFileMetaList = jsonArray.toJavaList(KnowLedgeBaseFileMeta.class); - knowLedgeBaseFileMetaList.removeIf(knowLedgeBaseFileMeta -> !"latest".equals(knowLedgeBaseFileMeta.getVersion())); - break; - case "local": - log.error("get file Metadata error: fs_type local not suppoort !"); - break; - case "hdfs": - log.error("get file Metadata error: fs_type hdfs not suppoort !"); - break; - default: - log.error("get file Metadata error: fs_type not suppoort !"); - break; - } + + HttpClientPoolUtil httpClientPoolUtil = HttpClientPoolUtil.getInstance(); + String metadate = httpClientPoolUtil.httpGet(URI.create(path + "?kb_id=" + kbId)); + Map metaObject = JSON.parseObject(metadate, Map.class); + JSONArray jsonArray = (JSONArray) metaObject.get("data"); + knowLedgeBaseFileMetaList = jsonArray.toJavaList(KnowLedgeBaseFileMeta.class); + knowLedgeBaseFileMetaList.removeIf(knowLedgeBaseFileMeta -> !"latest".equals(knowLedgeBaseFileMeta.getVersion())); + + } catch (Exception e) { log.error("get file Metadata error! path = " + path + " " + e.getMessage()); return null; diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/AsnKnowledgeBaseHandler.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/AsnKnowledgeBaseHandler.java index 6d5860c..381b0aa 100644 --- a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/AsnKnowledgeBaseHandler.java +++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/AsnKnowledgeBaseHandler.java @@ -2,29 +2,38 @@ package com.geedgenetworks.core.udf.knowlegdebase.handler; import com.geedgenetworks.common.config.KnowledgeBaseConfig; - import com.geedgenetworks.core.pojo.KnowLedgeBaseFileMeta; import com.geedgenetworks.core.pojo.KnowledgeBaseInfoForIpLookUp; import com.geedgenetworks.utils.IpLookupV2; import lombok.Getter; import lombok.extern.slf4j.Slf4j; - import java.io.ByteArrayInputStream; import java.util.*; import java.util.concurrent.ConcurrentHashMap; + @Slf4j public class AsnKnowledgeBaseHandler extends AbstractKnowledgeBaseHandler { @Getter private static Map kbNameWithKbConfigAndFileMetas = new ConcurrentHashMap<>(); private static AsnKnowledgeBaseHandler instance; + private static int count = 0; // 私有构造函数，防止外部实例化 private AsnKnowledgeBaseHandler() { } + public static synchronized void increment() { + count++; + } + public static synchronized void decrement() { + count--; + } + public static synchronized int getCount() { + return count; + } public static synchronized AsnKnowledgeBaseHandler getInstance() { if (instance == null) { instance = new AsnKnowledgeBaseHandler(); @@ -34,8 +43,12 @@ public class AsnKnowledgeBaseHandler extends AbstractKnowledgeBaseHandler { @Override public boolean initKnowledgeBase(KnowledgeBaseConfig knowledgeBaseConfig) { - return buildKnowledgeBase(knowledgeBaseConfig); - } + + if("http".equals(knowledgeBaseConfig.getFsType())) + return buildKnowledgeBaseWithMeta(knowledgeBaseConfig); + else + return buildKnowledgeBaseWithoutMeta(knowledgeBaseConfig); + } @Override public void updateKnowledgeBase() { @@ -43,7 +56,7 @@ public class AsnKnowledgeBaseHandler extends AbstractKnowledgeBaseHandler { for (Map.Entry entry : kbNameWithKbConfigAndFileMetas.entrySet()) { try { if (ifNeedUpdate(entry.getValue())) { - boolean result = buildKnowledgeBase(entry.getValue().getKnowledgeConfig()); + boolean result = buildKnowledgeBaseWithMeta(entry.getValue().getKnowledgeConfig()); if(result) { log.warn("AsnKnowledgeBase " + entry.getKey() + " update success!"); } @@ -80,8 +93,35 @@ public class AsnKnowledgeBaseHandler extends AbstractKnowledgeBaseHandler { } return false; } + private boolean buildKnowledgeBaseWithoutMeta(KnowledgeBaseConfig knowledgeBaseConfig) { - private static boolean buildKnowledgeBase(KnowledgeBaseConfig knowledgeBaseConfig) { + IpLookupV2.Builder asnLookupBuilder = new IpLookupV2.Builder(false); + for (int i = 0; i < knowledgeBaseConfig.getFiles().size(); i++) { + try { + byte[] fileByte = getFileFromLocal(knowledgeBaseConfig.getFsPath() + knowledgeBaseConfig.getFiles().get(i)); + switch (i) { + case 0: + asnLookupBuilder.loadAsnDataFile(new ByteArrayInputStream(fileByte)); + break; + case 1: + asnLookupBuilder.loadDataFilePrivate(new ByteArrayInputStream(fileByte)); + break; + } + + } catch (Exception e) { + log.error("buildKnowledgeBase error " + e.getMessage()); + return false; + } + + } + IpLookupV2 ipLookup = asnLookupBuilder.build(); + KnowledgeBaseInfoForIpLookUp kbInfoForIpLookUp = new KnowledgeBaseInfoForIpLookUp(); + kbInfoForIpLookUp.setKnowledgeConfig(knowledgeBaseConfig); + kbInfoForIpLookUp.setIpLookupV2(ipLookup); + kbNameWithKbConfigAndFileMetas.put(knowledgeBaseConfig.getName(), kbInfoForIpLookUp); + return true; + } + private static boolean buildKnowledgeBaseWithMeta(KnowledgeBaseConfig knowledgeBaseConfig) { IpLookupV2.Builder asnLookupBuilder = new IpLookupV2.Builder(false); List knowLedgeBaseFileMetas = new ArrayList<>(); diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/GeoIpKnowledgeBaseHandler.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/GeoIpKnowledgeBaseHandler.java index a441369..5272ada 100644 --- a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/GeoIpKnowledgeBaseHandler.java +++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/GeoIpKnowledgeBaseHandler.java @@ -33,7 +33,11 @@ public class GeoIpKnowledgeBaseHandler extends AbstractKnowledgeBaseHandler { @Override public boolean initKnowledgeBase(KnowledgeBaseConfig knowledgeBaseConfig) { - return buildKnowledgeBase(knowledgeBaseConfig); + if("http".equals(knowledgeBaseConfig.getFsType())) + return buildKnowledgeBaseWithMeta(knowledgeBaseConfig); + else + return buildKnowledgeBaseWithoutMeta(knowledgeBaseConfig); + } @Override @@ -42,7 +46,7 @@ public class GeoIpKnowledgeBaseHandler extends AbstractKnowledgeBaseHandler { for (Map.Entry entry : kbNameWithKbConfigAndFileMetas.entrySet()) { try { if (ifNeedUpdate(entry.getValue())) { - boolean result = buildKnowledgeBase(entry.getValue().getKnowledgeConfig()); + boolean result = buildKnowledgeBaseWithMeta(entry.getValue().getKnowledgeConfig()); if(result) { log.warn("GeoIpKnowledgeBase " + entry.getKey() + " update success!"); } @@ -66,8 +70,35 @@ public class GeoIpKnowledgeBaseHandler extends AbstractKnowledgeBaseHandler { } return false; } + private boolean buildKnowledgeBaseWithoutMeta(KnowledgeBaseConfig knowledgeBaseConfig) { + + IpLookupV2.Builder ipLookupBuilder = new IpLookupV2.Builder(false); + for (int i = 0; i < knowledgeBaseConfig.getFiles().size(); i++) { + try { + byte[] fileByte = getFileFromLocal(knowledgeBaseConfig.getFsPath()+knowledgeBaseConfig.getFiles().get(i)); + switch (i) { + case 0: + ipLookupBuilder.loadDataFile(new ByteArrayInputStream(fileByte)); + break; + case 1: + ipLookupBuilder.loadDataFilePrivate(new ByteArrayInputStream(fileByte)); + break; + } + + } catch (Exception e) { + log.error("buildKnowledgeBase error " + e.getMessage()); + return false; + } - private static boolean buildKnowledgeBase(KnowledgeBaseConfig knowledgeBaseConfig) { + } + IpLookupV2 ipLookup = ipLookupBuilder.build(); + KnowledgeBaseInfoForIpLookUp kbInfoForIpLookUp = new KnowledgeBaseInfoForIpLookUp(); + kbInfoForIpLookUp.setKnowledgeConfig(knowledgeBaseConfig); + kbInfoForIpLookUp.setIpLookupV2(ipLookup); + kbNameWithKbConfigAndFileMetas.put(knowledgeBaseConfig.getName(), kbInfoForIpLookUp); + return true; + } + private static boolean buildKnowledgeBaseWithMeta(KnowledgeBaseConfig knowledgeBaseConfig) { IpLookupV2.Builder ipLookupBuilder = new IpLookupV2.Builder(false); List knowLedgeBaseFileMetas = new ArrayList<>(); -- cgit v1.2.3 From 7115979e16f6e9577df9063f25d678e9cfbc5aac Mon Sep 17 00:00:00 2001 From: wangkuan Date: Fri, 8 Mar 2024 18:40:21 +0800 Subject: [improve][core]删除hadoop相关依赖 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- groot-common/pom.xml | 16 +--------------- 1 file changed, 1 insertion(+), 15 deletions(-) diff --git a/groot-common/pom.xml b/groot-common/pom.xml index 4a728de..8ff4910 100644 --- a/groot-common/pom.xml +++ b/groot-common/pom.xml @@ -81,21 +81,7 @@ provided - - org.apache.hadoop - hadoop-client - 2.7.1 - - - slf4j-log4j12 - org.slf4j - - - log4j-over-slf4j - org.slf4j - - - + -- cgit v1.2.3 From e25026d4784e04873aebad09f6306bc0b3945dfa Mon Sep 17 00:00:00 2001 From: gujinkai Date: Mon, 11 Mar 2024 14:10:21 +0800 Subject: [Improve][core] CN knowledge Handler add option: download from local --- .../handler/AbstractKnowledgeBaseHandler.java | 3 -- .../AbstractMultipleKnowledgeBaseHandler.java | 47 ++++++++++++++++++++-- .../AbstractSingleKnowledgeBaseHandler.java | 17 +++++++- .../handler/AppCategoryKnowledgeBaseHandler.java | 7 +++- .../AppTagUserDefineKnowledgeBaseHandler.java | 13 +++--- .../handler/DnsServerInfoKnowledgeBaseHandler.java | 7 +++- .../DomainTagUserDefineKnowledgeBaseHandler.java | 16 +++++--- .../handler/DomainVpnKnowledgeBaseHandler.java | 7 +++- .../handler/FqdnCategoryKnowledgeBaseHandler.java | 7 +++- .../handler/FqdnIcpKnowledgeBaseHandler.java | 7 +++- .../handler/FqdnWhoisKnowledgeBaseHandler.java | 7 +++- .../handler/IdcRenterKnowledgeBaseHandler.java | 7 +++- .../handler/InternalIpKnowledgeBaseHandler.java | 7 +++- .../handler/IocDarkwebKnowledgeBaseHandler.java | 9 ++++- .../handler/IocMalwareKnowledgeBaseHandler.java | 10 ++++- .../IpTagUserDefineKnowledgeBaseHandler.java | 15 ++++--- .../handler/IpVpnKnowledgeBaseHandler.java | 8 +++- .../handler/LinkDirectionKnowledgeBaseHandler.java | 8 +++- .../handler/RuleKnowledgeBaseHandler.java | 28 +++++++++++-- .../core/udf/cn/LookupTestUtils.java | 10 ++++- 20 files changed, 196 insertions(+), 44 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 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 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 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 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> 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> appTags, Long id, KnowLedgeBaseFileMeta metadata) { + private void buildSingleKnowledgeBase(Map> appTags, Long id) { try { List 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 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> 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 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 newFqdnTagFuzzyMatchRules = new Trie<>(); Map> 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 fqdnTagsFuzzy, Map> fqdnTagsFull, Long id, KnowLedgeBaseFileMeta metadata) { + private void buildSingleKnowledgeBase(Trie fqdnTagsFuzzy, Map> fqdnTagsFull, Long id) { try { List 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 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 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 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 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 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 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 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 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 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 newIpMap = new HashMap<>((int) (highCsvReader.getLineNumber() / 0.75F + 1.0F)); HashMap 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 newIpMap = new HashMap<>((int) (highCsvReader.getLineNumber() / 0.75F + 1.0F)); HashMap 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> 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> treeRangeMap, Long id, KnowLedgeBaseFileMeta metadata) { + private void buildSingleKnowledgeBase(TreeRangeMap> treeRangeMap, Long id) { try { List 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 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 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 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..c50b1b0 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 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 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() { -- cgit v1.2.3 From 1f4ea178b36c81d82af82783fe664e4eac074585 Mon Sep 17 00:00:00 2001 From: gujinkai Date: Mon, 11 Mar 2024 14:10:21 +0800 Subject: [Improve][core] CN knowledge Handler add option: download from local --- .../handler/AbstractKnowledgeBaseHandler.java | 3 -- .../AbstractMultipleKnowledgeBaseHandler.java | 47 ++++++++++++++++++++-- .../AbstractSingleKnowledgeBaseHandler.java | 17 +++++++- .../handler/AppCategoryKnowledgeBaseHandler.java | 7 +++- .../AppTagUserDefineKnowledgeBaseHandler.java | 13 +++--- .../handler/DnsServerInfoKnowledgeBaseHandler.java | 7 +++- .../DomainTagUserDefineKnowledgeBaseHandler.java | 16 +++++--- .../handler/DomainVpnKnowledgeBaseHandler.java | 7 +++- .../handler/FqdnCategoryKnowledgeBaseHandler.java | 7 +++- .../handler/FqdnIcpKnowledgeBaseHandler.java | 7 +++- .../handler/FqdnWhoisKnowledgeBaseHandler.java | 7 +++- .../handler/IdcRenterKnowledgeBaseHandler.java | 7 +++- .../handler/InternalIpKnowledgeBaseHandler.java | 7 +++- .../handler/IocDarkwebKnowledgeBaseHandler.java | 9 ++++- .../handler/IocMalwareKnowledgeBaseHandler.java | 10 ++++- .../IpTagUserDefineKnowledgeBaseHandler.java | 15 ++++--- .../handler/IpVpnKnowledgeBaseHandler.java | 8 +++- .../handler/LinkDirectionKnowledgeBaseHandler.java | 8 +++- .../handler/RuleKnowledgeBaseHandler.java | 28 +++++++++++-- .../core/udf/cn/LookupTestUtils.java | 12 ++++-- 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 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 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 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 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> 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> appTags, Long id, KnowLedgeBaseFileMeta metadata) { + private void buildSingleKnowledgeBase(Map> appTags, Long id) { try { List 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 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> 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 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 newFqdnTagFuzzyMatchRules = new Trie<>(); Map> 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 fqdnTagsFuzzy, Map> fqdnTagsFull, Long id, KnowLedgeBaseFileMeta metadata) { + private void buildSingleKnowledgeBase(Trie fqdnTagsFuzzy, Map> fqdnTagsFull, Long id) { try { List 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 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 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 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 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 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 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 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 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 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 newIpMap = new HashMap<>((int) (highCsvReader.getLineNumber() / 0.75F + 1.0F)); HashMap 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 newIpMap = new HashMap<>((int) (highCsvReader.getLineNumber() / 0.75F + 1.0F)); HashMap 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> 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> treeRangeMap, Long id, KnowLedgeBaseFileMeta metadata) { + private void buildSingleKnowledgeBase(TreeRangeMap> treeRangeMap, Long id) { try { List 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 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 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 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 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 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 properties = new HashMap<>(); properties.put("token", "testToken"); ruleKnowledgeBaseConfig.setProperties(properties); -- cgit v1.2.3 From b9e20efd2384b8bc94f974a0e8d47b8b00736410 Mon Sep 17 00:00:00 2001 From: wangkuan Date: Thu, 14 Mar 2024 16:43:00 +0800 Subject: [improve][core]增加本地知识库文件单元测试 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- config/grootstream.yaml | 2 +- groot-bootstrap/pom.xml | 17 ----------- groot-common/src/main/resources/grootstream.yaml | 2 +- groot-common/src/test/resources/grootstream.yaml | 2 +- .../com/geedgenetworks/core/udf/AsnLookup.java | 4 +-- .../com/geedgenetworks/core/udf/GeoIpLookup.java | 8 +++++- .../handler/GeoIpKnowledgeBaseHandler.java | 11 +++++++- .../core/udf/test/AsnLookupFunctionTest.java | 10 +++---- .../core/udf/test/GeoIpLookupFunctionTest.java | 31 +++++---------------- groot-core/src/test/resources/asn.mmdb | Bin 0 -> 2316 bytes groot-core/src/test/resources/ip_builtin.mmdb | Bin 0 -> 4738 bytes 11 files changed, 34 insertions(+), 53 deletions(-) create mode 100644 groot-core/src/test/resources/asn.mmdb create mode 100644 groot-core/src/test/resources/ip_builtin.mmdb diff --git a/config/grootstream.yaml b/config/grootstream.yaml index fdd240a..acb19b1 100644 --- a/config/grootstream.yaml +++ b/config/grootstream.yaml @@ -12,7 +12,7 @@ grootstream: - 64af7077-eb9b-4b8f-80cf-2ceebc89bea9 - 004390bc-3135-4a6f-a492-3662ecb9e289 properties: - hos.path: http://192.168.44.12:8089 + hos.path: http://192.168.44.12:9098/hos hos.bucket.name.traffic_file: traffic_file_bucket hos.bucket.name.troubleshooting_file: troubleshooting_file_bucket scheduler.knowledge_base.update.interval.minutes: 5 diff --git a/groot-bootstrap/pom.xml b/groot-bootstrap/pom.xml index 62c2dc8..0a1fd9c 100644 --- a/groot-bootstrap/pom.xml +++ b/groot-bootstrap/pom.xml @@ -121,23 +121,6 @@ ${scope} - - org.apache.hbase - hbase-client - ${scope} - - - org.slf4j - slf4j-log4j12 - - - org.slf4j - log4j-over-slf4j - - - - - org.apache.flink diff --git a/groot-common/src/main/resources/grootstream.yaml b/groot-common/src/main/resources/grootstream.yaml index efc5b52..1a9a974 100644 --- a/groot-common/src/main/resources/grootstream.yaml +++ b/groot-common/src/main/resources/grootstream.yaml @@ -12,6 +12,6 @@ grootstream: - 64af7077-eb9b-4b8f-80cf-2ceebc89bea9 - 004390bc-3135-4a6f-a492-3662ecb9e289 properties: - hos.path: http://192.168.44.12:8089 + hos.path: http://192.168.44.12:9098/hos hos.bucket.name.traffic_file: traffic_file_bucket hos.bucket.name.troubleshooting_file: troubleshooting_file_bucket \ No newline at end of file diff --git a/groot-common/src/test/resources/grootstream.yaml b/groot-common/src/test/resources/grootstream.yaml index efc5b52..1a9a974 100644 --- a/groot-common/src/test/resources/grootstream.yaml +++ b/groot-common/src/test/resources/grootstream.yaml @@ -12,6 +12,6 @@ grootstream: - 64af7077-eb9b-4b8f-80cf-2ceebc89bea9 - 004390bc-3135-4a6f-a492-3662ecb9e289 properties: - hos.path: http://192.168.44.12:8089 + hos.path: http://192.168.44.12:9098/hos hos.bucket.name.traffic_file: traffic_file_bucket hos.bucket.name.troubleshooting_file: troubleshooting_file_bucket \ No newline at end of file diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/AsnLookup.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/AsnLookup.java index 422251f..33c0653 100644 --- a/groot-core/src/main/java/com/geedgenetworks/core/udf/AsnLookup.java +++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/AsnLookup.java @@ -49,7 +49,7 @@ public class AsnLookup implements UDF { this.lookupFieldName = udfContext.getLookup_fields().get(0); this.outputFieldName = udfContext.getOutput_fields().get(0); AsnKnowledgeBaseHandler.increment(); - log.warn("AsnKnowledgeBaseHandlerCountIn"+AsnKnowledgeBaseHandler.getCount()); + log.warn("AsnKnowledgeBaseHandlerCount "+AsnKnowledgeBaseHandler.getCount()); } @@ -74,9 +74,9 @@ public class AsnLookup implements UDF { @Override public void close() { AsnKnowledgeBaseHandler.decrement(); - log.warn("+++++++++++++++++++++++++++++++++++AsnKnowledgeBaseHandlerCountDe"+AsnKnowledgeBaseHandler.getCount()); if (AsnKnowledgeBaseHandler.getCount()<=0) { KnowledgeBaseUpdateJob.removeKnowledgeBase(kbName); + log.warn("AsnKnowledgeBaseHandler removeKnowledgeBase "+kbName); } } diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/GeoIpLookup.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/GeoIpLookup.java index 9763758..442d499 100644 --- a/groot-core/src/main/java/com/geedgenetworks/core/udf/GeoIpLookup.java +++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/GeoIpLookup.java @@ -56,6 +56,8 @@ public class GeoIpLookup implements UDF { if(udfContext.getOutput_fields()!=null && !udfContext.getOutput_fields().isEmpty()){ this.outputFieldName = udfContext.getOutput_fields().get(0); } + GeoIpKnowledgeBaseHandler.increment(); + log.warn("GeoIpKnowledgeBaseHandler "+GeoIpKnowledgeBaseHandler.getCount()); } @@ -163,6 +165,10 @@ public class GeoIpLookup implements UDF { @Override public void close() { - KnowledgeBaseUpdateJob.removeKnowledgeBase(kbName); + GeoIpKnowledgeBaseHandler.decrement(); + if (GeoIpKnowledgeBaseHandler.getCount()<=0) { + KnowledgeBaseUpdateJob.removeKnowledgeBase(kbName); + log.warn("GeoIpKnowledgeBaseHandler removeKnowledgeBase "+kbName); + } } } diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/GeoIpKnowledgeBaseHandler.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/GeoIpKnowledgeBaseHandler.java index 5272ada..4d207a2 100644 --- a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/GeoIpKnowledgeBaseHandler.java +++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/GeoIpKnowledgeBaseHandler.java @@ -20,10 +20,19 @@ public class GeoIpKnowledgeBaseHandler extends AbstractKnowledgeBaseHandler { private static Map kbNameWithKbConfigAndFileMetas = new ConcurrentHashMap<>(); private static GeoIpKnowledgeBaseHandler instance; + private static int count = 0; private GeoIpKnowledgeBaseHandler() { } - + public static synchronized void increment() { + count++; + } + public static synchronized void decrement() { + count--; + } + public static synchronized int getCount() { + return count; + } public static synchronized GeoIpKnowledgeBaseHandler getInstance() { if (instance == null) { instance = new GeoIpKnowledgeBaseHandler(); diff --git a/groot-core/src/test/java/com/geedgenetworks/core/udf/test/AsnLookupFunctionTest.java b/groot-core/src/test/java/com/geedgenetworks/core/udf/test/AsnLookupFunctionTest.java index 69a074b..5bcf89e 100644 --- a/groot-core/src/test/java/com/geedgenetworks/core/udf/test/AsnLookupFunctionTest.java +++ b/groot-core/src/test/java/com/geedgenetworks/core/udf/test/AsnLookupFunctionTest.java @@ -70,14 +70,14 @@ public class AsnLookupFunctionTest { knowledgeBaseConfig.setName("tsg_ip_asn"); Map parameters = new HashMap<>(); knowledgeBaseConfig.setProperties(parameters); - knowledgeBaseConfig.setFsType("http"); - knowledgeBaseConfig.setFsPath("http://192.168.44.67:9999/v1/knowledge_base"); - knowledgeBaseConfig.setFiles(Collections.singletonList("f9f6bc91-2142-4673-8249-e097c00fe1ea")); + knowledgeBaseConfig.setFsType("local"); + knowledgeBaseConfig.setFsPath(""); + knowledgeBaseConfig.setFiles(Collections.singletonList("src/test/resources/asn.mmdb")); knowledgeBaseConfigList.add(knowledgeBaseConfig); commonConfig.setKnowledgeBaseConfig(knowledgeBaseConfigList); KnowledgeBaseUpdateJob.registerKnowledgeBase(AsnKnowledgeBaseHandler.getInstance(), knowledgeBaseConfig); - String asn = AsnKnowledgeBaseHandler.lookUp("tsg_ip_asn","IP_TO_ASN","2600:1015:b002::"); - assertEquals("6167", asn); + String asn = AsnKnowledgeBaseHandler.lookUp("tsg_ip_asn","IP_TO_ASN","12.16.98.1"); + assertEquals("31978", asn); } diff --git a/groot-core/src/test/java/com/geedgenetworks/core/udf/test/GeoIpLookupFunctionTest.java b/groot-core/src/test/java/com/geedgenetworks/core/udf/test/GeoIpLookupFunctionTest.java index faa65f4..f487500 100644 --- a/groot-core/src/test/java/com/geedgenetworks/core/udf/test/GeoIpLookupFunctionTest.java +++ b/groot-core/src/test/java/com/geedgenetworks/core/udf/test/GeoIpLookupFunctionTest.java @@ -29,7 +29,6 @@ public class GeoIpLookupFunctionTest { udfContext.setOutput_fields(Collections.singletonList("iplocation")); } -//com.geedgenetworks.core.udf.GeoIpLookup单元测试 @Test public void testInit(){ GeoIpLookup geoIpLookup = new GeoIpLookup(); @@ -69,33 +68,17 @@ public class GeoIpLookupFunctionTest { knowledgeBaseConfig.setName("tsg_ip_location"); Map parameters = new HashMap<>(); knowledgeBaseConfig.setProperties(parameters); - knowledgeBaseConfig.setFsType("http"); - knowledgeBaseConfig.setFsPath("http://192.168.44.67:9999/v1/knowledge_base"); - knowledgeBaseConfig.setFiles(Arrays.asList("64af7077-eb9b-4b8f-80cf-2ceebc89bea9","004390bc-3135-4a6f-a492-3662ecb9e289")); + knowledgeBaseConfig.setFsType("local"); + knowledgeBaseConfig.setFsPath(""); + knowledgeBaseConfig.setFiles(Collections.singletonList("src/test/resources/ip_builtin.mmdb")); knowledgeBaseConfigList.add(knowledgeBaseConfig); commonConfig.setKnowledgeBaseConfig(knowledgeBaseConfigList); KnowledgeBaseUpdateJob.registerKnowledgeBase(GeoIpKnowledgeBaseHandler.getInstance(), knowledgeBaseConfig); - String countryLookup = GeoIpKnowledgeBaseHandler.lookUp("tsg_ip_location","IP_TO_COUNTRY","2600:1015:b002::"); - assertEquals("美国", countryLookup); - String provinceLookup = GeoIpKnowledgeBaseHandler.lookUp("tsg_ip_location","IP_TO_PROVINCE","2600:1015:b002::"); - assertEquals("印第安纳州", provinceLookup); - String cityLookup = GeoIpKnowledgeBaseHandler.lookUp("tsg_ip_location","IP_TO_CITY","2600:1015:b002::"); - assertEquals("Unknown", cityLookup); - String cityLookupDetail = GeoIpKnowledgeBaseHandler.lookUp("tsg_ip_location","IP_TO_SUBDIVISION_ADDR","2600:1015:b002::"); - assertEquals("美国.印第安纳州.Unknown", cityLookupDetail); - String locationLookupDetail = GeoIpKnowledgeBaseHandler.lookUp("tsg_ip_location","IP_TO_DETAIL","2600:1015:b002::"); - assertEquals("美国.印第安纳州.Unknown.", locationLookupDetail.trim()); - String latLngLookup = GeoIpKnowledgeBaseHandler.lookUp("tsg_ip_location","IP_TO_LATLNG","2600:1015:b002::"); - assertEquals("39.9388,-86.1137", latLngLookup); - String ispLookup = GeoIpKnowledgeBaseHandler.lookUp("tsg_ip_location","IP_TO_PROVIDER","2600:1015:b002::"); - assertEquals("", ispLookup); - String infoLookupToJSONString = GeoIpKnowledgeBaseHandler.lookUp("tsg_ip_location","IP_TO_JSON","2600:1015:b002::"); - assertEquals(JSONObject.parseObject(infoLookupToJSONString).getString("country"), "美国"); - String country = GeoIpKnowledgeBaseHandler.lookUpObject("tsg_ip_location","2600:1015:b002::").getCountry(); - assertEquals("美国", country); + String countryLookup = GeoIpKnowledgeBaseHandler.lookUp("tsg_ip_location","IP_TO_COUNTRY","86.111.175.128"); + assertEquals("United Kingdom", countryLookup); + String country = GeoIpKnowledgeBaseHandler.lookUpObject("tsg_ip_location","86.111.175.128").getCountry(); + assertEquals("United Kingdom", country); assertNull(GeoIpKnowledgeBaseHandler.lookUpObject("tsg_ip_location", "255.255.255.255")); } - - } diff --git a/groot-core/src/test/resources/asn.mmdb b/groot-core/src/test/resources/asn.mmdb new file mode 100644 index 0000000..a20e238 Binary files /dev/null and b/groot-core/src/test/resources/asn.mmdb differ diff --git a/groot-core/src/test/resources/ip_builtin.mmdb b/groot-core/src/test/resources/ip_builtin.mmdb new file mode 100644 index 0000000..b333c64 Binary files /dev/null and b/groot-core/src/test/resources/ip_builtin.mmdb differ -- cgit v1.2.3