diff --git a/mythx_cli/fuzz/ide/brownie.py b/mythx_cli/fuzz/ide/brownie.py index 1dc60c7..07a8cca 100644 --- a/mythx_cli/fuzz/ide/brownie.py +++ b/mythx_cli/fuzz/ide/brownie.py @@ -7,17 +7,25 @@ from mythx_cli.fuzz.ide.generic import IDEArtifacts, JobBuilder from ...util import sol_files_by_directory +from ...util import files_by_directory LOGGER = logging.getLogger("mythx-cli") class BrownieArtifacts(IDEArtifacts): - def __init__(self, build_dir=None, targets=None): + def __init__(self, build_dir=None, targets=None, map_to_original_source=False): self._include = [] if targets: include = [] for target in targets: - include.extend(sol_files_by_directory(target)) + if not map_to_original_source: + LOGGER.debug(f"Mapping instrumented code") + include.extend(files_by_directory(target, ".sol")) + else: + # We replace .sol with .sol.original in case the target is a file and not a directory + target = target.replace(".sol", ".sol.original") + LOGGER.debug(f"Mapping original code, {target}") + include.extend(files_by_directory(target, ".sol.original")) self._include = include self._build_dir = build_dir or Path("./build/contracts") @@ -122,8 +130,8 @@ def _get_build_artifacts(build_dir) -> Dict: class BrownieJob: - def __init__(self, target: List[str], build_dir: Path): - artifacts = BrownieArtifacts(build_dir, targets=target) + def __init__(self, target: List[str], build_dir: Path, map_to_original_source: bool): + artifacts = BrownieArtifacts(build_dir, targets=target, map_to_original_source=map_to_original_source) self._jb = JobBuilder(artifacts) self.payload = None diff --git a/mythx_cli/fuzz/run.py b/mythx_cli/fuzz/run.py index e2b42c9..8ddf235 100644 --- a/mythx_cli/fuzz/run.py +++ b/mythx_cli/fuzz/run.py @@ -61,8 +61,16 @@ def determine_ide() -> IDE: default=None, required=False, ) +@click.option( + "-s", + "--map-to-original-source", + is_flag=True, + default=False, + help="Map the analyses results to the original source code, instead of the instrumented one. " + "This is meant to be used with Scribble.", +) @click.pass_obj -def fuzz_run(ctx, address, more_addresses, target, corpus_target): +def fuzz_run(ctx, address, more_addresses, corpus_target, map_to_original_source, target): # read YAML config params from ctx dict, e.g. ganache rpc url # Introduce a separate `fuzz` section in the YAML file @@ -87,6 +95,7 @@ def fuzz_run(ctx, address, more_addresses, target, corpus_target): "faas_url": "http://localhost:8080", "harvey_num_cores": 2, "campaign_name_prefix": "untitled", + "map_to_original_source": False, } config_options = analyze_config.keys() # Mandatory config parameters verification @@ -109,6 +118,12 @@ def fuzz_run(ctx, address, more_addresses, target, corpus_target): ) if not target: target = analyze_config["targets"] + if not map_to_original_source: + map_to_original_source = ( + analyze_config["map_to_original_source"] + if "map_to_original_source" in config_options + else default_config["map_to_original_source"] + ) # Optional config parameters # Here we parse the config parameters from the config file and use defaults for non available values contract_address = analyze_config["deployed_contract_address"] @@ -159,7 +174,7 @@ def fuzz_run(ctx, address, more_addresses, target, corpus_target): ide = determine_ide() if ide == IDE.BROWNIE: - artifacts = BrownieJob(target, analyze_config["build_directory"]) + artifacts = BrownieJob(target, analyze_config["build_directory"], map_to_original_source=map_to_original_source) artifacts.generate_payload() elif ide == IDE.HARDHAT: artifacts = HardhatJob(target, analyze_config["build_directory"]) diff --git a/mythx_cli/util.py b/mythx_cli/util.py index 66a30d7..c85bbc4 100644 --- a/mythx_cli/util.py +++ b/mythx_cli/util.py @@ -250,19 +250,28 @@ def write_or_print(ctx, data: str, mode="a+") -> None: LOGGER.debug(f"Writing data to {ctx['output']}") outfile.write(data + "\n") - def sol_files_by_directory(target_path: AnyStr) -> List: - """Gathers all the solidity files inside the target path + """Gathers all the .sol files inside the target path including sub-directories and returns them as a List. - Non solidity files are ignored. + Non .sol files are ignored. :param target_path: The directory to look for .sol files :return: """ - sol_files = [] - # We start by checking if the target_path is potentially a .sol file - if target_path.endswith(".sol"): - # If .sol file we check if the target exists or if it's a user input error + return files_by_directory(target_path, ".sol") + +def files_by_directory(target_path: AnyStr, extension: AnyStr) -> List: + """Gathers all the target extension files inside the target path + including sub-directories and returns them as a List. + Non target extension files are ignored. + + :param target_path: The directory to look for target extension files + :return: + """ + target_files = [] + # We start by checking if the target_path is potentially a target extension file + if target_path.endswith(extension): + # If target extension file we check if the target exists or if it's a user input error if not os.path.isfile(target_path): raise click.exceptions.UsageError( "Could not find " @@ -270,23 +279,24 @@ def sol_files_by_directory(target_path: AnyStr) -> List: + ". Did you pass the correct directory?" ) else: - # If it's a valid .sol file there is no need to search further and we just append it to our - # List to be returned - sol_files.append(target_path) + """ If it's a valid target extension file there is no need to search further and we just append it to our + list to be returned, removing the .original extension, leaving only the .sol """ + target_files.append(target_path.replace(".original","")) source_dir = os.walk(target_path) for sub_dir in source_dir: if len(sub_dir[2]) > 0: - # sub directory with .sol files + # sub directory with target extension files file_prefix = sub_dir[0] for file in sub_dir[2]: if "__scribble_" in file: LOGGER.debug(f"Skipped for being a scribble file {file}") continue - if not file.endswith(".sol"): + if not file.endswith(extension): LOGGER.debug(f"Skipped for not being a solidity file: {file}") continue file_name = file_prefix + "/" + file - LOGGER.debug(f"Found solidity file: {file_name}") - sol_files.append(file_name) - return sol_files + LOGGER.debug(f"Found target extension file: {file_name}") + # We remove the .original extension, added by Scribble + target_files.append(file_name.replace(".original","")) + return target_files