[improve][core] Indicator lookup add subscriber's tags

4 files changed, 47 insertions, 1 deletions

diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/IntelligenceIndicatorLookup.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/IntelligenceIndicatorLookup.java
index 9f14bd2..857ae74 100644
--- a/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/IntelligenceIndicatorLookup.java
+++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/cn/IntelligenceIndicatorLookup.java
@@ -57,6 +57,17 @@ public class IntelligenceIndicatorLookup extends AbstractKnowledgeScalarFunction
                         }
                     }
                     break;
+                case "SUBSCRIBER_TO_TAG":
+                    List<String> subscriberTags = knowledgeBaseHandler.lookupBySubscriber(lookupValue);
+                    if (subscriberTags != null && subscriberTags.size() > 0) {
+                        hitCounter.inc();
+                        if (event.getExtractedFields().get(outputFieldName) != null && event.getExtractedFields().get(outputFieldName) instanceof List) {
+                            ((List<String>) event.getExtractedFields().get(outputFieldName)).addAll(subscriberTags);
+                        } else {
+                            event.getExtractedFields().put(outputFieldName, subscriberTags);
+                        }
+                    }
+                    break;
                 default:
                     logger.error("unknown option :" + option);
                     break;
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 f66cbe0..759ab38 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
@@ -106,6 +106,10 @@ public abstract class AbstractSingleKnowledgeBaseHandler extends AbstractKnowled
             knowledgeMetedataCache = knowledgeMetedata;
             return true;
         }
+        if (knowledgeMetedata.getIsValid() != knowledgeMetedataCache.getIsValid()) {
+            knowledgeMetedataCache = knowledgeMetedata;
+            return true;
+        }
         if (knowledgeMetedataCache.getSha256().equals(knowledgeMetedata.getSha256())) {
             return false;
         } else {
diff --git a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/IntelligenceIndicatorKnowledgeBaseHandler.java b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/IntelligenceIndicatorKnowledgeBaseHandler.java
index 94fdae1..093dd1d 100644
--- a/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/IntelligenceIndicatorKnowledgeBaseHandler.java
+++ b/groot-core/src/main/java/com/geedgenetworks/core/udf/knowlegdebase/handler/IntelligenceIndicatorKnowledgeBaseHandler.java
@@ -31,6 +31,8 @@ public class IntelligenceIndicatorKnowledgeBaseHandler extends AbstractSingleKno
     // *开头，模糊匹配
     private Trie<String> domainSuffix = new Trie<>();
 
+    private HashMap<String, List<String>> subscriberTagMap = new HashMap<>();
+
     private IntelligenceIndicatorKnowledgeBaseHandler() {
     }
 
@@ -51,12 +53,14 @@ public class IntelligenceIndicatorKnowledgeBaseHandler extends AbstractSingleKno
             needColumns.add("ip1");
             needColumns.add("ip2");
             needColumns.add("domain");
+            needColumns.add("subscriber");
             needColumns.add("tags");
             byte[] content = downloadFile();
             HighCsvReader highCsvReader = new HighCsvReader(new InputStreamReader(new ByteArrayInputStream(content)), needColumns);
             TreeRangeMap<IPAddress, List<String>> newIpTagMap = TreeRangeMap.create();
             HashMap<String, List<String>> newDomainMap = new HashMap<>((int) (highCsvReader.getLineNumber() / 0.75F + 1.0F));
             Trie<String> newDomainSuffix = new Trie<>();
+            HashMap<String, List<String>> newSubscriberTagMap = new HashMap<>();
             HighCsvReader.CsvIterator iterator = highCsvReader.getIterator();
             while (iterator.hasNext()) {
                 Map<String, String> line = iterator.next();
@@ -66,6 +70,7 @@ public class IntelligenceIndicatorKnowledgeBaseHandler extends AbstractSingleKno
                     String ip1 = line.get("ip1");
                     String ip2 = line.get("ip2");
                     String domain = line.get("domain");
+                    String subscriberId = line.get("subscriber");
                     List<String> tags = Arrays.asList(line.get("tags").split(","));
 
                     if ("IP".equals(type)) {
@@ -121,6 +126,8 @@ public class IntelligenceIndicatorKnowledgeBaseHandler extends AbstractSingleKno
                         } else {
                             logger.warn("intelligence indicator find unknown domain: " + domain);
                         }
+                    } else if ("Subscriber".equals(type)) {
+                        newSubscriberTagMap.computeIfAbsent(subscriberId, k -> new ArrayList<>()).addAll(tags);
                     }
                 } catch (Exception lineException) {
                     logger.error(this.getClass().getSimpleName() + " line: " + line.toString() + " parse error：" + lineException, lineException);
@@ -129,6 +136,7 @@ public class IntelligenceIndicatorKnowledgeBaseHandler extends AbstractSingleKno
             ipTagMap = newIpTagMap;
             domainTagMap = newDomainMap;
             domainSuffix = newDomainSuffix;
+            subscriberTagMap = newSubscriberTagMap;
         } catch (Exception e) {
             logger.error(this.getClass().getSimpleName() + " update error", e);
             return false;
@@ -155,6 +163,10 @@ public class IntelligenceIndicatorKnowledgeBaseHandler extends AbstractSingleKno
         return result;
     }
 
+    public List<String> lookupBySubscriber(String subscriber) {
+        return subscriberTagMap.computeIfAbsent(subscriber, k -> new ArrayList<>());
+    }
+
     @Override
     public void close() {
         ipTagMap.clear();
diff --git a/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/IntelligenceIndicatorLookupTest.java b/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/IntelligenceIndicatorLookupTest.java
index 34fef6b..9275812 100644
--- a/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/IntelligenceIndicatorLookupTest.java
+++ b/groot-core/src/test/java/com/geedgenetworks/core/udf/cn/IntelligenceIndicatorLookupTest.java
@@ -27,7 +27,7 @@ public class IntelligenceIndicatorLookupTest {
     void setUp() {
         runtimeContext = mockRuntimeContext();
 
-        String content = "type,ip_addr_format,ip1,ip2,domain,tags\nIP,CIDR,116.178.65.0,25,ali.com,\"阿里1,云服务1\"\nDomain,CIDR,116.178.65.0,25,$ali.com,\"阿里2,云服务2\"\nDomain,CIDR,116.178.65.0,25,*baidu.com,\"阿里3,云服务3\"\nIP,Single,116.178.65.64,116.178.65.64,ali.com,\"test\"";
+        String content = "type,ip_addr_format,ip1,ip2,domain,subscriber,tags\nIP,CIDR,116.178.65.0,25,ali.com,,\"阿里1,云服务1\"\nDomain,CIDR,116.178.65.0,25,$ali.com,,\"阿里2,云服务2\"\nDomain,CIDR,116.178.65.0,25,*baidu.com,,\"阿里3,云服务3\"\nIP,Single,116.178.65.64,116.178.65.64,ali.com,,\"test\"\nSubscriber,Single,116.178.65.64,116.178.65.64,ali.com,1234567,\"test_subscriber\"";
         mockKnowledgeBaseHandler(content);
 
         intelligenceIndicatorLookup = new IntelligenceIndicatorLookup();
@@ -155,6 +155,25 @@ public class IntelligenceIndicatorLookupTest {
         assertEquals(Arrays.asList("阿里1", "云服务1", "test"), evaluate1.getExtractedFields().get("server_ip_tags"));
     }
 
+    @Test
+    void evaluate7() {
+        UDFContext udfContext = new UDFContext();
+        Map<String, Object> parameters = new HashMap<>();
+        parameters.put("kb_name", kbName);
+        parameters.put("option", "SUBSCRIBER_TO_TAG");
+        udfContext.setParameters(parameters);
+        udfContext.setLookup_fields(Collections.singletonList("subscriber_id"));
+        udfContext.setOutput_fields(Collections.singletonList("subscriber_tags"));
+        intelligenceIndicatorLookup.open(runtimeContext, udfContext);
+
+        Event event = new Event();
+        Map<String, Object> fields = new HashMap<>();
+        fields.put("subscriber_id", "1234567");
+        event.setExtractedFields(fields);
+        Event evaluate1 = intelligenceIndicatorLookup.evaluate(event);
+        assertEquals(Arrays.asList("test_subscriber"), evaluate1.getExtractedFields().get("subscriber_tags"));
+    }
+
     @AfterEach
     void afterAll() {
         clearState();