From 28efcf92b42e8d3f7f91163f4b02c94aedcf64e9 Mon Sep 17 00:00:00 2001
From: Bernhard Mueller <bernhard.mueller11@gmail.com>
Date: Fri, 29 Sep 2017 10:20:53 +0700
Subject: [PATCH] Trace now returns a string

---
 .gitignore         |  1 +
 contractstorage.py |  6 ++----
 ether/evm.py       | 15 ++++++++++++---
 mythril            | 11 ++++++++---
 4 files changed, 23 insertions(+), 10 deletions(-)

diff --git a/.gitignore b/.gitignore
index d080669a..ebc80678 100644
--- a/.gitignore
+++ b/.gitignore
@@ -8,3 +8,4 @@ mythril.egg-info
 build
 dist
 contracts.json
+hunt*
diff --git a/contractstorage.py b/contractstorage.py
index 8ac72b43..0e9a90ee 100644
--- a/contractstorage.py
+++ b/contractstorage.py
@@ -76,7 +76,7 @@ class ContractStorage:
         return ret[0]['code']
 
 
-    def search(self, expression):
+    def search(self, expression, callback_func):
 
         all_contracts = self.db.all()
 
@@ -87,8 +87,6 @@ class ContractStorage:
 
             sign_hash = utils.sha3(m)[:4].hex()
 
-            print(sign_hash)
-
             expression = expression.replace(m, sign_hash)
 
         for c in all_contracts:
@@ -98,5 +96,5 @@ class ContractStorage:
                 contract = ETHContract(c['code'], instance['balance'])
 
                 if (contract.matches_expression(expression)):
-                    print("Found contract:" + instance['address'])
+                    callback_func(instance['address'])
 
diff --git a/ether/evm.py b/ether/evm.py
index cae8069c..c6234385 100644
--- a/ether/evm.py
+++ b/ether/evm.py
@@ -2,7 +2,7 @@ from ethereum import vm, messages, transactions
 from ethereum.state import State
 from ethereum.slogging import get_logger
 from logging import StreamHandler
-import sys
+from io import StringIO
 import codecs
 from .util import safe_decode
 
@@ -11,7 +11,8 @@ def trace(code, address = "", calldata = ""):
 
 	logHandlers = ['eth.vm.op', 'eth.vm.op.stack', 'eth.vm.op.memory', 'eth.vm.op.storage']
 
-	streamHandler = StreamHandler(sys.stdout)
+	output = StringIO()
+	streamHandler = StreamHandler(output)
 
 	for handler in logHandlers:
 		log_vm_op = get_logger(handler)
@@ -28,4 +29,12 @@ def trace(code, address = "", calldata = ""):
 
 	message = vm.Message(addr_from, addr_to, 0, 21000, data, code_address=addr_to)
 
-	res, gas, dat = vm.vm_execute(ext, message, code)
\ No newline at end of file
+	res, gas, dat = vm.vm_execute(ext, message, code)
+
+	streamHandler.flush()
+
+	# print(output.getvalue())
+
+	ret = output.getvalue()
+
+	return ret
\ No newline at end of file
diff --git a/mythril b/mythril
index 8a97903c..b7014b8c 100755
--- a/mythril
+++ b/mythril
@@ -11,6 +11,9 @@ import argparse
 from rpc.client import EthJsonRpc
 
 
+def searchCallback(address):
+    print(address)
+
 def exitWithError(message):
     print(message)
     sys.exit()
@@ -86,15 +89,17 @@ elif (args.trace):
         exitWithError("Disassembler: Provide the input bytecode via -c BYTECODE or --id ID")
 
     if (args.data):
-        evm.trace(util.safe_decode(encoded_bytecode), args.address, args.data)
+        output = evm.trace(util.safe_decode(encoded_bytecode), args.address, args.data)
 
     else:
-        evm.trace(util.safe_decode(encoded_bytecode), args.address)
+        output = evm.trace(util.safe_decode(encoded_bytecode), args.address)
+
+    print(output)
 
 
 elif (args.search):
 
-    storage.search(args.search)
+    storage.search(args.search, searchCallback)
 
 elif (args.init_db):
         storage.initialize(args.rpchost, args.rpcport)