diff --git a/mythril/mythril/mythril_analyzer.py b/mythril/mythril/mythril_analyzer.py index 95cf9667..d1f9b807 100644 --- a/mythril/mythril/mythril_analyzer.py +++ b/mythril/mythril/mythril_analyzer.py @@ -7,6 +7,9 @@ import logging import traceback +from typing import Optional, List + +from . import MythrilDisassembler from mythril.support.source_support import Source from mythril.support.loader import DynLoader from mythril.analysis.symbolic import SymExecWrapper @@ -14,17 +17,23 @@ from mythril.analysis.callgraph import generate_graph from mythril.analysis.traceexplore import get_serializable_statespace from mythril.analysis.security import fire_lasers, retrieve_callback_issues from mythril.analysis.report import Report +from mythril.ethereum.evmcontract import EVMContract log = logging.getLogger(__name__) -class MythrilAnalyzer(object): +class MythrilAnalyzer: """ The Mythril Analyzer class Responsible for the analysis of the smart contracts """ - def __init__(self, disassembler, requires_dynld=False, onchain_storage_access=True): + def __init__( + self, + disassembler: MythrilDisassembler, + requires_dynld: bool = False, + onchain_storage_access: bool = True, + ): """ :param disassembler: The MythrilDisassembler class @@ -32,21 +41,21 @@ class MythrilAnalyzer(object): :param onchain_storage_access: Whether onchain access should be done or not """ self.eth = disassembler.eth - self.contracts = disassembler.contracts or [] + self.contracts = disassembler.contracts or [] # type: List[EVMContract] self.enable_online_lookup = disassembler.enable_online_lookup self.dynld = requires_dynld self.onchain_storage_access = onchain_storage_access def dump_statespace( self, - strategy, - contract, - address=None, - max_depth=None, - execution_timeout=None, - create_timeout=None, - enable_iprof=False, - ): + strategy: str, + contract: List[EVMContract], + address: Optional[str] = None, + max_depth: Optional[int] = None, + execution_timeout: Optional[int] = None, + create_timeout: Optional[int] = None, + enable_iprof: bool = False, + ) -> str: """ Returns serializable statespace of the contract :param strategy: The search strategy to go through the CFG @@ -77,16 +86,16 @@ class MythrilAnalyzer(object): def graph_html( self, - strategy, - contract, - address, - max_depth=None, - enable_physics=False, - phrackify=False, - execution_timeout=None, - create_timeout=None, - enable_iprof=False, - ): + strategy: str, + contract: List[EVMContract], + address: str, + max_depth: Optional[int] = None, + enable_physics: bool = False, + phrackify: bool = False, + execution_timeout: Optional[int] = None, + create_timeout: Optional[int] = None, + enable_iprof: bool = False, + ) -> str: """ :param strategy: The search strategy to go through the CFG @@ -118,17 +127,17 @@ class MythrilAnalyzer(object): def fire_lasers( self, - strategy, - contracts=None, - address=None, - modules=None, - verbose_report=False, - max_depth=None, - execution_timeout=None, - create_timeout=None, - transaction_count=None, - enable_iprof=False, - ): + strategy: str, + contracts: Optional[List[EVMContract]] = None, + address: Optional[str] = None, + modules: Optional[List[str]] = None, + verbose_report: bool = False, + max_depth: Optional[int] = None, + execution_timeout: Optional[int] = None, + create_timeout: Optional[int] = None, + transaction_count: Optional[int] = None, + enable_iprof: bool = False, + ) -> Report: """ :param strategy: The search strategy to go through the CFG :param contracts: The Contracts list on which the analysis should be done diff --git a/mythril/mythril/mythril_config.py b/mythril/mythril/mythril_config.py index 4e9217bd..55c04149 100644 --- a/mythril/mythril/mythril_config.py +++ b/mythril/mythril/mythril_config.py @@ -7,6 +7,7 @@ import re from pathlib import Path from shutil import copyfile from configparser import ConfigParser +from typing import Optional from mythril.exceptions import CriticalError from mythril.ethereum.interface.rpc.client import EthJsonRpc @@ -15,16 +16,16 @@ from mythril.ethereum.interface.leveldb.client import EthLevelDB log = logging.getLogger(__name__) -class MythrilConfig(object): +class MythrilConfig: def __init__(self): self.mythril_dir = self._init_mythril_dir() self.config_path = os.path.join(self.mythril_dir, "config.ini") self.leveldb_dir = self._init_config() - self.eth = None - self.eth_db = None + self.eth = None # type: Optional[EthJsonRpc] + self.eth_db = None # type: Optional[EthLevelDB] @staticmethod - def _init_mythril_dir(): + def _init_mythril_dir() -> str: """ Initializes the mythril dir and config.ini file :return: The mythril dir's path @@ -48,7 +49,7 @@ class MythrilConfig(object): return mythril_dir - def _init_config(self): + def _init_config(self) -> str: """If no config file exists, create it and add default options. Default LevelDB path is specified based on OS @@ -83,7 +84,7 @@ class MythrilConfig(object): return os.path.expanduser(leveldb_dir) @staticmethod - def _get_fallback_dir(): + def _get_fallback_dir() -> str: """ Returns the LevelDB path :return: The LevelDB path @@ -103,7 +104,7 @@ class MythrilConfig(object): return os.path.join(leveldb_fallback_dir, "geth", "chaindata") @staticmethod - def _add_default_options(config): + def _add_default_options(config: ConfigParser) -> None: """ Adds defaults option to config.ini :param config: The config file object @@ -112,7 +113,7 @@ class MythrilConfig(object): config.add_section("defaults") @staticmethod - def _add_leveldb_option(config, leveldb_fallback_dir): + def _add_leveldb_option(config: ConfigParser, leveldb_fallback_dir: str) -> None: """ Sets a default leveldb path in .mythril/config.ini file :param config: The config file object @@ -129,7 +130,7 @@ class MythrilConfig(object): config.set("defaults", "leveldb_dir", leveldb_fallback_dir) @staticmethod - def _add_dynamic_loading_option(config): + def _add_dynamic_loading_option(config: ConfigParser) -> None: """ Sets the dynamic loading config option in .mythril/config.ini file :param config: The config file object @@ -146,17 +147,17 @@ class MythrilConfig(object): ) config.set("defaults", "dynamic_loading", "infura") - def set_api_leveldb(self, leveldb_path): + def set_api_leveldb(self, leveldb_path: str) -> None: """ """ self.eth_db = EthLevelDB(leveldb_path) - def set_api_rpc_infura(self): + def set_api_rpc_infura(self) -> None: """Set the RPC mode to INFURA on Mainnet.""" log.info("Using INFURA Main Net for RPC queries") self.eth = EthJsonRpc("mainnet.infura.io", 443, True) - def set_api_rpc(self, rpc=None, rpctls=False): + def set_api_rpc(self, rpc: str = None, rpctls: bool = False) -> None: """ Sets the RPC mode to either ganache or infura """ @@ -181,12 +182,12 @@ class MythrilConfig(object): else: raise CriticalError("Invalid RPC settings, check help for details.") - def set_api_rpc_localhost(self): + def set_api_rpc_localhost(self) -> None: """Set the RPC mode to a local instance.""" log.info("Using default RPC settings: http://localhost:8545") self.eth = EthJsonRpc("localhost", 8545) - def set_api_from_config_path(self): + def set_api_from_config_path(self) -> None: """Set the RPC mode based on a given config file.""" config = ConfigParser(allow_no_value=False) config.optionxform = str diff --git a/mythril/mythril/mythril_disassembler.py b/mythril/mythril/mythril_disassembler.py index bdda55e5..3e455115 100644 --- a/mythril/mythril/mythril_disassembler.py +++ b/mythril/mythril/mythril_disassembler.py @@ -5,8 +5,9 @@ import os from ethereum import utils from solc.exceptions import SolcError -from typing import List, Optional +from typing import List, Tuple, Optional from mythril.ethereum import util +from mythril.ethereum.interface.rpc.client import EthJsonRpc from mythril.exceptions import CriticalError, CompilerError, NoContractFoundError from mythril.support import signatures from mythril.support.truffle import analyze_truffle_project @@ -19,17 +20,21 @@ log = logging.getLogger(__name__) class MythrilDisassembler: def __init__( - self, eth, solc_version=None, solc_args=None, enable_online_lookup=False - ): + self, + eth: Optional[EthJsonRpc], + solc_version: str = None, + solc_args: str = None, + enable_online_lookup: bool = False, + ) -> None: self.solc_binary = self._init_solc_binary(solc_version) self.solc_args = solc_args self.eth = eth self.enable_online_lookup = enable_online_lookup self.sigs = signatures.SignatureDB(enable_online_lookup=enable_online_lookup) - self.contracts = [] + self.contracts = [] # type: List[EVMContract] @staticmethod - def _init_solc_binary(version): + def _init_solc_binary(version: str) -> str: """ Only proper versions are supported. No nightlies, commits etc (such as available in remix). :param version: Version of the solc binary required @@ -68,7 +73,9 @@ class MythrilDisassembler: return solc_binary - def load_from_bytecode(self, code, bin_runtime=False, address=None): + def load_from_bytecode( + self, code: str, bin_runtime: bool = False, address: Optional[str] = None + ) -> Tuple[str, EVMContract]: """ Returns the address and the contract class for the given bytecode :param code: Bytecode @@ -96,7 +103,7 @@ class MythrilDisassembler: ) return address, self.contracts[-1] # return address and contract object - def load_from_address(self, address): + def load_from_address(self, address: str) -> Tuple[str, EVMContract]: """ Returns the contract given it's on chain address :param address: The on chain address of a contract @@ -130,7 +137,9 @@ class MythrilDisassembler: ) return address, self.contracts[-1] # return address and contract object - def load_from_solidity(self, solidity_files): + def load_from_solidity( + self, solidity_files: List[str] + ) -> Tuple[str, List[SolidityContract]]: """ :param solidity_files: List of solidity_files @@ -180,18 +189,18 @@ class MythrilDisassembler: return address, contracts - def analyze_truffle_project(self, *args, **kwargs): + def analyze_truffle_project(self, *args, **kwargs) -> None: """ :param args: :param kwargs: :return: """ - return analyze_truffle_project( + analyze_truffle_project( self.sigs, *args, **kwargs ) # just passthru by passing signatures for now @staticmethod - def hash_for_function_signature(func): + def hash_for_function_signature(func: str) -> str: """ Return function name's corresponding signature hash :param func: function name diff --git a/mythril/mythril/mythril_leveldb.py b/mythril/mythril/mythril_leveldb.py index 5ab2e29b..8220a614 100644 --- a/mythril/mythril/mythril_leveldb.py +++ b/mythril/mythril/mythril_leveldb.py @@ -2,7 +2,7 @@ import re from mythril.exceptions import CriticalError -class MythrilLevelDB(object): +class MythrilLevelDB: """ Class which does search operations on leveldb There are two DBs diff --git a/tests/mythril/mythril_analyzer_test.py b/tests/mythril/mythril_analyzer_test.py index 730c3046..90423dcc 100644 --- a/tests/mythril/mythril_analyzer_test.py +++ b/tests/mythril/mythril_analyzer_test.py @@ -1,6 +1,5 @@ -import pytest from pathlib import Path -from mythril.mythril import * +from mythril.mythril import MythrilDisassembler, MythrilAnalyzer def test_fire_lasers(): diff --git a/tests/mythril/mythril_config_test.py b/tests/mythril/mythril_config_test.py new file mode 100644 index 00000000..5a78a1e1 --- /dev/null +++ b/tests/mythril/mythril_config_test.py @@ -0,0 +1,2 @@ +from mythril.mythril import MythrilConfig +