diff --git a/volatility/framework/layers/scanners/__init__.py b/volatility/framework/layers/scanners/__init__.py index c8f39c6d7..21e78f559 100644 --- a/volatility/framework/layers/scanners/__init__.py +++ b/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 diff --git a/volatility/framework/layers/scanners/wumanber.py b/volatility/framework/layers/scanners/wumanber.py index 91800338d..9adbc8a4e 100644 --- a/volatility/framework/layers/scanners/wumanber.py +++ b/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")