From 001d53e7684b0a33f74f8ff67aa08a970ca58690 Mon Sep 17 00:00:00 2001 From: Mike Auty Date: Thu, 1 Nov 2018 11:17:15 +0000 Subject: [PATCH] Rework some code to help with typing. --- volatility/framework/automagic/pdbscan.py | 23 +++++++++++------------ volatility/framework/objects/__init__.py | 8 ++++---- 2 files changed, 15 insertions(+), 16 deletions(-) diff --git a/volatility/framework/automagic/pdbscan.py b/volatility/framework/automagic/pdbscan.py index 04edf1611..b3b2ff052 100644 --- a/volatility/framework/automagic/pdbscan.py +++ b/volatility/framework/automagic/pdbscan.py @@ -10,7 +10,6 @@ import os import struct import typing -from volatility import symbols from volatility.framework import constants, exceptions, layers, validity from volatility.framework.configuration import requirements from volatility.framework.interfaces import configuration @@ -27,7 +26,7 @@ from volatility.framework import interfaces vollog = logging.getLogger(__name__) ValidKernelsType = typing.Dict[str, typing.Tuple[int, typing.Dict]] -KernelsType = typing.List[typing.Dict[str, typing.Any]] +KernelsType = typing.Iterable[typing.Dict[str, typing.Any]] class PdbSignatureScanner(interfaces.layers.ScannerInterface): @@ -237,12 +236,12 @@ class KernelPDBScanner(interfaces.automagic.AutomagicInterface): def method_fixed_mapping(self, context: interfaces.context.ContextInterface, - virtual_layer_name: str, + vlayer: layers.intel.Intel, kernels: KernelsType, progress_callback: validity.ProgressCallback = None) -> ValidKernelsType: # TODO: Verify this is a windows image valid_kernels = {} - vlayer = context.memory[virtual_layer_name] # type: layers.intel.Intel + virtual_layer_name = vlayer.name physical_layer_name = self.get_physical_layer_name(context, vlayer) kvo_path = interfaces.configuration.path_join(vlayer.config_path, 'kernel_virtual_offset') for kernel in kernels: @@ -272,14 +271,14 @@ class KernelPDBScanner(interfaces.automagic.AutomagicInterface): def method_module_offset(self, context: interfaces.context.ContextInterface, - virtual_layer_name: str, + vlayer: layers.intel.Intel, kernels: KernelsType, progress_callback: validity.ProgressCallback = None) -> ValidKernelsType: """Method for finding a suitable kernel offset based on a module table""" valid_kernels = {} vollog.debug("Kernel base randomized, searching layer for base address offset") # If we're here, chances are high we're in a Win10 x64 image with kernel base randomization - vlayer = context.memory[virtual_layer_name] # type: layers.intel.Intel + virtual_layer_name = vlayer.name physical_layer_name = self.get_physical_layer_name(context, vlayer) physical_layer = context.memory[physical_layer_name] # TODO: On older windows, this might be \WINDOWS\system32\nt rather than \SystemRoot\system32\nt @@ -308,12 +307,12 @@ class KernelPDBScanner(interfaces.automagic.AutomagicInterface): def method_kdbg_offset(self, context: interfaces.context.ContextInterface, - virtual_layer_name: str, + vlayer: layers.intel.Intel, kernels: KernelsType, progress_callback: validity.ProgressCallback = None) -> ValidKernelsType: valid_kernels = {} vollog.debug("Kernel base randomized, using KDBG structure for kernel offset") - vlayer = context.memory[virtual_layer_name] # type: layers.intel.Intel + virtual_layer_name = vlayer.name physical_layer_name = self.get_physical_layer_name(context, vlayer) physical_layer = context.memory[physical_layer_name] results = physical_layer.scan(context, scanners.BytesScanner(b"KDBG"), progress_callback = progress_callback) @@ -363,13 +362,13 @@ class KernelPDBScanner(interfaces.automagic.AutomagicInterface): Returns: A dictionary of valid kernels """ - valid_kernels = {} + valid_kernels = {} # type: ValidKernelsType for virtual_layer_name in potential_kernels: kernels = list(potential_kernels[virtual_layer_name]) - vlayer = context.memory[virtual_layer_name] - if virtual_layer_name and isinstance(vlayer, layers.intel.Intel): + vlayer = context.memory.get(virtual_layer_name, None) + if isinstance(vlayer, layers.intel.Intel): for method in self.methods: - valid_kernels = method(self, context, virtual_layer_name, kernels, progress_callback) + valid_kernels = method(self, context, vlayer, kernels, progress_callback) if valid_kernels: break if not valid_kernels: diff --git a/volatility/framework/objects/__init__.py b/volatility/framework/objects/__init__.py index de913c1c4..12affe022 100644 --- a/volatility/framework/objects/__init__.py +++ b/volatility/framework/objects/__init__.py @@ -43,13 +43,13 @@ class PrimitiveObject(interfaces.objects.ObjectInterface): struct_format = struct_format) self._struct_format = struct_format - def __new__(cls, + def __new__(cls: typing.Type, context: interfaces.context.ContextInterface, type_name: str, object_info: interfaces.objects.ObjectInformation, struct_format: str, new_value: typing.Union[int, float, bool, bytes, str] = None, - **kwargs) -> 'PrimitiveObject': + **kwargs) -> typing.Type['PrimitiveObject']: """Creates the appropriate class and returns it so that the native type is inherited The only reason the **kwargs is added, is so that the inherriting types can override __init__ @@ -136,12 +136,12 @@ class Bytes(PrimitiveObject, bytes): struct_format = str(length) + "s") self._vol['length'] = length - def __new__(cls, + def __new__(cls: typing.Type, context: interfaces.context.ContextInterface, type_name: str, object_info: interfaces.objects.ObjectInformation, length: int = 1, - **kwargs) -> 'Bytes': + **kwargs) -> typing.Type['Bytes']: """Creates the appropriate class and returns it so that the native type is inherritted The only reason the **kwargs is added, is so that the inherriting types can override __init__