Update detect_duplicate_features.py

Using get_rules menthod to get set of all existing rules.

scripts/detect_duplicate_features.py

@@ -1,11 +1,16 @@
-import os
 import argparse
 
+import capa.main
 import capa.rules
 import capa.engine as ceng
 
 
-def get_child_features(feature):
+def get_child_features(feature) -> list:
+    """
+    args:
+    \tfeature : capa.rule.Rule.statement containing feature statements
+    returns a list containg all the features in the rule
+    """
     children = []
 
     if isinstance(feature, (ceng.And, ceng.Or, ceng.Some)):
@@ -19,43 +24,45 @@ def get_child_features(feature):
 
 
 def get_features(rule_path):
-    error = ""
+    """
+    args:
+    \tfeature : rule path
+    returns a list containg all the features in the rule
+    """
     feature_list = []
     with open(rule_path, "r") as f:
         try:
             new_rule = capa.rules.Rule.from_yaml(f.read())
             feature_list = get_child_features(new_rule.statement)
         except Exception as e:
-            error = "rule :" + rule_path + " " + str(type(e)) + " " + str(e)
-    return feature_list, error
+            raise Warning("Error: " + rule_path + " " + str(type(e)) + " " + str(e))
+    return feature_list
 
 
 def find_overlapping_rules(new_rule_path, rules_path):
     if not new_rule_path.endswith(".yml"):
         raise FileNotFoundError("FileNotFoundError ! New rule file name doesn't end with yml")
 
-    new_rule_features, error = get_features(new_rule_path)
-    if error:
-        raise Warning(error)
+    # Loads features of new rule in a list.
+    new_rule_features = get_features(new_rule_path)
 
-    errors: list = []
     count = 0
     overlapping_rules = []
-    for rules in rules_path:
-        for dirpath, dirnames, filenames in os.walk(rules):
-            for filename in filenames:
-                if filename.endswith(".yml"):
-                    rule_path = os.path.join(dirpath, filename)
-                    rule_features, error = get_features(rule_path)
-                    if error:
-                        errors.append(error)
-                    if not len(rule_features):
-                        continue
-                    count += 1
-                    if any([feature in rule_features for feature in new_rule_features]):
-                        overlapping_rules.append(rule_path)
 
-    result = {"overlapping_rules": overlapping_rules, "count": count, "errors": errors}
+    # capa.rules.RuleSet stores all rules in given paths
+    ruleset = capa.main.get_rules(rules_path)
+
+    for rule_name, rule in ruleset.rules.items():
+        rule_features = get_child_features(rule.statement)
+
+        if not len(rule_features):
+            continue
+        count += 1
+        # Checks if any features match between existing and new rule.
+        if any([feature in rule_features for feature in new_rule_features]):
+            overlapping_rules.append(rule_name)
+
+    result = {"overlapping_rules": overlapping_rules, "count": count}
     return result
 
 
@@ -73,15 +80,14 @@ def main():
         result = find_overlapping_rules(new_rule_path, rules_path)
         print("\nNew rule path : %s" % new_rule_path)
         print("Number of rules checked : %s " % result["count"])
-        print("Paths to overlapping rules : ")
-        for r in result["overlapping_rules"]:
-            print(r)
+        if result["overlapping_rules"]:
+            print("Paths to overlapping rules : ")
+            for r in result["overlapping_rules"]:
+                print("- %s" % r)
+        else:
+            print("Paths to overlapping rules : None")
         print("Number of rules containing same features : %s" % len(result["overlapping_rules"]))
-        if result["errors"]:
-            print("\nWhile checking following .yml files error occured:")
-            for error in result["errors"]:
-                print(error)
-            print("\n")
+        print("\n")
     except Exception as e:
         print(e)