Moved the check_argument_list_var() funcitonality

- removed it from the Helper class and added it to MetadataAttributes
CybercentreCanada · Sep 10, 2020 · 0819629 · 0819629
1 parent 261cf2c
commit 0819629
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 17 deletions.
diff --git a/yara-validator/validator_functions.py b/yara-validator/validator_functions.py
@@ -372,7 +372,7 @@ def valid_mitre_att(self, rule_to_validate_mitre_att, metadata_index, metadata_k
             self.required_fields[MITRE_ATT].attributeinvalid()
 
         if self.required_fields[MITRE_ATT].valid and mitre_att_to_validate.startswith('S'):
-            soft_codes_found = Helper.check_argument_list_var(self.required_fields, MITRE_ATT, MITRE_SOFTWAREID_FOUND)
+            soft_codes_found = self.required_fields[MITRE_ATT].check_argument_list_var(MITRE_SOFTWAREID_FOUND)
             soft_codes_found.append(mitre_att_to_validate)
 
         return self.required_fields[MITRE_ATT].valid
@@ -490,8 +490,8 @@ def valid_category_type(self, rule_to_validate_type, metadata_index, metadata_ke
                                                                                     rule_category_key_to_check):
             malware_id = Helper.get_software_id_by_name(rule_category_value_to_check)
             if malware_id:
-                malware_ids_found = Helper.check_argument_list_var(self.required_fields, child_metadata_place_holder,
-                                                                   MITRE_SOFTWAREID_GEN)
+                malware_ids_found = self.required_fields[child_metadata_place_holder] \
+                                        .check_argument_list_var(MITRE_SOFTWAREID_GEN)
                 malware_ids_found.append(malware_id)
 
         return self.required_fields[child_metadata_place_holder].valid
@@ -602,8 +602,9 @@ def mitre_software_generator(self, rule_to_generate_mitre_att, category_key, mit
         MITRE_SOFTWAREID_FOUND = 'mitre_softwareid_found'
         child_metadata_place_holder = self.required_fields[CATEGORY].argument.get(CHILD_PLACE_HOLDER)
 
-        malware_ids_found = Helper.check_argument_list_var(self.required_fields, child_metadata_place_holder, MITRE_SOFTWAREID_GEN)
-        soft_codes_found = Helper.check_argument_list_var(self.required_fields, MITRE_ATT, MITRE_SOFTWAREID_FOUND)
+        malware_ids_found = self.required_fields[child_metadata_place_holder]\
+                                .check_argument_list_var(MITRE_SOFTWAREID_GEN)
+        soft_codes_found = self.required_fields[MITRE_ATT].check_argument_list_var(MITRE_SOFTWAREID_FOUND)
 
         for malware_id_found in malware_ids_found:
             if malware_id_found not in soft_codes_found:
@@ -941,13 +942,3 @@ def get_software_id_by_name(software_name):
         else:
             return ''
 
-    @staticmethod
-    def check_argument_list_var(required_fields, metadata, variable_name):
-        if not required_fields[metadata].argument:
-            required_fields[metadata].argument = {variable_name: []}
-        elif not required_fields[metadata].argument.get(variable_name):
-            required_fields[metadata].argument.update({variable_name: []})
-        elif not isinstance(required_fields[metadata].argument.get(variable_name), list):
-            required_fields[metadata].argument.update({variable_name: []})
-
-        return required_fields[metadata].argument.get(variable_name)
diff --git a/yara-validator/yara_validator.py b/yara-validator/yara_validator.py
@@ -8,7 +8,7 @@
 # for querying the MITRE ATT&CK data
 from stix2 import FileSystemSource
 
-from validator_functions import Validators, MetadataOpt, StringEncoding, check_encoding, Helper
+from validator_functions import Validators, MetadataOpt, StringEncoding, check_encoding
 from yara_file_processor import YaraFileProcessor, YaraRule
 
 # set current working directory
@@ -304,6 +304,14 @@ def attributereset(self):
         self.found = False
         self.valid = False
 
+    def check_argument_list_var(self, variable_name):
+        if not self.argument or not isinstance(self.argument, dict):
+            self.argument = {variable_name: []}
+        elif not self.argument.get(variable_name) or not isinstance(self.argument.get(variable_name), list):
+            self.argument.update({variable_name: []})
+
+        return self.argument.get(variable_name)
+
 
 class Positional:
     """
@@ -631,7 +639,7 @@ def return_req_optional(self, rule_to_validate):
             keys_to_return.append(self.required_fields[ACTOR].argument.get(CHILD_PLACE_HOLDER))
 
         category_type = self.required_fields[CATEGORY].argument.get(CHILD_PLACE_HOLDER)
-        if Helper.check_argument_list_var(self.required_fields, category_type, MITRE_SOFTWAREID_GEN):
+        if self.required_fields[category_type].check_argument_list_var(MITRE_SOFTWAREID_GEN):
             self.validators.mitre_software_generator(rule_to_validate, CATEGORY, MITRE_ATT)
 
         return keys_to_return