smallworld-re 1.0.0__py3-none-any.whl

smallworld/analyses/unstable/angr/nwbt.py

@@ -0,0 +1,63 @@
+import angr
+
+from ....emulators.angr.exploration import (
+    BoundedExplorationMixin,
+    TerminationExplorationMixin,
+)
+from ....emulators.angr.memory import TrackerMemoryMixin
+from .divergence import DivergenceExplorationMixin, DivergenceMemoryMixin
+from .model import ModelMemoryMixin
+from .typedefs import TypeDefPlugin
+
+
+class NWBTMemoryPlugin(  # type: ignore[misc]
+    DivergenceMemoryMixin,
+    TrackerMemoryMixin,
+    ModelMemoryMixin,
+    angr.storage.DefaultMemory,
+):
+    pass
+
+
+class NWBTExplorationTechnique(
+    TerminationExplorationMixin,
+    BoundedExplorationMixin,
+    DivergenceExplorationMixin,
+    angr.exploration_techniques.suggestions.Suggestions,
+):
+    pass
+
+
+def configure_nwbt_plugins(emu):
+    """Configure NWBT analysis plugins.
+
+    This creates a new plugin preset that overrides
+    angr's default symbolic memory plugin.
+    This preset is passed to the entry state constructor,
+    and can't be changed afterward.
+    Thus, this needs to get called in a preinit callback.
+    """
+    preset = angr.SimState._presets["default"].copy()
+    preset.add_default_plugin("sym_memory", NWBTMemoryPlugin)
+    preset.add_default_plugin("typedefs", TypeDefPlugin)
+    emu._plugin_preset = preset
+
+
+def configure_nwbt_strategy(emu):
+    """Configure NWBT analysis strategies
+
+    This overrides the default angr exploration strategy.
+    This needs to access the instantiated exploration manager,
+    so it needs to get called in an init callback.
+    """
+    emu.mgr.use_technique(NWBTExplorationTechnique())
+    # angr bug: states don't inherit plugin presets.
+    #
+    # Normally, they copy all plugins from their parents.
+    # If you define a custom plugin and don't touch it,
+    # it never gets initialized, and won't get inherited.
+    #
+    # If you try to touch that plugin on a successor,
+    # it can't be initialized, since it doesn't have
+    # your custom preset.
+    emu.state.get_plugin("typedefs")

smallworld/analyses/unstable/angr/typedefs.py

