Merge pull request #441 from InQuest/issue/440-mass-import-not-commit…

…ting-yara-rules #440 : Fixing issue with import not committing due to numeric metadat…
InQuest · Oct 26, 2022 · c6effc0 · c6effc0
2 parents 164bae4 + d2e34df
commit c6effc0
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 5 deletions.
diff --git a/app/models/cfg_category_range_mapping.py b/app/models/cfg_category_range_mapping.py
@@ -1,5 +1,6 @@
 from app import db
 from . import yara_rule
+from sqlalchemy import bindparam
 
 
 class CfgCategoryRangeMapping(db.Model):
@@ -36,7 +37,7 @@ def to_dict(self, include_inactive=False):
         )
 
     @staticmethod
-    def get_next_category_eventid(category=None):
+    def get_next_category_eventid(category=None, connection=None):
         default_category_min = 10000
         default_category_max = 20000
 
@@ -49,6 +50,22 @@ def get_next_category_eventid(category=None):
                     category = CfgCategoryRangeMapping(category=CfgCategoryRangeMapping.DEFAULT_CATEGORY,
                                                        range_max=default_category_max,
                                                        range_min=default_category_min, current=default_category_min)
+                    if connection:
+                        transaction = connection.begin()
+                        connection.execute(CfgCategoryRangeMapping.__table__.insert().values(
+                            category=bindparam("category"),
+                            range_max=bindparam("range_max"),
+                            range_min=bindparam("range_min"),
+                            current=bindparam("current")
+                        ), {
+                            "category": category.category,
+                            "range_max": category.range_max,
+                            "range_min": category.range_min,
+                            "current": category.current
+                        })
+                        transaction.commit()
+                        category = CfgCategoryRangeMapping.query.filter(
+                            CfgCategoryRangeMapping.category == CfgCategoryRangeMapping.DEFAULT_CATEGORY).first()
                     CfgCategoryRangeMapping.COMMITTED_DEFAULT = category
                 else:
                     category = CfgCategoryRangeMapping.COMMITTED_DEFAULT
@@ -67,6 +84,5 @@ def get_next_category_eventid(category=None):
             "update `cfg_category_range_mapping` set current=%s where id=%s" % (category.current, category.id))
         return eventid
 
-
     def __repr__(self):
         return '<CfgCategoryRangeMapping %r>' % self.id
diff --git a/app/models/cfg_states.py b/app/models/cfg_states.py
@@ -53,7 +53,7 @@ def verify_state(state_to_verify):
 
 @listens_for(Cfg_states, "before_insert")
 def generate_eventid(mapper, connect, target):
-    if target.is_release_state > 0:
+    if not target.is_release_state or target.is_release_state > 0:
         Cfg_states.query.update(dict(is_release_state=0))
 
 

diff --git a/app/models/yara_rule.py b/app/models/yara_rule.py
@@ -370,7 +370,7 @@ def get_yara_rule_from_yara_dict(cls, yara_dict, metadata_field_mapping={}):
         if type(yara_dict["metadata"]) is list:
             yara_dict["metadata"] = {list(m.keys())[0]: list(m.values())[0] for m in yara_dict["metadata"]}
 
-        yara_metadata = {key.lower(): val.strip().strip("\"") for key, val in
+        yara_metadata = {key.lower(): str(val).strip().strip("\"") for key, val in
                          yara_dict["metadata"].items()} if "metadata" in yara_dict else {}
         for possible_field, mapped_to in metadata_field_mapping.items():
             mapped_to = mapped_to.lower()
@@ -431,7 +431,7 @@ def generate_eventid(mapper, connect, target):
     target.name = re.sub("[^A-Za-z0-9_]", "", target.name)
 
     if not target.eventid:
-        target.eventid = CfgCategoryRangeMapping.get_next_category_eventid(target.category)
+        target.eventid = CfgCategoryRangeMapping.get_next_category_eventid(target.category, connect)
 
 
 @listens_for(Yara_rule, "before_update")

diff --git a/app/routes/import_.py b/app/routes/import_.py
@@ -30,6 +30,12 @@ def save_artifacts(extract_ip, extract_dns, extract_signature, artifacts, shared
     unique_rule_name_enforcement = cfg_settings.Cfg_settings.get_setting("ENFORCE_UNIQUE_YARA_RULE_NAMES")
     retired_state = cfg_states.Cfg_states.query.filter(cfg_states.Cfg_states.is_retired_state > 0).first()
 
+    if not retired_state:
+        retired_state = cfg_states.Cfg_states(state="Retired", is_retired_state=1)
+        db.session.add(retired_state)
+        db.session.commit()
+        db.session.refresh(retired_state)
+
     if not cfg_states.Cfg_states.query.filter_by(state=default_state).first():
         db.session.add(cfg_states.Cfg_states(state=default_state))
         db.session.commit()