1 files changed, 41 insertions, 10 deletions

diff --git a/miasm_util.py b/miasm_util.py
index 3ceebea..0d3ab3d 100644
--- a/miasm_util.py
+++ b/miasm_util.py
@@ -41,13 +41,13 @@ class MiasmConcreteState:
         self.state = state
         self.loc_db = loc_db
 
-    def resolve_register(self, regname: str) -> int:
+    def resolve_register(self, regname: str) -> int | None:
         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:
+    def resolve_memory(self, addr: int, size: int) -> bytes | None:
         return self.state.read_memory(addr, size)
 
     def resolve_location(self, loc: LocKey) -> int | None:
@@ -95,14 +95,18 @@ def _eval_exprint(expr: ExprInt, _):
 def _eval_exprid(expr: ExprId, state: MiasmConcreteState):
     """Evaluate an ExprId using the current state"""
     val = state.resolve_register(expr.name)
-    return ExprInt(val, expr.size)
+    if val is None:
+        return expr
+    if isinstance(val, int):
+        return ExprInt(val, expr.size)
+    return val
 
 def _eval_exprloc(expr: ExprLoc, state: MiasmConcreteState):
     """Evaluate an ExprLoc using the current state"""
     offset = state.resolve_location(expr.loc_key)
-    if offset is not None:
-        return ExprInt(offset, expr.size)
-    return expr
+    if offset is None:
+        return expr
+    return ExprInt(offset, expr.size)
 
 def _eval_exprmem(expr: ExprMem, state: MiasmConcreteState):
     """Evaluate an ExprMem using the current state.
@@ -116,10 +120,15 @@ def _eval_exprmem(expr: ExprMem, state: MiasmConcreteState):
     assert(expr.size % 8 == 0)
 
     addr = eval_expr(expr.ptr, state)
-    ret = state.resolve_memory(int(addr), int(expr.size / 8))
-    assert(len(ret) * 8 == expr.size)
-    ival = ExprInt(int.from_bytes(ret, byteorder='little'), expr.size)
-    return ExprSlice(ival, 0, len(ret) * 8)
+    if not addr.is_int():
+        return expr
+
+    mem = state.resolve_memory(int(addr), int(expr.size / 8))
+    if mem is None:
+        return expr
+
+    assert(len(mem) * 8 == expr.size)
+    return ExprInt(int.from_bytes(mem, byteorder='little'), expr.size)
 
 def _eval_exprcond(expr, state: MiasmConcreteState):
     """Evaluate an ExprCond using the current state"""
@@ -133,12 +142,34 @@ def _eval_exprslice(expr, state: MiasmConcreteState):
     arg = eval_expr(expr.arg, state)
     return ExprSlice(arg, expr.start, expr.stop)
 
+def _eval_cpuid(rax: ExprInt, out_reg: ExprInt):
+    """Evaluate the `x86_cpuid` operator by performing a real invocation of
+    the CPUID instruction.
+
+    :param rax:     The current value of RAX. Must be concrete.
+    :param out_reg: An index in `[0, 4)` signaling which register's value
+                    shall be returned. Must be concrete.
+    """
+    from cpuid import cpuid
+
+    regs = cpuid.CPUID()(int(rax))
+
+    if int(out_reg) >= len(regs):
+        raise ValueError(f'Output register may not be {out_reg}.')
+    return ExprInt(regs[int(out_reg)], out_reg.size)
+
 def _eval_exprop(expr, state: MiasmConcreteState):
     """Evaluate an ExprOp using the current state"""
     args = []
     for oarg in expr.args:
         arg = eval_expr(oarg, state)
         args.append(arg)
+
+    if expr.op == 'x86_cpuid':
+        # Can't do this in an expression simplifier plugin because the
+        # arguments must be concrete.
+        assert(len(expr.args) == 2)
+        return _eval_cpuid(args[0], args[1])
     return ExprOp(expr.op, *args)
 
 def _eval_exprcompose(expr, state: MiasmConcreteState):