Merge pull request #792 from elicn/dev-improv

Clean and document loader and memory modules

Author: xwings

qiling/const.py

@@ -4,7 +4,7 @@
 #
 
 from enum import EnumMeta, IntEnum
-from typing import Mapping
+from typing import Mapping, TypeVar
 
 class QL_ENDIAN(IntEnum):
     EL = 1
@@ -64,16 +64,18 @@ class QL_INTERCEPT(IntEnum):
 QL_HOOK_BLOCK = 0b0001
 QL_CALL_BLOCK = 0b0010
 
-def __reverse_enum(e: EnumMeta) -> Mapping[str, int]:
+__QL_CE = TypeVar('__QL_CE', QL_DEBUGGER, QL_ARCH, QL_OS, QL_VERBOSE)
+
+def __reverse_enum(e: EnumMeta) -> Mapping[str, __QL_CE]:
     '''Create a reverse mapping for an enum.
     '''
 
-    return dict((k.lower(), v.value) for k, v in e.__members__.items())
+    return dict((v.name.lower(), v.value) for v in e.__members__.values())
 
-debugger_map = __reverse_enum(QL_DEBUGGER)
-arch_map     = __reverse_enum(QL_ARCH)
-os_map       = __reverse_enum(QL_OS)
-verbose_map  = __reverse_enum(QL_VERBOSE)
+debugger_map: Mapping[str, QL_DEBUGGER] = __reverse_enum(QL_DEBUGGER)
+arch_map    : Mapping[str, QL_ARCH]     = __reverse_enum(QL_ARCH)
+os_map      : Mapping[str, QL_OS]       = __reverse_enum(QL_OS)
+verbose_map : Mapping[str, QL_VERBOSE]  = __reverse_enum(QL_VERBOSE)
 
 loader_map = {
     QL_OS.LINUX   : "ELF",

qiling/core.py

@@ -691,7 +691,7 @@ def _init_stop_guard(self):
         # Allocate a guard page, we need this in both cases
         # On a negative stack pointer, we still need a return address (otherwise we end up at 0)
         # Make sure it is not close to the heap (PE), otherwise the heap cannot grow
-        self._exit_trap_addr = self.mem.find_free_space(0x1000, min_addr=0x9000000, alignment=0x10)
+        self._exit_trap_addr = self.mem.find_free_space(0x1000, minaddr=0x9000000, align=0x10)
         self.mem.map(self._exit_trap_addr, 0x1000, info='[Stop guard]')
 
         # Stop on a negative stack pointer

qiling/loader/elf.py

@@ -15,7 +15,7 @@
 from qiling.const import *
 
 from qiling.exception import *
-from .loader import QlLoader
+from .loader import QlLoader, Image
 from qiling.os.linux.function_hook import FunctionHook
 from qiling.os.linux.syscall_nums import SYSCALL_NR
 from qiling.os.linux.kernel_api.hook import *
@@ -422,7 +422,7 @@ def load_with_ld(self, stack_addr, load_address=-1, argv=[], env={}):
         self.ql.os.elf_entry = self.elf_entry = load_address + elfhead['e_entry']
         self.stack_address = new_stack
         self.load_address = load_address
-        self.images.append(self.coverage_image(load_address, load_address + mem_end, self.path))
+        self.images.append(Image(load_address, load_address + mem_end, self.path))
         self.ql.os.function_hook = FunctionHook(self.ql, self.elf_phdr + mem_start, self.elf_phnum, self.elf_phent,
                                                 load_address, load_address + mem_end)
         self.init_sp = self.ql.reg.arch_sp

qiling/loader/loader.py

@@ -2,27 +2,27 @@
 # 
 # Cross Platform and Multi Architecture Advanced Binary Emulation Framework
 #
-from collections import namedtuple
 
-from qiling.const import QL_OS, QL_ARCH, QL_ENDIAN
-from qiling.exception import QlErrorArch, QlErrorOsType, QlErrorOutput
+from typing import Any, Mapping, MutableSequence, NamedTuple
 
+from qiling import Qiling
+
+Image = NamedTuple('Image', (('base', int), ('end', int), ('path', str)))
 
 class QlLoader():
-    def __init__(self, ql):
-        self.ql     = ql
-        self.env    = self.ql.env
-        self.argv   = self.ql.argv
-        self.images = []
-        self.coverage_image = namedtuple('Image', 'base end path')
+    def __init__(self, ql: Qiling):
+        self.ql = ql
+        self.env = self.ql.env
+        self.argv = self.ql.argv
+        self.images: MutableSequence[Image] = []
         self.skip_exit_check = False
-    
-    def save(self):
-        saved_state = {}
-        saved_state['images'] = list(map(tuple, self.images))
-        return saved_state
 
-    def restore(self, saved_state):
-        for (base, end, path) in saved_state['images']:
-            self.images.append(self.coverage_image(base, end, path))
+    def save(self) -> Mapping[str, Any]:
+        saved_state = {
+            'images': [tuple(img) for img in self.images]
+        }
+
+        return saved_state
 
+    def restore(self, saved_state: Mapping[str, Any]):
+        self.images = [Image(*img) for img in saved_state['images']]

qiling/loader/pe.py

@@ -12,7 +12,7 @@
 from qiling.exception import *
 from qiling.const import *
 from qiling.arch.x86_const import *
-from .loader import QlLoader
+from .loader import QlLoader, Image
 from qiling.os.memory import QlMemoryHeap
 
 
@@ -141,7 +141,7 @@ def load_dll(self, dll_name, driver=False):
             self.add_ldr_data_table_entry(dll_name)
 
         # add DLL to coverage images
-        self.images.append(self.coverage_image(dll_base, dll_base+dll_len, path))
+        self.images.append(Image(dll_base, dll_base+dll_len, path))
 
         self.ql.log.info("Done with loading %s" % path)
 
@@ -472,7 +472,7 @@ def load(self):
             self.sizeOfStackReserve = self.pe.OPTIONAL_HEADER.SizeOfStackReserve
             self.ql.log.info("Loading %s to 0x%x" % (self.path, self.pe_image_address))
             self.ql.log.info("PE entry point at 0x%x" % self.entry_point)
-            self.images.append(self.coverage_image(self.pe_image_address, self.pe_image_address + self.pe.NT_HEADERS.OPTIONAL_HEADER.SizeOfImage, self.path))
+            self.images.append(Image(self.pe_image_address, self.pe_image_address + self.pe.NT_HEADERS.OPTIONAL_HEADER.SizeOfImage, self.path))
 
             # Stack should not init at the very bottom. Will cause errors with Dlls
             sp = self.stack_address + self.stack_size - 0x1000

qiling/loader/pe_uefi.py

@@ -9,7 +9,7 @@
 from qiling import Qiling
 from qiling.const import QL_ARCH
 from qiling.exception import QlErrorArch, QlMemoryMappedError
-from qiling.loader.loader import QlLoader
+from qiling.loader.loader import QlLoader, Image
 
 from qiling.os.uefi import context, st, smst
 from qiling.os.uefi.ProcessorBind import CPU_STACK_ALIGNMENT
@@ -122,7 +122,7 @@ def map_and_load(self, path: str, exec_now: bool=False):
         self.install_loaded_image_protocol(image_base, image_size)
 
         # this would be used later be os.find_containing_image
-        self.images.append(self.coverage_image(image_base, image_base + image_size, path))
+        self.images.append(Image(image_base, image_base + image_size, path))
 
         # update next memory slot to allow sequencial loading. its availability
         # is unknown though