@@ -0,0 +1,170 @@
+import ctypes
+import logging
+
+import angr
+import claripy
+
+log = logging.getLogger(__name__)
+
+
+class TypeDefPlugin(angr.SimStatePlugin):
+    """Angr state plugin for tracking typedefs
+
+    This captures two concepts for a state:
+
+    - Type bindings for addresses, registers and symbols
+    - Instance addresses for symbols
+    """
+
+    def __init__(self):
+        super().__init__()
+        self._structdefs = dict()
+
+        self._symbols = dict()
+        self._sym_values = dict()
+        self._registers = dict()
+        self._addrs = dict()
+        self._allocations = dict()
+
+    @angr.SimStatePlugin.memo
+    def copy(self, memo):
+        dup = TypeDefPlugin()
+        dup._structdefs.update(self._structdefs)
+        dup._symbols.update(self._symbols)
+        dup._sym_values.update(self._sym_values)
+        dup._registers.update(self._registers)
+
+        # heckin' multi-layer copy...
+        for typedef, addrs in self._allocations.items():
+            dup._allocations[typedef] = set(addrs)
+
+        return dup
+
+    def ctype_to_bv(self, typedef, prefix="sw"):
+        if hasattr(typedef, "_fields_"):
+            off = 0
+            fielddefs = list()
+            for name, fielddef in typedef._fields_:
+                field_off = getattr(typedef, name).offset
+                if field_off > off:
+                    # handle padding
+                    fielddefs.append(claripy.BVS("padding", (field_off - off) * 8))
+                    off = field_off
+                fielddefs.append(self.ctype_to_bv(fielddef, f"{prefix}_{name}"))
+                off += ctypes.sizeof(fielddef)
+
+            return claripy.Concat(*fielddefs)
+        else:
+            res = claripy.BVS(prefix, ctypes.sizeof(typedef) * 8)
+            self.bind_symbol(res, typedef)
+            res = claripy.Reverse(res)
+            return res
+
+    def add_struct(self, structdef):
+        """Add a struct type to this context for easy lookup"""
+        if structdef._name in self._structdefs:
+            raise ValueError(f"Already have a struct def for {structdef._name}")
+        self._structdefs[structdef._name] = structdef
+
+    def get_struct(self, name):
+        """Look up a struct type by its name"""
+        return self._structdefs[name]
+
+    def bind_symbol(self, sym, typedef):
+        """Bind a symbol to a typedef"""
+        if sym in self._symbols:
+            raise ValueError(f"Already have a binding for {sym}")
+        self._symbols[sym.args[0]] = typedef
+
+    def get_symbol_binding(self, sym):
+        """Get the type binding for a symbol"""
+        if sym in self._symbols:
+            return self._symbols[sym]
+        return None
+
+    def set_symbol(self, sym, val):
+        """Set a concrete value for a symbol
+
+        Symbols are static-single-assignment,
+        so if they are collapsed to a concrete value,
+        that assignment will stick.
+        The assignment is normally bound up in the solver,
+        making it annoying and expensive to recover.
+        This feature lets us quickly recover our stuff
+        """
+        log.info(f"Setting {sym} to {val}")
+        if sym in self._sym_values:
+            raise ValueError(f"Already have value {self._sym_values[sym]} for {sym}")
+        self._sym_values[sym] = val
+
+    def get_symbol(self, sym):
+        """Get a concrete value for a symbol"""
+        if sym in self._sym_values:
+            return self._sym_values[sym]
+        return None
+
+    def bind_register(self, name, typedef):
+        """Bind a register to a typedef
+
+        Registers themselves don't have specific types.
+        However, this only gets used in uninitialized value analysis;
+        it only covers the type a register has before it's
+        written by the program.
+        """
+        if name in self._registers:
+            raise ValueError(f"Already have a binding for {name}")
+        self._registers[name] = typedef
+
+    def get_register_binding(self, name):
+        """Get the type binding for a register
+
+        Registers themselves don't have specific types.
+        However, this only gets used in uninitialized value analysis;
+        it only covers the type a register has before it's
+        written by the program.
+        """
+        if name in self._registers:
+            return self._registers[name]
+        return None
+
+    def bind_address(self, addr, typedef):
+        if addr in self._addrs:
+            raise ValueError(f"Already have a binding for {addr}")
+        self._addrs[addr] = typedef
+
+    def get_address_binding(self, addr):
+        if addr in self._addrs:
+            return self._addrs[addr]
+        return None
+
+    def allocate(self, typedef):
+        addr = self.state.heap.allocate(ctypes.sizeof(typedef))
+
+        rep = self.ctype_to_bv(typedef)
+
+        self._allocations.setdefault(typedef, set()).add(addr)
+        log.debug(f"Storing {rep} at {addr:x}")
+        self.state.memory.store(addr, rep)
+        return addr
+
+    def get_allocs_for_type(self, typedef):
+        if typedef in self._allocations:
+            return list(self._allocations[typedef])
+        return list()
+
+    def pp(self, log, all_allocs=True):
+        log("Register Bindings")
+        for reg_name, reg_def in self._registers.items():
+            log(f"\t{reg_name} -> {reg_def}")
+        log("Symbol Bindings")
+        for sym_name, sym_def in self._symbols.items():
+            if sym_name in self._sym_values:
+                log(f"\t{sym_name} -> {sym_def} ({self._sym_values[sym_name]:x})")
+            else:
+                log(f"\t{sym_name} -> {sym_def} (No Value)")
+        if all_allocs:
+            log("Allocations:")
+            for typedef, addrs in self._allocations.items():
+                log(f"- {typedef}")
+                for addr in addrs:
+                    log(f"\t- {addr:x}")

smallworld/analyses/unstable/angr/utils.py

