Merge pull request #342 from b-mueller/master

Refactor search

mythril/interfaces/cli.py

@@ -51,7 +51,6 @@ def main():
     database = parser.add_argument_group('local contracts database')
     database.add_argument('-s', '--search', help='search the contract database', metavar='EXPRESSION')
     database.add_argument('--leveldb-dir', help='specify leveldb directory for search or direct access operations', metavar='LEVELDB_PATH')
-    database.add_argument('--search-all', action='store_true', help='search all contracts instead of active (non-zero balance) only')
 
     utilities = parser.add_argument_group('utilities')
     utilities.add_argument('--hash', help='calculate function signature hash', metavar='SIGNATURE')
@@ -135,7 +134,7 @@ def main():
 
         if args.search:
             # Database search ops
-            mythril.search_db(args.search, args.search_all)
+            mythril.search_db(args.search)
             sys.exit()
 
         if args.truffle:

mythril/leveldb/client.py

@@ -2,6 +2,7 @@ import plyvel
 import binascii
 import rlp
 import hashlib
+import logging
 from ethereum import utils
 from ethereum.block import BlockHeader, Block
 from mythril.leveldb.state import State, Account
@@ -40,27 +41,32 @@ class EthLevelDB(object):
         self.headBlockHeader = None
         self.headState = None
 
-    def get_contracts(self, search_all):
+    def get_contracts(self):
         '''
-        iterate through contracts with non-zero balance by default or all if search_all is set
+        iterate through all contracts
         '''
         for account in self._get_head_state().get_all_accounts():
-            if account.code is not None and (search_all or account.balance != 0):
+            if account.code is not None:
                 code = _encode_hex(account.code)
-                md5 = hashlib.md5()
+                contract = ETHContract(code)
-                md5.update(code.encode('UTF-8'))
-                contract_hash = md5.digest()
-                contract = ETHContract(code, name=contract_hash.hex())
                 yield contract, _encode_hex(account.address), account.balance
 
-    def search(self, expression, search_all, callback_func):
+    def search(self, expression, callback_func):
         '''
         searches through non-zero balance contracts
         '''
-        for contract, address, balance in self.get_contracts(search_all):
+        cnt = 0
+
+        for contract, address, balance in self.get_contracts():
+
             if contract.matches_expression(expression):
                 callback_func(contract.name, contract, [address], [balance])
 
+            cnt += 1
+
+            if not cnt % 1000:
+                logging.info("Searched %d contracts" % cnt)
+
     def eth_getBlockHeaderByNumber(self, number):
         '''
         gets block header by block number

mythril/mythril.py

@@ -226,15 +226,14 @@ class Mythril(object):
         self.eth = EthJsonRpc('localhost', 8545)
         logging.info("Using default RPC settings: http://localhost:8545")
 
-    def search_db(self, search, search_all):
+    def search_db(self, search):
 
         def search_callback(code_hash, code, addresses, balances):
-            print("Matched contract with code hash " + code_hash)
             for i in range(0, len(addresses)):
-                print("Address: " + addresses[i] + ", balance: " + str(balances[i]))
+                print("Address Hash: " + addresses[i] + ", balance: " + str(balances[i]))
 
         try:
-            self.ethDb.search(search, search_all, search_callback)
+            self.ethDb.search(search, search_callback)
 
         except SyntaxError:
             raise CriticalError("Syntax error in search expression.")

setup.py

@@ -166,9 +166,6 @@ in the `legendary "Mitch Brenner" blog post
 <https://medium.com/@rtaylor30/how-i-snatched-your-153-037-eth-after-a-bad-tinder-date-d1d84422a50b>`__
 in [STRIKEOUT:seconds] minutes instead of days.
 
-The default behavior is to search contracts with a non-zero balance.
-You can disable this behavior with the ``--search-all`` flag.
-
 You may also use geth database directly for fetching contracts instead of
 using IPC/RPC APIs by specifying ``--leveldb`` flag. This is useful
 because search will return hashed addresses which will not be accepted by