Add memory storage capabilities to `ProgramState`

The `SparseMemory` class represents a program's memory. While the user
can read from and write to arbitrary memory addresses, it manages its
memory in pages/chunks internally. This is a tradeoff between space
consumption (this solution might have a memory overhead) and lookup
speed of individual memory addresses.

Add two small unit tests for `SparseMemory`.

1 files changed, 48 insertions, 54 deletions

diff --git a/miasm_util.py b/miasm_util.py
index 4f20dd8..31083d9 100644
--- a/miasm_util.py
+++ b/miasm_util.py
@@ -1,43 +1,36 @@
-from angr_targets.memory_map import MemoryMap
-from miasm.core.locationdb import LocationDB
+from miasm.core.locationdb import LocationDB, LocKey
 from miasm.expression.expression import Expr, ExprOp, ExprId, ExprLoc, \
                                         ExprInt, ExprMem, ExprCompose, \
                                         ExprSlice, ExprCond
 from miasm.expression.simplifications import expr_simp_explicit
 
-class MiasmProgramState:
-    def __init__(self,
-                 regs: dict[ExprId, ExprInt],
-                 mem: list[tuple[MemoryMap, bytes]]
-                 ):
-        self.regs = regs
-        self.memory = mem
-
-    def _find_mem_map(self, addr: int) \
-            -> tuple[MemoryMap, bytes] | tuple[None, None]:
-        for m, data in self.memory:
-            if addr >= m.start_address and addr < m.end_address:
-                return m, data
-        return None, None
-
-    def read_memory(self, addr: int, size: int) -> bytes:
-        res = bytes()
-        while size > 0:
-            m, data = self._find_mem_map(addr)
-            if m is None:
-                raise AttributeError(f'No memory mapping contains the address {addr}.')
-
-            assert(m is not None and data is not None)
-            read_off = addr - m.start_address
-            read_size = min(size, m.end_address - addr)
-            assert(read_off + read_size <= len(data))
-            res += data[read_off:read_off+read_size]
-
-            size -= read_size
-            addr += read_size
-        return res
-
-def eval_expr(expr: Expr, conc_state: MiasmProgramState, loc_db) -> int:
+from snapshot import ProgramState
+
+class MiasmConcreteState:
+    miasm_flag_aliases = {
+        'NF':     'SF',
+        'I_F':    'IF',
+        'IOPL_F': 'IOPL',
+        'I_D':    'ID',
+    }
+
+    def __init__(self, state: ProgramState, loc_db: LocationDB):
+        self.state = state
+        self.loc_db = loc_db
+
+    def resolve_register(self, regname: str) -> int:
+        regname = regname.upper()
+        if regname in self.miasm_flag_aliases:
+            regname = self.miasm_flag_aliases[regname]
+        return self.state.read(regname)
+
+    def resolve_memory(self, addr: int, size: int) -> bytes:
+        return self.state.read_memory(addr, size)
+
+    def resolve_location(self, loc: LocKey) -> int | None:
+        return self.loc_db.get_location_offset(loc)
+
+def eval_expr(expr: Expr, conc_state: MiasmConcreteState) -> int:
     # Most of these implementation are just copy-pasted members of
     # `SymbolicExecutionEngine`.
     expr_to_visitor = {
@@ -55,28 +48,29 @@ def eval_expr(expr: Expr, conc_state: MiasmProgramState, loc_db) -> int:
     if visitor is None:
         raise TypeError("Unknown expr type")
 
-    ret = visitor(expr, conc_state, loc_db)
+    ret = visitor(expr, conc_state)
     ret = expr_simp_explicit(ret)
     assert(ret is not None)
 
     return ret
 
-def _eval_exprint(expr: ExprInt, _, __: LocationDB):
+def _eval_exprint(expr: ExprInt, _):
     """Evaluate an ExprInt using the current state"""
     return expr
 
-def _eval_exprid(expr: ExprId, state: MiasmProgramState, _):
+def _eval_exprid(expr: ExprId, state: MiasmConcreteState):
     """Evaluate an ExprId using the current state"""
-    return state.regs[expr]
+    val = state.resolve_register(expr.name)
+    return ExprInt(val, expr.size)
 
-def _eval_exprloc(expr: ExprLoc, _, loc_db: LocationDB):
+def _eval_exprloc(expr: ExprLoc, state: MiasmConcreteState):
     """Evaluate an ExprLoc using the current state"""
-    offset = loc_db.get_location_offset(expr.loc_key)
+    offset = state.resolve_location(expr.loc_key)
     if offset is not None:
         return ExprInt(offset, expr.size)
     return expr
 
-def _eval_exprmem(expr: ExprMem, state: MiasmProgramState, loc_db: LocationDB):
+def _eval_exprmem(expr: ExprMem, state: MiasmConcreteState):
     """Evaluate an ExprMem using the current state.
     This function first evaluates the memory pointer value.
     """
@@ -87,34 +81,34 @@ def _eval_exprmem(expr: ExprMem, state: MiasmProgramState, loc_db: LocationDB):
     assert(expr.size <= 64)
     assert(expr.size % 8 == 0)
 
-    addr = eval_expr(expr.ptr, state, loc_db)
-    ret = state.read_memory(int(addr), int(expr.size / 8))
+    addr = eval_expr(expr.ptr, state)
+    ret = state.resolve_memory(int(addr), int(expr.size / 8))
     assert(len(ret) * 8 == expr.size)
     return ExprInt(int.from_bytes(ret, byteorder='little'), expr.size)
 
-def _eval_exprcond(expr, state: MiasmProgramState, loc_db: LocationDB):
+def _eval_exprcond(expr, state: MiasmConcreteState):
     """Evaluate an ExprCond using the current state"""
-    cond = eval_expr(expr.cond, state, loc_db)
-    src1 = eval_expr(expr.src1, state, loc_db)
-    src2 = eval_expr(expr.src2, state, loc_db)
+    cond = eval_expr(expr.cond, state)
+    src1 = eval_expr(expr.src1, state)
+    src2 = eval_expr(expr.src2, state)
     return ExprCond(cond, src1, src2)
 
-def _eval_exprslice(expr, state: MiasmProgramState, loc_db: LocationDB):
+def _eval_exprslice(expr, state: MiasmConcreteState):
     """Evaluate an ExprSlice using the current state"""
-    arg = eval_expr(expr.arg, state, loc_db)
+    arg = eval_expr(expr.arg, state)
     return ExprSlice(arg, expr.start, expr.stop)
 
-def _eval_exprop(expr, state: MiasmProgramState, loc_db: LocationDB):
+def _eval_exprop(expr, state: MiasmConcreteState):
     """Evaluate an ExprOp using the current state"""
     args = []
     for oarg in expr.args:
-        arg = eval_expr(oarg, state, loc_db)
+        arg = eval_expr(oarg, state)
         args.append(arg)
     return ExprOp(expr.op, *args)
 
-def _eval_exprcompose(expr, state: MiasmProgramState, loc_db: LocationDB):
+def _eval_exprcompose(expr, state: MiasmConcreteState):
     """Evaluate an ExprCompose using the current state"""
     args = []
     for arg in expr.args:
-        args.append(eval_expr(arg, state, loc_db))
+        args.append(eval_expr(arg, state))
     return ExprCompose(*args)