refactor: refactor the node of tags

5 files changed, 46 insertions, 33 deletions

diff --git a/platform-etl/src/main/java/com/zdjizhi/etl/utils/tag/AppTagUtils.java b/platform-etl/src/main/java/com/zdjizhi/etl/utils/tag/AppTagUtils.java
index 6f4bc09..040d55e 100644
--- a/platform-etl/src/main/java/com/zdjizhi/etl/utils/tag/AppTagUtils.java
+++ b/platform-etl/src/main/java/com/zdjizhi/etl/utils/tag/AppTagUtils.java
@@ -21,6 +21,7 @@ public class AppTagUtils {
     public static void updateTag(Map<Long, byte[]> contents) {
         Map<String, Node> newAppLabelTagRules = new HashMap<>();
         for (Long s : contents.keySet()) {
+            BaseTagUtils.addLibraryType(s, BaseTagUtils.LibraryType.APP);
             parseAppLabelTagRules(newAppLabelTagRules, s, contents.get(s));
         }
         appLabelTagRules = newAppLabelTagRules;
@@ -63,12 +64,10 @@ public class AppTagUtils {
 
                 if (appTags.containsKey(appName)) {
                     Node node = appTags.get(appName);
-                    node.getTags().add(tagValue);
-                    node.put(id, "app");
+                    node.add(tagValue, id);
                 } else {
                     Node node = new Node();
-                    node.getTags().add(tagValue);
-                    node.put(id, "app");
+                    node.add(tagValue, id);
                     appTags.put(appName, node);
                 }
             }
diff --git a/platform-etl/src/main/java/com/zdjizhi/etl/utils/tag/BaseTagUtils.java b/platform-etl/src/main/java/com/zdjizhi/etl/utils/tag/BaseTagUtils.java
index 3963db5..c488f02 100644
--- a/platform-etl/src/main/java/com/zdjizhi/etl/utils/tag/BaseTagUtils.java
+++ b/platform-etl/src/main/java/com/zdjizhi/etl/utils/tag/BaseTagUtils.java
@@ -15,6 +15,12 @@ public class BaseTagUtils {
     private static final Logger logger = LoggerFactory.getLogger(BaseTagUtils.class);
     private static final Map<Long, List<Long>> libraryIdRuleIdMap = new HashMap<>();
 
+    private static final Map<Long, LibraryType> libraryTypeMap = new HashMap<>();
+
+    public static void addLibraryType(Long libraryId, LibraryType libraryType) {
+        libraryTypeMap.put(libraryId, libraryType);
+    }
+
     public static void updateInternalRuleId(Long key, RuleInfo ruleInfo) {
         logger.info("tagUtils update rule: " + ruleInfo.getName() + ":" + ruleInfo.getRuleId());
         long ruleId = ruleInfo.getRuleId();
@@ -43,16 +49,31 @@ public class BaseTagUtils {
     }
 
     protected static void setRuleIdAndIocType(CnRecordLog entity, Node node) {
-        for (int i = 0; i < node.knowledgeIds.size(); i++) {
-            Long libraryId = node.knowledgeIds.get(i);
-            String libraryType = node.knowledgeTypes.get(i);
+        node.knowledgeIds.forEach(libraryId -> {
+            String libraryType = libraryTypeMap.get(libraryId).name;
             if (libraryIdRuleIdMap.containsKey(libraryId)) {
                 for (Long ruleId : libraryIdRuleIdMap.get(libraryId)) {
                     entity.getRule_id_list().add(ruleId);
                     entity.getIoc_type_list().add(libraryType);
                 }
             }
+        });
+    }
+
+
+    protected enum LibraryType {
+        APP("app"),
+        DOMAIN("domain"),
+        IP("ip");
+
+        private final String name;
+
+        LibraryType(String name) {
+            this.name = name;
         }
 
+        public String getName() {
+            return this.name;
+        }
     }
 }
diff --git a/platform-etl/src/main/java/com/zdjizhi/etl/utils/tag/DomainTagUtils.java b/platform-etl/src/main/java/com/zdjizhi/etl/utils/tag/DomainTagUtils.java
index ceac90d..3244d07 100644
--- a/platform-etl/src/main/java/com/zdjizhi/etl/utils/tag/DomainTagUtils.java
+++ b/platform-etl/src/main/java/com/zdjizhi/etl/utils/tag/DomainTagUtils.java
@@ -25,6 +25,7 @@ public class DomainTagUtils {
         Trie<Node> newFqdnTagFuzzyMatchRules = new Trie<>();
         Map<String, Node> newFqdnTagFullMatchRules = new HashMap<>();
         for (Long s : contents.keySet()) {
+            BaseTagUtils.addLibraryType(s, BaseTagUtils.LibraryType.DOMAIN);
             parseFqdnTagRules(newFqdnTagFuzzyMatchRules, newFqdnTagFullMatchRules, s, contents.get(s));
         }
         fqdnTagFuzzyMatchRules = newFqdnTagFuzzyMatchRules;
@@ -80,18 +81,15 @@ public class DomainTagUtils {
                 String newRule = rule.substring(1);
                 if (rule.startsWith("*")) {
                     Node node = new Node();
-                    node.getTags().add(tagValue);
-                    node.put(id, "domain");
+                    node.add(tagValue, id);
                     fqdnTagsFuzzy.put(StringUtils.reverse(newRule), node);
                 } else {
                     if (fqdnTagsFull.containsKey(newRule)) {
                         Node node = fqdnTagsFull.get(newRule);
-                        node.getTags().add(tagValue);
-                        node.put(id, "domain");
+                        node.add(tagValue, id);
                     } else {
                         Node node = new Node();
-                        node.getTags().add(tagValue);
-                        node.put(id, "domain");
+                        node.add(tagValue, id);
                         fqdnTagsFull.put(newRule, node);
                     }
                 }
diff --git a/platform-etl/src/main/java/com/zdjizhi/etl/utils/tag/IpTagUtils.java b/platform-etl/src/main/java/com/zdjizhi/etl/utils/tag/IpTagUtils.java
index e94711b..a070695 100644
--- a/platform-etl/src/main/java/com/zdjizhi/etl/utils/tag/IpTagUtils.java
+++ b/platform-etl/src/main/java/com/zdjizhi/etl/utils/tag/IpTagUtils.java
@@ -5,7 +5,6 @@ import com.zdjizhi.etl.utils.csv.HighCsvReader;
 import inet.ipaddr.IPAddress;
 import inet.ipaddr.IPAddressString;
 import org.apache.flink.shaded.guava18.com.google.common.collect.Range;
-import org.apache.flink.shaded.guava18.com.google.common.collect.RangeMap;
 import org.apache.flink.shaded.guava18.com.google.common.collect.TreeRangeMap;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -25,6 +24,7 @@ public class IpTagUtils {
     public static void updateTag(Map<Long, byte[]> contents) {
         TreeRangeMap<IPAddress, Node> newIpTagRules = TreeRangeMap.create();
         for (Long s : contents.keySet()) {
+            BaseTagUtils.addLibraryType(s, BaseTagUtils.LibraryType.IP);
             parseIpTagRules(newIpTagRules, s, contents.get(s));
         }
         ipTagRules = newIpTagRules;
@@ -71,8 +71,8 @@ public class IpTagUtils {
                     String ip1 = line.get("ip1");
                     String ip2 = line.get("ip2");
 
-                    IPAddress startIpAddress = null;
-                    IPAddress endIpAddress = null;
+                    IPAddress startIpAddress;
+                    IPAddress endIpAddress;
                     if ("Single".equals(addrFormat)) {
                         IPAddress ipAddress = new IPAddressString(ip1).getAddress();
                         if (ipAddress == null) {
@@ -102,19 +102,15 @@ public class IpTagUtils {
                         continue;
                     }
 
-                    RangeMap<IPAddress, Node> ipAddressListRangeMap = treeRangeMap.subRangeMap(Range.closed(startIpAddress, endIpAddress));
-                    Map<Range<IPAddress>, Node> rangeListMap = ipAddressListRangeMap.asMapOfRanges();
+                    Map<Range<IPAddress>, Node> rangeListMap = treeRangeMap.subRangeMap(Range.closed(startIpAddress, endIpAddress)).asMapOfRanges();
                     TreeRangeMap<IPAddress, Node> subRangeMap = TreeRangeMap.create();
                     Node node = new Node();
-                    node.getTags().add(tagValue);
-                    node.put(id, "ip");
+                    node.add(tagValue, id);
                     subRangeMap.put(Range.closed(startIpAddress, endIpAddress), node);
-                    for (Range<IPAddress> ipAddressRange : rangeListMap.keySet()) {
-                        Node ipAddressNode = ipAddressListRangeMap.get(ipAddressRange.lowerEndpoint());
-                        ipAddressNode.getTags().add(tagValue);
-                        ipAddressNode.put(id, "ip");
+                    rangeListMap.forEach((ipAddressRange, ipAddressNode) -> {
+                        ipAddressNode.add(tagValue, id);
                         subRangeMap.put(ipAddressRange, ipAddressNode);
-                    }
+                    });
                     treeRangeMap.putAll(subRangeMap);
                 } catch (Exception lineException) {
                     logger.error("ip tag line: " + line.toString() + " parse error：" + lineException, lineException);
diff --git a/platform-etl/src/main/java/com/zdjizhi/etl/utils/tag/Node.java b/platform-etl/src/main/java/com/zdjizhi/etl/utils/tag/Node.java
index af2f4dc..f22d931 100644
--- a/platform-etl/src/main/java/com/zdjizhi/etl/utils/tag/Node.java
+++ b/platform-etl/src/main/java/com/zdjizhi/etl/utils/tag/Node.java
@@ -1,30 +1,29 @@
 package com.zdjizhi.etl.utils.tag;
 
 import java.util.ArrayList;
+import java.util.HashSet;
 import java.util.List;
+import java.util.Set;
 
 public class Node {
-    private List<String> tags = new ArrayList<>();
-    protected List<Long> knowledgeIds = new ArrayList<>();
-
-    protected List<String> knowledgeTypes = new ArrayList<>();
+    private final List<String> tags = new ArrayList<>();
+    protected Set<Long> knowledgeIds = new HashSet<>();
 
     public List<String> getTags() {
         return tags;
     }
 
-    public void put(Long id, String type) {
-        if (id == null || type == null) {
+    public void add(String tag, Long id) {
+        if (tag == null) {
             return;
         }
+        tags.add(tag);
         knowledgeIds.add(id);
-        knowledgeTypes.add(type);
     }
 
     public void merge(Node node) {
         this.tags.addAll(node.getTags());
         this.knowledgeIds.addAll(node.knowledgeIds);
-        this.knowledgeTypes.addAll(node.knowledgeTypes);
     }
 
     public void mergeAll(List<Node> nodes) {