From 3540bd402c57c92c0fdcc82288b7563c3dd43ff7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C5=82=20Praszmo?= <michalpr@cert.pl>
Date: Mon, 19 Jun 2023 15:29:26 +0200
Subject: [PATCH] Fix yara-python compatibilty break (#94)

---
 malduck/yara.py  | 42 +++++++++++++++++++++++++++---------------
 malduck/yara.pyi |  2 +-
 requirements.txt |  2 +-
 3 files changed, 29 insertions(+), 17 deletions(-)

diff --git a/malduck/yara.py b/malduck/yara.py
index c4ac5bc..06986f3 100644
--- a/malduck/yara.py
+++ b/malduck/yara.py
@@ -259,25 +259,37 @@ def _map_matches(self, matches, offset_mapper):
 
     def _map_strings(self, strings, offset_mapper):
         mapped_strings = defaultdict(list)
-        for offset, identifier, content in strings:
+        for yara_string in strings:
+            # yara-python 4.3.0 broke compatibilty and started returning a StringMatch object
+            if type(yara_string) is tuple:
+                offsets = [yara_string[0]]
+                identifier = yara_string[1]
+                contents = [yara_string[2]]
+            else:
+                offsets = [x.offset for x in yara_string.instances]
+                identifier = yara_string.identifier
+                contents = [x.matched_data for x in yara_string.instances]
+
             # Get identifier without "$" and group identifier
             real_ident, group_ident = self._parse_string_identifier(identifier)
-            # Map offset if offset_mapper is provided
-            if offset_mapper is not None:
-                _offset = offset_mapper(offset, len(content))
-                if _offset is None:
-                    # Ignore match for unmapped region
-                    continue
-                offset = _offset
-            # Register offset for full identifier
-            mapped_strings[real_ident].append(
-                YaraStringMatch(real_ident, offset, content)
-            )
-            # Register offset for grouped identifier
-            if real_ident != group_ident:
-                mapped_strings[group_ident].append(
+
+            for offset, content in zip(offsets, contents):
+                # Map offset if offset_mapper is provided
+                if offset_mapper is not None:
+                    _offset = offset_mapper(offset, len(content))
+                    if _offset is None:
+                        # Ignore match for unmapped region
+                        continue
+                    offset = _offset
+                # Register offset for full identifier
+                mapped_strings[real_ident].append(
                     YaraStringMatch(real_ident, offset, content)
                 )
+                # Register offset for grouped identifier
+                if real_ident != group_ident:
+                    mapped_strings[group_ident].append(
+                        YaraStringMatch(real_ident, offset, content)
+                    )
         return mapped_strings
 
     def _parse_string_identifier(self, identifier):
diff --git a/malduck/yara.pyi b/malduck/yara.pyi
index 6d4124f..c3ea6d6 100644
--- a/malduck/yara.pyi
+++ b/malduck/yara.pyi
@@ -20,7 +20,7 @@ from typing_extensions import Literal, Protocol
 T = TypeVar("T")
 OffsetMapper = Callable[[Optional[int], Optional[int]], Optional[int]]
 
-YaraRulesString = Tuple[int, str, bytes]
+YaraRulesString = Union[Tuple[int, str, bytes], Any]
 
 class YaraRulesMatch(Protocol):
     meta: Dict[str, str]
diff --git a/requirements.txt b/requirements.txt
index 1f910a3..b3b9436 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -3,7 +3,7 @@ pefile>=2022.5.30
 pyelftools
 pycryptodomex>=3.8.2
 capstone>=4.0.1
-yara-python==4.2.3
+yara-python
 typing-extensions>=3.7.4.2
 cryptography>=3.1
 dnfile>=0.11.0