@@ -0,0 +1,25 @@
+import angr
+
+
+def print_state(log, state, tag):
+    """Pretty-print data contained in an execution state
+
+    This prints useful registers, memory, and state constraints.
+    """
+    log(f"{state}: {tag}")
+    if not state.regs.rip.symbolic:
+        try:
+            block = state.block()
+            block.pp()
+        except angr.errors.SimEngineError:
+            log("(Invalid IP!)")
+    else:
+        log("(Exit state)")
+    log("Registers")
+    state.registers.pp(log)
+    log("Memory")
+    state.memory.pp(log)
+    log("Constraints:")
+    for expr in state.solver.constraints:
+        log(f"\t{expr}")
+    state.typedefs.pp(log)

smallworld/analyses/unstable/angr/visitor.py

@@ -0,0 +1,315 @@
+import logging
+
+import claripy
+
+log = logging.getLogger(__name__)
+
+
+class ClaripyVisitor:
+    """Abstract class for a claripy AST visitor"""
+
+    def visit_int(self, v, **kwargs):
+        raise NotImplementedError(f"{type(self)} missing Integer")
+
+    def visit_bvv(self, v, **kwargs):
+        raise NotImplementedError(f"{type(self)} missing BVV")
+
+    def visit_bvs(self, v, **kwargs):
+        raise NotImplementedError(f"{type(self)} missing BVS")
+
+    def visit_add(self, v, **kwargs):
+        raise NotImplementedError(f"{type(self)} missing Add")
+
+    def visit_mul(self, v, **kwargs):
+        raise NotImplementedError(f"{type(self)} missing Mul")
+
+    def visit_bwor(self, v, **kwargs):
+        raise NotImplementedError(f"{type(self)} missing BWOr")
+
+    def visit_concat(self, v, **kwargs):
+        raise NotImplementedError(f"{type(self)} missing Concat")
+
+    def visit_extract(self, v, **kwargs):
+        raise NotImplementedError(f"{type(self)} missing Extract")
+
+    def visit_if(self, v, **kwargs):
+        raise NotImplementedError(f"{type(self)} missing If")
+
+    def visit_reverse(self, v, **kwargs):
+        raise NotImplementedError(f"{type(self)} missing Reverse")
+
+    def visit_lshift(self, v, **kwargs):
+        raise NotImplementedError(f"{type(self)} missing LShift")
+
+    def visit_rshift(self, v, **kwargs):
+        raise NotImplementedError(f"{type(self)} missing RShift")
+
+    def visit_zeroext(self, v, **kwargs):
+        raise NotImplementedError(f"{type(self)} missing ZeroExt")
+
+    def visit(self, v, **kwargs):
+        """Perform the visitor pattern on a claripy expression."""
+        # TODO: This is very far from complete.
+        # claripy supports a dizzying array of bitvector operations.
+        # It's easier to populate them as I find them.
+        if isinstance(v, int):
+            # v is a python integer.
+            # I occasionally see an int instead of a BVV.
+            # I don't exactly know why.
+            return self.visit_int(v, **kwargs)
+        elif v.op == "BVV":
+            # v is a concrete value
+            return self.visit_bvv(v, **kwargs)
+        elif v.op == "BVS":
+            # v is a symbol
+            return self.visit_bvs(v, **kwargs)
+        elif v.op == "__add__":
+            # v = sum(*X), for two or more expressions x in X.
+            return self.visit_add(v, **kwargs)
+        elif v.op == "__mul__":
+            # v = prod(*X), for two or more expressions x in X.
+            return self.visit_mul(v, **kwargs)
+        elif v.op == "__or__":
+            # v = BWOr(*X) for two or more expressions x in X.
+            return self.visit_bwor(v, **kwargs)
+        elif v.op == "Concat":
+            # v = Concat(*X), for two or more expressions x in X.
+            return self.visit_concat(v, **kwargs)
+        elif v.op == "If":
+            # v is ITE(x, y, z)
+            return self.visit_if(v, **kwargs)
+        elif v.op == "Reverse":
+            # v is Reverse(x)
+            return self.visit_reverse(v, **kwargs)
+        elif v.op == "Extract":
+            # v is x[a:b], for expression x and ints a and b.
+            return self.visit_extract(v, **kwargs)
+        elif v.op == "__lshift__":
+            # v is x << y, for expressions x and y
+            return self.visit_lshift(v, **kwargs)
+        elif v.op == "LShR":
+            # v is x >>> y, for expressions x and y
+            return self.visit_rshift(v, **kwargs)
+        elif v.op == "ZeroExt":
+            # v is ZeroExt(v, a) for expression x and int a
+            return self.visit_zeroext(v, **kwargs)
+        else:
+            # v is something I haven't thought of yet.
+            raise NotImplementedError(f"Unknown op {v.op}")
+
+
+class ConditionalVisitor(ClaripyVisitor):
+    """Visitor that splits a conditional expression into its possible evaluations."""
+
+    def visit_int(self, v):
+        # Integers are leaf ASTs, and never conditional.
+        return [(v, claripy.BoolV(True))]
+
+    def visit_bvv(self, v):
+        # BVVs are leaf ASTs, and never conditional.
+        return [(v, claripy.BoolV(True))]
+
+    def visit_bvs(self, v):
+        # BVSs are leaf ASTs, and never conditional.
+        return [(v, claripy.BoolV(True))]
+
+    def visit_add(self, v):
+        # Addition can produce all combinations of evaluations
+        # of the argument expressions.
+        out = self.visit(v.args[0])
+        for arg in v.args[1:]:
+            old_out = out
+            out = list()
+            for res2, guard2 in self.visit(arg):
+                for res1, guard1 in old_out:
+                    res = res1 + res2
+                    guard = claripy.And(guard1, guard2)
+                    out.append((res, guard))
+        return out
+
+    def visit_mul(self, v):
+        # Multiplication can produce all combinations of evaluations
+        # of the argument expressions.
+        out = self.visit(v.args[0])
+        for arg in v.args[1:]:
+            old_out = out
+            out = list()
+            for res2, guard2 in self.visit(arg):
+                for res1, guard1 in old_out:
+                    res = res1 * res2
+                    guard = claripy.And(guard1, guard2)
+                    out.append((res, guard))
+        return out
+
+    def visit_concat(self, v):
+        # Concatenation can produce all combinations of evaluations
+        # of the argument expressions.
+        out = self.visit(v.args[0])
+        for arg in v.args[1:]:
+            old_out = out
+            out = list()
+            for res2, guard2 in self.visit(arg):
+                for res1, guard1 in old_out:
+                    res = res1.concat(res2)
+                    guard = claripy.And(guard1, guard2)
+                    out.append((res, guard))
+        return out
+
+    def visit_extract(self, v):
+        # Extraction produces one expression per evaluation
+        # of the main argument.  The other two are always ints.
+        a = v.args[0]
+        b = v.args[1]
+        res = self.visit(v.args[2])
+        res = list(map(lambda x: (x[0][a:b], x[1]), res))
+        return res
+
+    def visit_if(self, v):
+        # ITE produces the union of the results of
+        # the "then" and "else" expressions.
+        # The condition itself is ignored.
+        out = list()
+        out.extend(
+            map(lambda x: (x[0], claripy.And(x[1], v.args[0])), self.visit(v.args[1]))
+        )
+        out.extend(
+            map(
+                lambda x: (x[0], claripy.And(x[1], claripy.Not(v.args[0]))),
+                self.visit(v.args[2]),
+            )
+        )
+        return out
+
+    def visit_reverse(self, v):
+        # Reversal produces one expression per evaluation of the argument.
+        return list(map(lambda x: (claripy.Reverse(x[0]), x[1]), self.visit(v.args[0])))
+
+
+class EvalVisitor(ClaripyVisitor):
+    """Visitor that concretizes a symbolic expression.
+
+    Uses a mapping from variables to concrete values.
+    This implementation requires a complete mapping;
+    it will raise an exception if it encounters
+    a symbol for which it does not have a value.
+    """
+
+    def visit_int(self, v, bindings=None):
+        return v
+
+    def visit_bvv(self, v, bindings=None):
+        return v
+
+    def visit_bvs(self, v, bindings=None):
+        if v.args[0] not in bindings:
+            raise KeyError("Missing binding for {v}")
+        return bindings[v.args[0]]
+
+    def visit_add(self, v, bindings=None):
+        # Eval of sum is sum of evals
+        out = self.visit(v.args[0], bindings=bindings)
+        for x in v.args[1:]:
+            out += self.visit(x, bindings=bindings)
+        return out
+
+    def visit_concat(self, v, bindings=None):
+        # Eval of concat is concat of evals
+        return claripy.Concat(map(lambda x: self.visit(x, bindings=bindings), v.args))
+
+    def visit_extract(self, v, bindings=None):
+        # Extract only hase one BV argument; the range limits are ints.
+        a = v.args[0]
+        b = v.args[1]
+        return self.visit(v.args[2], bindings=bindings)[a:b]
+
+    def visit_if(self, v, bindings=None):
+        # Concretize all three args of the ITE expression.
+        i = self.visit(v.args[0], bindings=bindings)
+        t = self.visit(v.args[1], bindings=bindings)
+        e = self.visit(v.args[2], bindings=bindings)
+        return claripy.If(i, t, e)
+
+    def visit_reverse(self, v, bindings=None):
+        # Reverse is a simple unary; reverse the result from the arg.
+        res = claripy.Reverse(self.visit(v.args[0], bindings=bindings))
+        return res
+
+
+class PPrintVisitor(ClaripyVisitor):
+    def visit_int(self, v, **kwargs):
+        raise NotImplementedError("Why?")
+
+    def visit_bvv(self, v, **kwargs):
+        indent = kwargs.get("indent", 0)
+        out = kwargs.get("out", print)
+        out("  " * indent + str(v))
+
+    def visit_bvs(self, v, **kwargs):
+        indent = kwargs.get("indent", 0)
+        out = kwargs.get("out", print)
+        out("  " * indent + str(v))
+
+    def visit_add(self, v, **kwargs):
+        indent = kwargs.get("indent", 0)
+        out = kwargs.get("out", print)
+        out("  " * indent + "Add:")
+        for arg in v.args:
+            self.visit(arg, indent=indent + 1, out=out)
+
+    def visit_bwor(self, v, **kwargs):
+        indent = kwargs.get("indent", 0)
+        out = kwargs.get("out", print)
+        out("  " * indent + "BWOr:")
+        for arg in v.args:
+            self.visit(arg, indent=indent + 1, out=out)
+
+    def visit_concat(self, v, **kwargs):
+        indent = kwargs.get("indent", 0)
+        out = kwargs.get("out", print)
+        out("  " * indent + "Concat:")
+        for arg in v.args:
+            self.visit(arg, indent=indent + 1, out=out)
+
+    def visit_extract(self, v, **kwargs):
+        indent = kwargs.get("indent", 0)
+        out = kwargs.get("out", print)
+        out("  " * indent + f"Extract [{v.args[0]}:{v.args[1]}]:")
+        self.visit(v.args[2], indent=indent + 1, out=out)
+
+    def visit_if(self, v, **kwargs):
+        indent = kwargs.get("indent", 0)
+        out = kwargs.get("out", print)
+        out("  " * indent + "If:")
+        self.visit(v.args[0], indent=indent + 1, out=out)
+        out("  " * indent + "Then:")
+        self.visit(v.args[1], indent=indent + 1, out=out)
+        out("  " * indent + "Else:")
+        self.visit(v.args[2], indent=indent + 1, out=out)
+
+    def visit_reverse(self, v, **kwargs):
+        indent = kwargs.get("indent", 0)
+        out = kwargs.get("out", print)
+        out("  " * indent + "Reverse:")
+        self.visit(v.args[0], indent=indent + 1, out=out)
+
+    def visit_lshift(self, v, **kwargs):
+        indent = kwargs.get("indent", 0)
+        out = kwargs.get("out", print)
+        out("  " * indent + "LShift:")
+        self.visit(v.args[0], indent=indent + 1, out=out)
+        out("  " * indent + "By:")
+        self.visit(v.args[1], indent=indent + 1, out=out)
+
+    def visit_rshift(self, v, **kwargs):
+        indent = kwargs.get("indent", 0)
+        out = kwargs.get("out", print)
+        out("  " * indent + "RShift (Logical):")
+        self.visit(v.args[0], indent=indent + 1, out=out)
+        out("  " * indent + "By:")
+        self.visit(v.args[1], indent=indent + 1, out=out)
+
+    def visit_zeroext(self, v, **kwargs):
+        indent = kwargs.get("indent", 0)
+        out = kwargs.get("out", print)
+        out("  " * indent + f"ZeroExt[{v.args[0]}]:")
+        self.visit(v.args[1], indent=indent + 1, out=out)

