From 2f80f78b82fc88aeb51446ec79873428c85de824 Mon Sep 17 00:00:00 2001
From: Nikhil <nikhilparasaram@gmail.com>
Date: Mon, 14 Oct 2019 15:26:49 +0100
Subject: [PATCH] Use concrete hashes too

---
 mythril/analysis/solver.py                    | 23 +++++++++++++++----
 .../laser/ethereum/keccak_function_manager.py | 23 ++++++++++++++-----
 tests/laser/evm_testsuite/evm_test.py         | 12 +---------
 3 files changed, 37 insertions(+), 21 deletions(-)

diff --git a/mythril/analysis/solver.py b/mythril/analysis/solver.py
index 74d838f1..69688768 100644
--- a/mythril/analysis/solver.py
+++ b/mythril/analysis/solver.py
@@ -145,10 +145,25 @@ def _replace_with_actual_sha(
             find_input = symbol_factory.BitVecVal(
                 int(tx["input"][10 + i : 74 + i], 16), 256
             )
-            _, inverse = keccak_function_manager.get_function(160)
-            input_ = symbol_factory.BitVecVal(
-                model.eval(inverse(find_input).raw).as_long(), 160
-            )
+            input_ = None
+            for size in (160, 256, 512):
+                _, inverse = keccak_function_manager.get_function(size)
+                try:
+                    input_ = symbol_factory.BitVecVal(
+                        model.eval(inverse(find_input).raw).as_long(), size
+                    )
+                except AttributeError:
+                    continue
+                hex_input = hex(input_.value)[2:]
+                found = False
+                for new_tx in concrete_transactions:
+                    if hex_input in new_tx["input"]:
+                        found = True
+                        break
+                if found:
+                    break
+            if input_ is None:
+                continue
             keccak = keccak_function_manager.find_keccak(input_)
             tx["input"] = tx["input"].replace(
                 tx["input"][10 + i : 74 + i], hex(keccak.value)[2:]
diff --git a/mythril/laser/ethereum/keccak_function_manager.py b/mythril/laser/ethereum/keccak_function_manager.py
index e3ed72d3..3bc3f26d 100644
--- a/mythril/laser/ethereum/keccak_function_manager.py
+++ b/mythril/laser/ethereum/keccak_function_manager.py
@@ -8,8 +8,9 @@ from mythril.laser.smt import (
     And,
     ULT,
     Bool,
+    Or,
 )
-from typing import Dict, Tuple
+from typing import Dict, Tuple, List
 
 TOTAL_PARTS = 10 ** 40
 PART = (2 ** 256 - 1) // TOTAL_PARTS
@@ -22,7 +23,7 @@ class KeccakFunctionManager:
         self.store_function = {}  # type: Dict[int, Tuple[Function, Function]]
         self.interval_hook_for_size = {}  # type: Dict[int, int]
         self._index_counter = TOTAL_PARTS - 34534
-        self.quick_inverse = {}  # type: Dict[BitVec, BitVec]  # This is for VMTests
+        self.concrete_hash_vals = {}  # type: Dict[int, List[BitVec]]
 
     def find_keccak(self, data: BitVec) -> BitVec:
         """
@@ -50,6 +51,7 @@ class KeccakFunctionManager:
         except KeyError:
             func = Function("keccak256_{}".format(length), length, 256)
             inverse = Function("keccak256_{}-1".format(length), 256, length)
+            self.concrete_hash_vals[length] = []
             self.store_function[length] = (func, inverse)
         return func, inverse
 
@@ -60,10 +62,17 @@ class KeccakFunctionManager:
         :return: Tuple of keccak and the condition it should satisfy
         """
         length = data.size()
-        condition = self._create_condition(func_input=data)
-        func, _ = self.get_function(length)
-        self.quick_inverse[func(data)] = data
-        return func(data), condition
+        func, inverse = self.get_function(length)
+
+        if data.symbolic:
+            condition = self._create_condition(func_input=data)
+            output = func(data)
+        else:
+            concrete_val = self.find_keccak(data)
+            condition = And(func(data) == concrete_val, inverse(func(data)) == data)
+            self.concrete_hash_vals[length].append(concrete_val)
+            output = concrete_val
+        return output, condition
 
     def _create_condition(self, func_input: BitVec) -> Bool:
         """
@@ -89,6 +98,8 @@ class KeccakFunctionManager:
             ULT(func(func_input), symbol_factory.BitVecVal(upper_bound, 256)),
             URem(func(func_input), symbol_factory.BitVecVal(64, 256)) == 0,
         )
+        for val in self.concrete_hash_vals[length]:
+            cond = Or(cond, func(func_input) == val)
         return cond
 
 
diff --git a/tests/laser/evm_testsuite/evm_test.py b/tests/laser/evm_testsuite/evm_test.py
index 8afe3d14..9fbb9dfa 100644
--- a/tests/laser/evm_testsuite/evm_test.py
+++ b/tests/laser/evm_testsuite/evm_test.py
@@ -1,11 +1,9 @@
 from mythril.laser.ethereum.svm import LaserEVM
 from mythril.laser.ethereum.state.account import Account
-from mythril.laser.ethereum.keccak_function_manager import keccak_function_manager
 from mythril.laser.ethereum.state.world_state import WorldState
 from mythril.disassembler.disassembly import Disassembly
 from mythril.laser.ethereum.transaction.concolic import execute_message_call
 from mythril.laser.smt import Expression, BitVec, symbol_factory
-from mythril.analysis.solver import get_model
 from datetime import datetime
 
 import binascii
@@ -178,15 +176,7 @@ def test_vmtest(
                 expected = int(value, 16)
                 actual = account.storage[symbol_factory.BitVecVal(int(index, 16), 256)]
                 if isinstance(actual, Expression):
-                    if (
-                        actual.symbolic
-                        and actual in keccak_function_manager.quick_inverse
-                    ):
-                        actual = keccak_function_manager.find_keccak(
-                            keccak_function_manager.quick_inverse[actual]
-                        )
-                    else:
-                        actual = actual.value
+                    actual = actual.value
                     actual = 1 if actual is True else 0 if actual is False else actual
                 else:
                     if type(actual) == bytes: