0.6.1 release

dod-cyber-crime-center · Jan 4, 2023 · aee8fc4 · aee8fc4
1 parent 7e90369
commit aee8fc4
Show file tree

Hide file tree

Showing 8 changed files with 80 additions and 49 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,6 +1,11 @@
 # Changelog
 All notable changes to this project will be documented in this file.
 
+## [0.6.1] - 2022-12-20
+- Fix EOFError that can occur when running consecutive Emulator instances due to stale caching.
+- Fix IndexError bug in IMUL emulation.
+
+
 ## [0.6.0] - 2022-12-02
 - Support all instruction operands, both implied and explicit.
 - Fix bug in ROL opcode implementation.
@@ -70,7 +75,8 @@ All notable changes to this project will be documented in this file.
 - Migrated the majority of Kordesii functionality to work with Dragodis.
 
 
-[Unreleased]: https://github.com/dod-cyber-crime-center/rugosa/compare/0.6.0...HEAD
+[Unreleased]: https://github.com/dod-cyber-crime-center/rugosa/compare/0.6.1...HEAD
+[0.6.1]: https://github.com/dod-cyber-crime-center/rugosa/compare/0.6.0...0.6.1
 [0.6.0]: https://github.com/dod-cyber-crime-center/rugosa/compare/0.5.1...0.6.0
 [0.5.1]: https://github.com/dod-cyber-crime-center/rugosa/compare/0.5.0...0.5.1
 [0.5.0]: https://github.com/dod-cyber-crime-center/rugosa/compare/0.4.0...0.5.0

diff --git a/rugosa/__init__.py b/rugosa/__init__.py
@@ -5,4 +5,4 @@
 import rugosa.re as re
 import rugosa.yara as yara
 
-__version__ = "0.6.0"
+__version__ = "0.6.1"
diff --git a/rugosa/emulation/emulator.py b/rugosa/emulation/emulator.py
@@ -83,6 +83,8 @@ def __init__(self, disassembler: dragodis.Disassembler, max_instructions=10000,
         else:
             raise NotImplementedError(f"Architecture not supported: {self.arch}")
 
+        self._flowchart_cache = {}
+        self._memory_cache = {}
         self._call_hooks = call_hooks.BUILTINS.copy()
         self._instruction_hooks = collections.defaultdict(list)
         self._opcode_hooks = self._context_class.OPCODES.copy()
@@ -91,16 +93,14 @@ def __init__(self, disassembler: dragodis.Disassembler, max_instructions=10000,
         self.disabled_rep = False
         self.teleported = False
 
-    @classmethod
-    def clear_cache(cls):
+    def clear_cache(self):
         """
         Clears any internal caching during previous emulation runs.
         Calling this will be necessary if you have patched in new bytes into the disassembler and would like
         the emulator's memory object to reflect this change.
         """
-        # Necessary to import PageMap here because we could be teleported.
-        from rugosa.emulation.memory import PageMap
-        PageMap._segment_cache = {}
+        self._flowchart_cache.clear()
+        self._memory_cache.clear()
 
     def disable(self, name: str):
         """
@@ -485,7 +485,7 @@ def init_contexts():
         else:
             for context in init_contexts():
                 flowchart = self.disassembler.get_flowchart(address)
-                for path in iter_paths(flowchart, address):
+                for path in iter_paths(flowchart, address, _cache=self._flowchart_cache):
                     yield path.cpu_context(address, call_depth=call_depth, init_context=deepcopy(context))
 
                     # Don't process other paths if we are at the user call level and exhaustive wasn't chosen.

diff --git a/rugosa/emulation/flowchart.py b/rugosa/emulation/flowchart.py
@@ -24,21 +24,6 @@ class PathNode:
     This object can also track cpu context up to a certain address.
     """
 
-    _cache = {}
-
-    def __new__(cls, block: BasicBlock, prev: Optional[PathNode]):
-        """
-        Constructor that caches and reuses existing instances.
-        """
-        key = (block, prev)
-        try:
-            return cls._cache[key]
-        except KeyError:
-            self = super().__new__(cls)
-            self.__init__(block, prev)
-            cls._cache[key] = self
-            return self
-
     def __init__(self, block: BasicBlock, prev: Optional[PathNode]):
         """
         Initialize a path node.
@@ -54,15 +39,19 @@ def __init__(self, block: BasicBlock, prev: Optional[PathNode]):
         self._call_depth = 0          # the number of calls deep we are allowed to emulated
 
     @classmethod
-    def iter_all(cls, block: BasicBlock, _visited=None) -> Iterable[PathNode]:
+    def iter_all(cls, block: BasicBlock, _visited=None, _cache=None) -> Iterable[PathNode]:
         """
         Iterates all tail path nodes from a given block.
 
         :param block: Block to obtain all path nodes.
         :param _visited: Internally used.
+        :param _cache: Internally used.
 
         :yields: PathNode objects that represent the last entry of the path linked list.
         """
+        if _cache is None:
+            _cache = {}
+
         if _visited is None:
             _visited = set()
 
@@ -78,8 +67,14 @@ def iter_all(cls, block: BasicBlock, _visited=None) -> Iterable[PathNode]:
                     continue
 
                 # Create path nodes for each path of parent.
-                for parent_path in cls.iter_all(parent, _visited=_visited):
-                    yield cls(block, prev=parent_path)
+                for parent_path in cls.iter_all(parent, _visited=_visited, _cache=_cache):
+                    key = (block, parent_path)
+                    try:
+                        yield _cache[key]
+                    except KeyError:
+                        path_node = cls(block, prev=parent_path)
+                        _cache[key] = path_node
+                        yield path_node
 
         _visited.remove(block.start)
 
@@ -188,7 +183,7 @@ def cpu_context(self, addr: int = None, *, call_depth: int = 0, init_context: Pr
         return deepcopy(self._context)
 
 
-def iter_paths(flowchart: Flowchart, addr: int) -> Iterable[PathNode]:
+def iter_paths(flowchart: Flowchart, addr: int, _cache=None) -> Iterable[PathNode]:
     """
     Given an EA, iterate over the paths to the EA.
 
@@ -202,9 +197,13 @@ def iter_paths(flowchart: Flowchart, addr: int) -> Iterable[PathNode]:
 
     :param flowchart: Dragodis Flowchart to get basic blocks from.
     :param addr: Address of interest
+    :param _cache: Internally used.
 
     :yield: a path to the object
     """
+    if _cache is None:
+        _cache = {}
+
     # Obtain the block containing the address of interest
     try:
         block = flowchart.get_block(addr)
@@ -213,5 +212,4 @@ def iter_paths(flowchart: Flowchart, addr: int) -> Iterable[PathNode]:
         logger.debug(f"Unable to find block with ea: 0x{addr:08X}")
         return
 
-    yield from PathNode.iter_all(block)
-
+    yield from PathNode.iter_all(block, _cache=_cache)
diff --git a/rugosa/emulation/memory.py b/rugosa/emulation/memory.py
@@ -109,19 +109,22 @@ class PageMap(collections.defaultdict):
 
     PAGE_SIZE = 0x1000
 
-    # Cache of segment pages.
-    # Used to prevent multiple calls to pull data from the disassembler.
-    _segment_cache = {}
-
-    def __init__(self, dis: dragodis.Disassembler, map_segments=True):
+    def __init__(self, dis: dragodis.Disassembler, map_segments=True, _cache=None):
         # Setting default_factory to None, because we have overwritten it in __missing__()
         super().__init__(None)
+
+        # Cache of segment pages.
+        # Used to prevent multiple calls to pull data from the disassembler.
+        if _cache is None:
+            _cache = {}
+        self._segment_cache = _cache
+
         self._dis = dis
         if map_segments:
             self.map_segments()
 
     def __deepcopy__(self, memo):
-        copy = PageMap(self._dis, map_segments=False)
+        copy = PageMap(self._dis, map_segments=False, _cache=self._segment_cache)
         memo[id(self)] = copy
         copy.update({index: (page[:] if page is not None else None) for index, page in self.items()})
         return copy
@@ -200,16 +203,6 @@ def peek(self, page_index: int) -> bytearray:
         return self._new_page(page_index)
 
 
-def clear_cache():
-    """
-    Clears the internal cache of segment bytes.
-    Calling this will be necessary if you have patched in new bytes into the disassembler.
-    """
-    warnings.warn("clear_cache() as been moved to Emulator.clear_cache()", DeprecationWarning)
-    from rugosa.emulation.emulator import Emulator
-    Emulator.clear_cache()
-
-
 class Memory:
     """
     Class which implements the CPU memory controller backed by the segment data in the input file.
@@ -228,10 +221,10 @@ class Memory:
     MAX_MEM_READ = 0x10000000
     MAX_MEM_WRITE = 0x10000000
 
-    def __init__(self, cpu_context: ProcessorContext):
+    def __init__(self, cpu_context: ProcessorContext, _cache=None):
         """Initializes Memory object."""
         self._cpu_context = cpu_context
-        self._pages = PageMap(cpu_context.emulator.disassembler)
+        self._pages = PageMap(cpu_context.emulator.disassembler, _cache=cpu_context.emulator._memory_cache)
         # A map of base addresses to size for heap allocations.
         self._heap_base = cpu_context.emulator.disassembler.max_address
         self._heap_allocations = {}

diff --git a/rugosa/emulation/x86_64/opcodes.py b/rugosa/emulation/x86_64/opcodes.py
@@ -491,8 +491,9 @@ def _mul(cpu_context: ProcessorContext, instruction: Instruction):
     operands = instruction.operands
     width = get_max_operand_size(operands)
     mask = utils.get_mask(width)
-    multiplier1 = cpu_context.registers[RAX_REG_SIZE_MAP[width]]
-    multiplier2 = operands[1].value
+    multiplier1 = cpu_context.registers[RAX_REG_SIZE_MAP[width]]  # implied operand
+    # Pull right-most operand to handle both 1 and 2 operand instructions
+    multiplier2 = operands[-1].value
     result = multiplier1 * multiplier2
     flags = ["cf", "of"]
 

diff --git a/tests/conftest.py b/tests/conftest.py
@@ -15,6 +15,16 @@ def _get_strings_path(arch, tmp_path_factory) -> pathlib.Path:
     return new_strings_path
 
 
+@pytest.fixture
+def strings_x86(tmp_path_factory) -> pathlib.Path:
+    return _get_strings_path("x86", tmp_path_factory)
+
+
+@pytest.fixture
+def strings_arm(tmp_path_factory) -> pathlib.Path:
+    return _get_strings_path("arm", tmp_path_factory)
+
+
 @pytest.fixture(scope="function")
 def disassembler(request, tmp_path_factory) -> dragodis.Disassembler:
     """

diff --git a/tests/test_issues.py b/tests/test_issues.py
@@ -2,6 +2,9 @@
 Tests for bug fixes reported on GitHub.
 """
 
+import pytest
+
+import dragodis
 from rugosa.emulation.emulator import Emulator
 from rugosa.emulation.constants import WIDE_STRING
 
@@ -43,3 +46,23 @@ def dummy(ctx, func_name, func_args):
     assert emulator.get_call_hook("SuperFunc").__name__ == "dummy"
     assert emulator.get_call_hook("SUPERfunc").__name__ == "dummy"
     assert emulator.get_call_hook("superfunc").__name__ == "dummy"
+
+
+def test_multiple_emulations(strings_x86):
+    """
+    Tests issue with stale caching when multiple emulations happen in same process but different dragodis disassemblies.
+    """
+    address = 0x401024
+
+    try:
+        with dragodis.open_program(strings_x86, "ida") as dis:
+            emulator = Emulator(dis, teleport=False)
+            ctx = emulator.context_at(address)
+
+        with dragodis.open_program(strings_x86, "ida") as dis:
+            emulator = Emulator(dis, teleport=False)
+            ctx = emulator.context_at(address)
+    except dragodis.NotInstalledError as e:
+        pytest.skip(str(e))
+
+    # We pass if we don't get an EOFError raised.