diff --git a/mythril/interfaces/cli.py b/mythril/interfaces/cli.py index aee8a637..6210060c 100644 --- a/mythril/interfaces/cli.py +++ b/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: diff --git a/mythril/leveldb/client.py b/mythril/leveldb/client.py index 83a80d24..5f1e4221 100644 --- a/mythril/leveldb/client.py +++ b/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() - md5.update(code.encode('UTF-8')) - contract_hash = md5.digest() - contract = ETHContract(code, name=contract_hash.hex()) + contract = ETHContract(code) 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 diff --git a/mythril/mythril.py b/mythril/mythril.py index ff0d5b8a..69d75356 100644 --- a/mythril/mythril.py +++ b/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.") diff --git a/setup.py b/setup.py index e1a623a2..a8e9b12e 100755 --- a/setup.py +++ b/setup.py @@ -166,9 +166,6 @@ in the `legendary "Mitch Brenner" blog post `__ 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