Add in type annotations for scanners.

volatility/framework/layers/scanners/__init__.py

@@ -1,4 +1,5 @@
 import re
+import typing
 
 from volatility.framework.interfaces import layers
 from volatility.framework.layers.scanners import wumanber
@@ -7,11 +8,11 @@ from volatility.framework.layers.scanners import wumanber
 class BytesScanner(layers.ScannerInterface):
     thread_safe = True
 
-    def __init__(self, needle):
+    def __init__(self, needle: bytes) -> None:
         super().__init__()
         self.needle = self._check_type(needle, bytes)
 
-    def __call__(self, data, data_offset):
+    def __call__(self, data: bytes, data_offset: int) -> typing.Generator[int, None, None]:
         """Runs through the data looking for the needle, and yields all offsets where the needle is found
         """
         find_pos = data.find(self.needle)
@@ -24,15 +25,14 @@ class RegExScanner(layers.ScannerInterface):
     # TODO: Document why this isn't thread safe?
     thread_safe = False
 
-    def __init__(self, pattern, flags = 0):
+    def __init__(self, pattern: bytes, flags: int = 0) -> None:
         super().__init__()
         self.regex = re.compile(self._check_type(pattern, bytes), self._check_type(flags, int))
 
-    def __call__(self, data, data_offset):
+    def __call__(self, data: bytes, data_offset: int) -> typing.Generator[int, None, None]:
         """Runs through the data looking for the needle, and yields all offsets where the needle is found
         """
         find_pos = self.regex.finditer(data)
-        find_pos = list(find_pos)
         for match in find_pos:
             offset = match.start()
             yield offset + data_offset
@@ -41,7 +41,7 @@ class RegExScanner(layers.ScannerInterface):
 class MultiStringScanner(layers.ScannerInterface):
     thread_safe = True
 
-    def __init__(self, patterns):
+    def __init__(self, patterns: typing.List[bytes]) -> None:
         super().__init__()
         self._check_type(patterns, list)
         self._patterns = wumanber.WuManber()
@@ -50,7 +50,8 @@ class MultiStringScanner(layers.ScannerInterface):
             self._patterns.add_pattern(pattern)
         self._patterns.preprocess()
 
-    def __call__(self, data, data_offset):
+    def __call__(self, data: bytes, data_offset: int) \
+            -> typing.Generator[typing.Tuple[int, typing.Union[str, bytes]], None, None]:
         """Runs through the data looking for the needles"""
         for offset, pattern in self._patterns.search(data):
             yield offset + data_offset, pattern

volatility/framework/layers/scanners/wumanber.py

@@ -1,17 +1,20 @@
+import typing
+
+
 class WuManber(object):
     """Algorithm for multi-string matching"""
 
-    def __init__(self, block_size = 3):
-        self.minimum_pattern_length = None
+    def __init__(self, block_size: int = 3) -> None:
+        self.minimum_pattern_length = None  # type: typing.Optional[int]
 
         self._block_size = block_size
-        self._maximum_hash = self._hash_function([0xff, 0xff, 0xff]) + 1  # This depends on the hash function used
+        self._maximum_hash = self._hash_function(b"\xff\xff\xff") + 1  # This depends on the hash function used
 
-        self._patterns = []
-        self._shift = None  # This gets generated by preprocess
-        self._hashes = [set() for _ in range(self._maximum_hash)]
+        self._patterns = []  # type: typing.List[bytes]
+        self._shift = None  # type: typing.Optional[typing.List[int]]
+        self._hashes = [set() for _ in range(self._maximum_hash)]  # type: typing.List[typing.Set[bytes]]
 
-    def add_pattern(self, pattern):
+    def add_pattern(self, pattern: bytes) -> None:
         if not isinstance(pattern, bytes):
             raise TypeError("Pattern must be a byte string")
 
@@ -20,7 +23,7 @@ class WuManber(object):
 
         self._patterns.append(pattern)
 
-    def preprocess(self):
+    def preprocess(self) -> None:
         """Preprocesses the patterns by populating the three arrays"""
 
         if not self._patterns:
@@ -31,7 +34,7 @@ class WuManber(object):
 
         max_jump = self.minimum_pattern_length - self._block_size + 1
         self._shift = [max_jump] * self._maximum_hash
-        self.hashes = [set() for _ in range(self._maximum_hash)]
+        self.hashes = [set() for _ in range(self._maximum_hash)]  # type: typing.List[typing.Set[bytes]]
 
         for pattern in self._patterns:
             for i in range(self._block_size, self.minimum_pattern_length + 1):
@@ -41,14 +44,15 @@ class WuManber(object):
                 if self.minimum_pattern_length - i == 0:
                     self._hashes[hashval].add(pattern)
 
-    def _hash_function(self, value_bytes):
+    def _hash_function(self, value_bytes: bytes) -> int:
         """Hash function to bucket _block_size number of bytes into sets
 
         If this hash_function changes, the maximum number of responses must be set in self._maximum_hash
         """
         return (value_bytes[0] << 5) + (value_bytes[1] << 3) + value_bytes[2]
 
-    def search(self, haystack):
+    def search(self, haystack: bytes) \
+            -> typing.Generator[typing.Tuple[int, typing.Union[str, bytes]], None, None]:
         """Search through a large body of data for patterns previously added with add_pattern"""
         if not isinstance(haystack, bytes):
             raise TypeError("Search haystack must be a byte string")