diff --git a/fuzz_utils/fuzzers/Echidna.py b/fuzz_utils/fuzzers/Echidna.py index 6a7b27b..135e99e 100644 --- a/fuzz_utils/fuzzers/Echidna.py +++ b/fuzz_utils/fuzzers/Echidna.py @@ -23,12 +23,15 @@ class Echidna: Handles the generation of Foundry test files from Echidna reproducers """ - def __init__(self, target_name: str, corpus_path: str, slither: Slither) -> None: + def __init__( + self, target_name: str, corpus_path: str, slither: Slither, named_inputs: bool + ) -> None: self.name = "Echidna" self.target_name = target_name self.slither = slither self.target = self.get_target_contract() self.reproducer_dir = f"{corpus_path}/reproducers" + self.named_inputs = named_inputs def get_target_contract(self) -> Contract: """Finds and returns Slither Contract""" @@ -58,6 +61,7 @@ def parse_reproducer(self, calls: Any, index: int) -> str: template = jinja2.Template(templates["TEST"]) return template.render(function_name=function_name, call_list=call_list) + # pylint: disable=too-many-locals,too-many-branches def _parse_call_object(self, call_dict: dict[Any, Any]) -> tuple[str, str]: """ Takes a single call dictionary, parses it, and returns the series of function calls as a string, along with @@ -103,6 +107,15 @@ def _parse_call_object(self, call_dict: dict[Any, Any]) -> tuple[str, str]: variable_definition, call_definition = self._decode_function_params( function_parameters, False, slither_entry_point ) + parameters_str: str = "" + if isinstance(slither_entry_point.parameters, list): + if self.named_inputs and len(slither_entry_point.parameters) > 0: + for idx, input_param in enumerate(slither_entry_point.parameters): + call_definition[idx] = input_param.name + ": " + call_definition[idx] + parameters_str = "{" + ", ".join(call_definition) + "}" + print(parameters_str) + else: + parameters_str = ", ".join(call_definition) # 3. Generate a call string and return it template = jinja2.Template(templates["CALL"]) @@ -112,7 +125,7 @@ def _parse_call_object(self, call_dict: dict[Any, Any]) -> tuple[str, str]: block_delay=block_delay, caller=caller, value=value, - function_parameters=", ".join(call_definition), + function_parameters=parameters_str, function_name=function_name, contract_name=self.target_name, ) diff --git a/fuzz_utils/fuzzers/Medusa.py b/fuzz_utils/fuzzers/Medusa.py index 2745a45..57e17ea 100644 --- a/fuzz_utils/fuzzers/Medusa.py +++ b/fuzz_utils/fuzzers/Medusa.py @@ -23,13 +23,16 @@ class Medusa: Handles the generation of Foundry test files from Medusa reproducers """ - def __init__(self, target_name: str, corpus_path: str, slither: Slither) -> None: + def __init__( + self, target_name: str, corpus_path: str, slither: Slither, named_inputs: bool + ) -> None: self.name = "Medusa" self.target_name = target_name self.corpus_path = corpus_path self.slither = slither self.target = self.get_target_contract() self.reproducer_dir = f"{corpus_path}/test_results" + self.named_inputs = named_inputs def get_target_contract(self) -> Contract: """Finds and returns Slither Contract""" @@ -61,6 +64,7 @@ def parse_reproducer(self, calls: Any, index: int) -> str: # 3. Using the call list to generate a test string # 4. Return the test string + # pylint: disable=too-many-locals,too-many-branches def _parse_call_object(self, call_dict: dict) -> tuple[str, str]: """ Takes a single call dictionary, parses it, and returns the series of function calls as a string, along with @@ -96,6 +100,16 @@ def _parse_call_object(self, call_dict: dict) -> tuple[str, str]: function_parameters, False, slither_entry_point ) + parameters_str: str = "" + if isinstance(slither_entry_point.parameters, list): + if self.named_inputs and len(slither_entry_point.parameters) > 0: + for idx, input_param in enumerate(slither_entry_point.parameters): + call_definition[idx] = input_param.name + ": " + call_definition[idx] + parameters_str = "{" + ", ".join(call_definition) + "}" + print(parameters_str) + else: + parameters_str = ", ".join(call_definition) + # 3. Generate a call string and return it template = jinja2.Template(templates["CALL"]) call_str = template.render( @@ -104,7 +118,7 @@ def _parse_call_object(self, call_dict: dict) -> tuple[str, str]: block_delay=block_delay, caller=caller, value=value, - function_parameters=", ".join(call_definition), + function_parameters=parameters_str, function_name=function_name, contract_name=self.target_name, ) diff --git a/fuzz_utils/main.py b/fuzz_utils/main.py index 2b93933..961035e 100644 --- a/fuzz_utils/main.py +++ b/fuzz_utils/main.py @@ -128,6 +128,13 @@ def main() -> None: # type: ignore[func-returns-value] version=require("fuzz-utils")[0].version, action="version", ) + parser.add_argument( + "--named-inputs", + dest="named_inputs", + help="Include function input names when making calls.", + default=False, + action="store_true", + ) args = parser.parse_args() @@ -146,9 +153,9 @@ def main() -> None: # type: ignore[func-returns-value] match args.selected_fuzzer.lower(): case "echidna": - fuzzer = Echidna(target_contract, corpus_dir, slither) + fuzzer = Echidna(target_contract, corpus_dir, slither, args.named_inputs) case "medusa": - fuzzer = Medusa(target_contract, corpus_dir, slither) + fuzzer = Medusa(target_contract, corpus_dir, slither, args.named_inputs) case _: handle_exit( f"\n* The requested fuzzer {args.selected_fuzzer} is not supported. Supported fuzzers: echidna, medusa." diff --git a/tests/conftest.py b/tests/conftest.py index e855b09..d47f554 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -14,8 +14,8 @@ class TestGenerator: def __init__(self, target: str, target_path: str, corpus_dir: str): slither = Slither(target_path) - echidna = Echidna(target, f"echidna-corpora/{corpus_dir}", slither) - medusa = Medusa(target, f"medusa-corpora/{corpus_dir}", slither) + echidna = Echidna(target, f"echidna-corpora/{corpus_dir}", slither, False) + medusa = Medusa(target, f"medusa-corpora/{corpus_dir}", slither, False) self.echidna_generator = FoundryTest( "../src/", target, f"echidna-corpora/{corpus_dir}", "./test/", slither, echidna )