smallworld/analyses/unstable/angr_nwbt.py

@@ -0,0 +1,106 @@
+import logging
+
+from ... import emulators, hinting, state
+from .. import analysis
+from .angr.nwbt import configure_nwbt_plugins, configure_nwbt_strategy
+from .angr.utils import print_state
+
+log = logging.getLogger(__name__)
+hinter = hinting.get_hinter(__name__)
+
+
+class AngrNWBTAnalysis(analysis.Analysis):
+    name = "angr-nwbt"
+    description = "Angr-based Not-Written-By-Trace detection and correction"
+    version = "0.0.1"
+
+    def __init__(self, *args, initfunc=None, max_steps=500, **kwargs):
+        self.steps_left = max_steps
+        self.initfunc = initfunc
+
+    def analysis_preinit(self, emu):
+        log.info("Registering plugins")
+        configure_nwbt_plugins(emu)
+
+    def analysis_init(self, emu):
+        configure_nwbt_strategy(emu)
+        if self.initfunc is not None:
+            self.initfunc(self, emu.entry)
+
+    def run(self, machine: state.Machine):
+        cpu = machine.get_cpu()
+        emu = emulators.AngrEmulator(
+            cpu.platform, preinit=self.analysis_preinit, init=self.analysis_init
+        )
+        machine.apply(emu)
+
+        # Extract typedef info from the CPU state,
+        # and bind it to the machine state
+        for item in cpu:
+            if isinstance(item, state.Register):
+                if item.get_type() is not None:
+                    log.debug(f"Applying type for {item}")
+                    emu.state.typedefs.bind_register(item.name, item.get_type())
+        for item in machine:
+            if isinstance(item, state.memory.code.Executable):
+                # TODO: Code is also memory, so it should have types...?
+                pass
+            elif isinstance(item, state.memory.Memory):
+                for offset, value in item.items():
+                    if value.get_type() is not None:
+                        addr = item.address + offset
+                        log.debug(f"Applying type for {hex(addr)}")
+                        emu.state.typedefs.bind_address(addr, value.get_type())
+            else:
+                if not isinstance(item, state.cpus.CPU) and item.type is not None:
+                    raise NotImplementedError(
+                        f"Applying typedef {item.get_type()} of type {type(item)} not implemented"
+                    )
+
+        while (self.steps_left is None or self.steps_left > 0) and not self.step(emu):
+            if self.steps_left is not None:
+                self.steps_left -= 1
+
+    def _report_status(self, emu):
+        for st in emu.mgr.active:
+            dis = "\r".join(map(str, st.block().disassembly.insns))
+            log.debug(f"Active state: {dis}")
+        for st in emu.mgr.unconstrained:
+            hint = hinting.OutputHint(
+                message="State left the program",
+                registers=st.registers.create_hint(),
+                memory=st.memory.create_hint(),
+            )
+            hinter.info(hint)
+        for st in emu.mgr.deadended:
+            hint = hinting.OutputHint(
+                message="State exited due to breakpoint",
+                registers=st.registers.create_hint(),
+                memory=st.memory.create_hint(),
+            )
+            hinter.info(hint)
+        for st in emu.mgr.unsat:
+            hint = hinting.OutputHint(
+                message="State cannot continue; constraints unsat",
+                registers=st.registers.create_hint(),
+                memory=st.memory.create_hint(),
+            )
+            hinter.info(hint)
+        for err in emu.mgr.errored:
+            print_state(log.info, err.state, "error")
+            log.error(
+                "\tError:",
+                exc_info=(type(err.error), err.error, err.error.__traceback__),
+            )
+        log.info(f"Summary: {emu.mgr}")
+
+    def step(self, emu):
+        # Report our current status
+        self._report_status(emu)
+        # In our case, return states are unconstrained.
+        emu.mgr.move(from_stash="deadended", to_stash="done")
+        emu.mgr.move(from_stash="unconstrained", to_stash="done")
+        # Drop unsat states once we've logged them.
+        emu.mgr.drop(stash="unsat")
+        self._report_status(emu)
+        return emu.step()