From e27fc115dcf4cec3f0f4f1dc47c88aa0c209e218 Mon Sep 17 00:00:00 2001
From: Mike Auty <mike.auty@gmail.com>
Date: Wed, 24 Feb 2021 17:39:43 +0000
Subject: [PATCH] Layers: Improve regex scanning for linux/mac

---
 .../framework/layers/scanners/__init__.py     | 73 +++++++++++++++++--
 1 file changed, 66 insertions(+), 7 deletions(-)

diff --git a/volatility3/framework/layers/scanners/__init__.py b/volatility3/framework/layers/scanners/__init__.py
index 2364c77c9..2a167fc4e 100644
--- a/volatility3/framework/layers/scanners/__init__.py
+++ b/volatility3/framework/layers/scanners/__init__.py
@@ -1,9 +1,8 @@
 # This file is Copyright 2019 Volatility Foundation and licensed under the Volatility Software License 1.0
 # which is available at https://www.volatilityfoundation.org/license/vsl-v1.0
 #
-
 import re
-from typing import Generator, List, Tuple
+from typing import Generator, List, Tuple, Dict, Union, Optional
 
 from volatility3.framework.interfaces import layers
 from volatility3.framework.layers.scanners import multiregexp
@@ -44,19 +43,79 @@ class RegExScanner(layers.ScannerInterface):
             if offset < self.chunk_size:
                 yield offset + data_offset
 
-
 class MultiStringScanner(layers.ScannerInterface):
     thread_safe = True
 
     def __init__(self, patterns: List[bytes]) -> None:
         super().__init__()
-        self._patterns = multiregexp.MultiRegexp()
+        self._pattern_trie = {}  # type: Optional[Dict[int, Optional[Dict]]]
         for pattern in patterns:
-            self._patterns.add_pattern(pattern)
-        self._patterns.preprocess()
+            self._process_pattern(pattern)
+        self._regex = self._process_trie(self._pattern_trie)
+
+    def _process_pattern(self, value: bytes) -> None:
+        trie = self._pattern_trie
+        if trie is None:
+            return None
+
+        for char in value:
+            trie[char] = trie.get(char, {})
+            trie = trie[char]
+
+        # Mark the end of a string
+        trie[-1] = None
+
+    def _process_trie(self, trie: Optional[Dict[int, Optional[Dict]]]) -> bytes:
+        if trie is None or len(trie) == 1 and -1 in trie:
+            # We've reached the end of this path, return the empty byte string
+            return b''
+
+        choices = []
+        suffixes = []
+        finished = False
+
+        for entry in sorted(trie):
+            # Clump together different paths
+            if entry >= 0:
+                remainder = self._process_trie(trie[entry])
+                if remainder:
+                    choices.append(re.escape(bytes([entry])) + remainder)
+                else:
+                    suffixes.append(re.escape(bytes([entry])))
+            else:
+                # If we've fininshed one of the strings at this point, remember it for later
+                finished = True
+
+        if len(suffixes) == 1:
+            choices.append(suffixes[0])
+        elif len(suffixes) > 1:
+            choices.append(b'[' + b''.join(suffixes) + b']')
+
+        if len(choices) == 0:
+            # If there's none, return the empty byte string
+            response = b''
+        elif len(choices) == 1:
+            # If there's only one return it
+            response = choices[0]
+        else:
+            response = b'(?:' + b'|'.join(choices) + b')'
+
+        if finished:
+            # We finished one string, so everything after this is optional
+            response = b"(?:" + response + b")?"
+
+        return response
 
     def __call__(self, data: bytes, data_offset: int) -> Generator[Tuple[int, bytes], None, None]:
         """Runs through the data looking for the needles."""
-        for offset, pattern in self._patterns.search(data):
+        for offset, pattern in self.search(data):
             if offset < self.chunk_size:
                 yield offset + data_offset, pattern
+
+    def search(self, haystack: bytes) -> Generator[Tuple[int, bytes], None, None]:
+        if not isinstance(haystack, bytes):
+            raise TypeError("Search haystack must be a byte string")
+        if not self._regex:
+            raise ValueError("MultiRegexp cannot be used with an empty set of search strings")
+        for match in re.finditer(self._regex, haystack):
+            yield match.start(0), match.group()