feat: add map_to_original_source flag

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
 

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"])

mythx_cli/util.py

@@ -250,43 +250,53 @@ 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 "
                 + str(target_path)
                 + ". 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