Streamline State Serialization Interface

manticore/core/manticore.py

@@ -311,6 +311,23 @@ def __init__(self, initial_state, workspace_url=None, policy="random", **kwargs)
     def __str__(self):
         return f"<{str(type(self))[8:-2]}| Alive States: {self.count_ready_states()}; Running States: {self.count_busy_states()} Terminated States: {self.count_terminated_states()} Killed States: {self.count_killed_states()} Started: {self._running.value} Killed: {self._killed.value}>"
 
+    @classmethod
+    def from_saved_state(cls, filename: str, *args, **kwargs):
+        """
+        Creates a Manticore object starting from a serialized state on the disk.
+
+        :param filename: File to load the state from
+        :param args: Arguments forwarded to the Manticore object
+        :param kwargs: Keyword args forwarded to the Manticore object
+        :return: An instance of a subclass of ManticoreBase with the given initial state
+        """
+        from ..utils.helpers import PickleSerializer
+
+        with open(filename, "rb") as fd:
+            deserialized = PickleSerializer().deserialize(fd)
+
+        return cls(deserialized, *args, **kwargs)
+
     def _fork(self, state, expression, policy="ALL", setstate=None):
         """
         Fork state on expression concretizations.

manticore/core/state.py

@@ -1,7 +1,7 @@
 import copy
 import logging
 
-from .smtlib import solver, Bool, issymbolic
+from .smtlib import solver, Bool, issymbolic, BitVecConstant
 from ..utils.event import Eventful
 
 logger = logging.getLogger(__name__)
@@ -54,6 +54,28 @@ def __init__(self, message, expression, setstate=None, policy=None, **kwargs):
         super().__init__(**kwargs)
 
 
+class SerializeState(Concretize):
+    """ Allows the user to save a copy of the current state somewhere on the
+        disk so that analysis can later be resumed from this point.
+    """
+
+    def setstate(self, state, _value):
+        from ..utils.helpers import PickleSerializer
+
+        with open(self.filename, "wb") as statef:
+            PickleSerializer().serialize(state, statef)
+
+    def __init__(self, filename, **kwargs):
+        super().__init__(
+            f"Saving state to {filename}",
+            BitVecConstant(32, 0),
+            setstate=self.setstate,
+            policy="ONE",
+            **kwargs,
+        )
+        self.filename = filename
+
+
 class ForkState(Concretize):
     """ Specialized concretization class for Bool expressions.
         It tries True and False as concrete solutions. /

tests/native/test_resume.py

@@ -0,0 +1,47 @@
+import unittest
+from manticore.native import Manticore
+from manticore.core.state import SerializeState, TerminateState
+from pathlib import Path
+
+
+ms_file = str(
+    Path(__file__).parent.parent.parent.joinpath("examples", "linux", "binaries", "multiple-styles")
+)
+
+
+class TestResume(unittest.TestCase):
+    def test_resume(self):
+        m = Manticore(ms_file, stdin_size=17)
+
+        # First instruction of `main`
+        @m.hook(0x4009AE)
+        def serialize(state):
+            with m.locked_context() as context:
+                if context.get("kill", False):
+                    raise TerminateState("Abandoning...")
+                context["kill"] = True
+            raise SerializeState("/tmp/ms_checkpoint.pkl")
+
+        m.run()
+        self.assertEqual(m.count_terminated_states(), 1)
+        for state in m.terminated_states:
+            self.assertEqual(state.cpu.PC, 0x4009AE)
+
+        m = Manticore.from_saved_state("/tmp/ms_checkpoint.pkl")
+        self.assertEqual(m.count_ready_states(), 1)
+        for st in m.ready_states:
+            self.assertEqual(state.cpu.PC, 0x4009AE)
+        m.run()
+
+        self.assertEqual(m.count_terminated_states(), 18)
+        self.assertTrue(
+            any("exit status: 0" in str(st._terminated_by) for st in m.terminated_states)
+        )
+        m.finalize()
+        for st in m.terminated_states:
+            if "exit status: 0" in str(st._terminated_by):
+                self.assertEqual(st.solve_one(st.input_symbols[0]), b"coldlikeminisodas")
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/wasm/test_state_saving.py

@@ -0,0 +1,76 @@
+import unittest
+from manticore.wasm import ManticoreWASM
+from manticore.wasm.types import I32
+from manticore.core.plugin import Plugin
+from manticore.core.state import SerializeState, TerminateState
+from pathlib import Path
+
+
+class CallCounterPlugin(Plugin):
+    def did_execute_instruction_callback(self, state, instruction):
+        with self.locked_context("counter", dict) as ctx:
+            val = ctx.setdefault(instruction.mnemonic, 0)
+            ctx[instruction.mnemonic] = val + 1
+
+
+class SerializerPlugin(Plugin):
+    killed = False
+
+    def did_execute_instruction_callback(self, state, instruction):
+        if self.killed:
+            raise TerminateState("Abandoning")
+        with self.locked_context("counter", dict) as ctx:
+            if instruction.mnemonic == "loop" and ctx.get(instruction.mnemonic, 0) == 24:
+                self.killed = True
+                raise SerializeState("/tmp/collatz_checkpoint.pkl")
+            val = ctx.setdefault(instruction.mnemonic, 0)
+            ctx[instruction.mnemonic] = val + 1
+
+
+collatz_file = str(
+    Path(__file__).parent.parent.parent.joinpath("examples", "wasm", "collatz", "collatz.wasm")
+)
+
+
+class TestResume(unittest.TestCase):
+    def test_resume(self):
+        m = ManticoreWASM(collatz_file)
+        plugin = CallCounterPlugin()
+        m.register_plugin(plugin)
+        m.collatz(lambda s: [I32(1337)])
+        m.run()
+
+        counts_canonical = plugin.context.get("counter")
+
+        m = ManticoreWASM(collatz_file)
+        plugin = SerializerPlugin()
+        m.register_plugin(plugin)
+        m.collatz(lambda s: [I32(1337)])
+        m.run()
+
+        counts_save = plugin.context.get("counter")
+
+        m = ManticoreWASM.from_saved_state("/tmp/collatz_checkpoint.pkl")
+        plugin = CallCounterPlugin()
+        m.register_plugin(plugin)
+        m.run()
+
+        counts_resume = plugin.context.get("counter")
+
+        for k in counts_canonical:
+            with self.subTest(k):
+                self.assertEqual(
+                    counts_save.get(k, 0) + counts_resume.get(k, 0),
+                    counts_canonical[k],
+                    f"Mismatched {k} count",
+                )
+
+        results = []
+        for idx, val_list in enumerate(m.collect_returns()):
+            results.append(val_list[0][0])
+
+        self.assertEqual(sorted(results), [44])
+
+
+if __name__ == "__main__":
+    unittest.main()