Add config merge in extract_manager

malduck/extractor/extract_manager.py

@@ -20,6 +20,42 @@
 
 __all__ = ["ExtractManager"]
 
+def merge_configs(base_config: Config, new_config: Config) -> Config:
+    """
+    Merge static configurations.
+
+    :param base_config: Base configuration
+    :param new_config: Changes to apply
+    :return: Merged configuration
+    """
+    config = dict(base_config)
+    for k, v in new_config.items():
+        if k == "family":
+            continue
+        if k not in config:
+            config[k] = v
+        elif config[k] == v:
+            continue
+        elif type(config[k]) == type(v):
+            if isinstance(config[k], list):
+                for el in v:
+                    if el not in config[k]:
+                        config[k] = config[k] + [el]
+            elif isinstance(config[k], dict):
+                config[k] = merge_configs(config[k], v)
+            else:
+                log.warning(
+                    f"Merge error for key '{k}' of type {type(v)}: Merging not implemented, "
+                    "using value from base_config"
+                )
+        else:
+            log.warning(
+                f"TypeError for key '{k}': Failed to merge existing value ({type(config[k])}) "
+                f"'{config[k]}' with new value (new type: {type(v)}) '{v}', "
+                "using value from base_config"
+            )
+    return config
+
 
 class ExtractManager:
     """
@@ -136,13 +172,18 @@ def push_config(self, config: Config) -> bool:
 
         family = config["family"]
         if family in self.configs:
+            base_config = None
+            new_config = None
             if is_config_better(base_config=self.configs[family], new_config=config):
-                self.configs[family] = config
-                log.debug("%s config looks better than previous one", family)
-                return True
+                log.debug(f"{family} new config looks better - use as base and merge existing")
+                base_config = config
+                new_config = self.configs[family]
             else:
-                log.debug("%s config doesn't look better than previous one", family)
-                return False
+                log.debug(f"{family} new config doesn't look better - try merging into existing.")
+                base_config = self.configs[family]
+                new_config = config
+            self.configs[family] = merge_configs(base_config=base_config, new_config=new_config)
+            return True
 
         if family in self.modules.override_paths:
             # 'citadel' > 'zeus'

tests/files/modules/configmerge/__init__.py

@@ -0,0 +1 @@
+from .configmerge import ConfigMerge

tests/files/modules/configmerge/configmerge.py

@@ -0,0 +1,16 @@
+from malduck.extractor import Extractor
+
+
+class ConfigMerge(Extractor):
+    yara_rules = "configmerge",
+    family = "ConfigMerge"
+
+    @Extractor.final
+    def final(self, p):
+        return {
+            "constant": "CONST",
+            "mem_types": [str(type(p))],
+            "dict": {
+                hex(p.imgbase): "imagebase"
+            }
+        }

tests/files/modules/configmerge/configmerge.yar

@@ -0,0 +1,6 @@
+rule configmerge {
+    strings:
+        $calc_exe_0x80000 = { E5 ED 2D 7A 0E 1D 32 DB 8E 73 56 1E 1C 95 93 A4 }
+    condition:
+        any of them
+}

tests/test_extractor.py

@@ -80,3 +80,24 @@ def test_multirules():
         'matched': ['v2'],
         'third': ['ThIrD string']
     }]
+
+
+def test_configmerge():
+    modules = ExtractorModules("tests/files/modules")
+    calc_exe_path = os.path.join(os.path.abspath(os.path.dirname(__file__)), "files", "calc.exe")
+    extractor = ExtractManager(modules)
+    extractor.push_file(calc_exe_path)
+    assert len(extractor.config) == 1
+
+    conf = extractor.config[0]
+    assert conf == {
+        'family': "ConfigMerge",
+        'constant': "CONST",
+        'mem_types': [str(procmem), str(procmempe)],
+        'dict': {
+            '0x0': "imagebase",
+            '0x1000000': "imagebase"
+        }
+    }
+
+