From 65322ee7f7f76937a539aa3611eee19ed62ccfd6 Mon Sep 17 00:00:00 2001
From: Joran Honig <joran.honig@gmail.com>
Date: Tue, 10 Jul 2018 14:26:06 +0200
Subject: [PATCH] Fix circular dependency

---
 mythril/laser/ethereum/call.py |  3 +--
 mythril/laser/ethereum/svm.py  | 40 +++++++++++++++++++++++++++++++++-
 2 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/mythril/laser/ethereum/call.py b/mythril/laser/ethereum/call.py
index aee2f30c..d94b71bc 100644
--- a/mythril/laser/ethereum/call.py
+++ b/mythril/laser/ethereum/call.py
@@ -1,8 +1,7 @@
 import logging
 from z3 import BitVec, simplify
 import mythril.laser.ethereum.util as util
-from mythril.laser.ethereum.state import Account
-from mythril.laser.ethereum.svm import CalldataType
+from mythril.laser.ethereum.state import Account, CalldataType
 import re
 
 """
diff --git a/mythril/laser/ethereum/svm.py b/mythril/laser/ethereum/svm.py
index 1ceb3740..d35e0d32 100644
--- a/mythril/laser/ethereum/svm.py
+++ b/mythril/laser/ethereum/svm.py
@@ -35,6 +35,9 @@ class LaserEVM:
         self.strategy = DepthFirstSearchStrategy(self.work_list, max_depth)
         self.max_depth = max_depth
 
+        self.pre_hooks = {}
+        self.post_hooks = {}
+
         logging.info("LASER EVM initialized with dynamic loader: " + str(dynamic_loader))
 
     def sym_exec(self, main_address):
@@ -81,7 +84,12 @@ class LaserEVM:
     def execute_state(self, global_state):
         instructions = global_state.environment.code.instruction_list
         op_code = instructions[global_state.mstate.pc]['opcode']
-        return Instruction(op_code, self.dynamic_loader).evaluate(global_state), op_code
+
+        self._execute_pre_hook(op_code, global_state)
+        new_global_states = Instruction(op_code, self.dynamic_loader).evaluate(global_state)
+        self._execute_post_hook(op_code, new_global_states)
+
+        return new_global_states, op_code
 
     def manage_cfg(self, opcode, new_states):
         if opcode == "JUMP":
@@ -130,3 +138,33 @@ class LaserEVM:
             logging.info("- Entering function " + environment.active_account.contract_name + ":" + new_node.function_name)
 
         new_node.function_name = environment.active_function_name
+
+    def _execute_pre_hook(self, op_code, global_state):
+        if op_code not in self.pre_hooks.keys():
+            return
+        for hook in self.pre_hooks[op_code]:
+            hook(global_state)
+
+    def _execute_post_hook(self, op_code, global_states):
+        if op_code not in self.post_hooks.keys():
+            return
+
+        for hook in self.post_hooks[op_code]:
+            for global_state in global_states:
+                hook(global_state)
+
+    def hook(self, op_code):
+        def hook_decorator(function):
+            if op_code not in self.pre_hooks.keys():
+                self.pre_hooks[op_code] = []
+            self.pre_hooks[op_code].append(function)
+            return function
+        return hook_decorator
+
+    def post_hook(self, op_code):
+        def hook_decorator(function):
+            if op_code not in self.post_hooks.keys():
+                self.post_hooks[op_code] = []
+            self.post_hooks[op_code].append(function)
+            return function
+        return hook_decorator