diff --git a/README.md b/README.md index a23d073..d0064dd 100644 --- a/README.md +++ b/README.md @@ -69,6 +69,13 @@ The `template` command is used to generate a fuzzing harness. The harness can in - `-c`/`--contracts` `target_contracts: list`: The name of the target contract. - `-o`/`--output-dir` `output_directory: str`: Output directory name. By default it is `fuzzing` - `--config`: Path to the `fuzz-utils` config JSON file +- `--mode`: The strategy to use when generating the harnesses. Valid options: `simple`, `prank`, `actor` + +**Generation modes** +The tool support three harness generation strategies: +- `simple` - The fuzzing harness will be generated with all of the state-changing functions from the target contracts. All function calls are performed directly, with the harness contract as the `msg.sender`. +- `prank` - Similar to `simple` mode, with the difference that function calls are made from different users by using `hevm.prank()`. The users can be defined in the configuration file as `"actors": ["0xb4b3", "0xb0b", ...]` +- `actor` - `Actor` contracts will be generated and all harness function calls will be proxied through these contracts. The `Actor` contracts can be considered as users of the target contracts and the functions included in these actors can be filtered by modifier, external calls, or by `payable`. This allows for granular control over user capabilities. **Example** diff --git a/fuzz_utils/parsing/commands/template.py b/fuzz_utils/parsing/commands/template.py index 4c3940b..3de8a3d 100644 --- a/fuzz_utils/parsing/commands/template.py +++ b/fuzz_utils/parsing/commands/template.py @@ -6,6 +6,7 @@ from fuzz_utils.template.HarnessGenerator import HarnessGenerator from fuzz_utils.utils.crytic_print import CryticPrint from fuzz_utils.utils.remappings import find_remappings +from fuzz_utils.utils.error_handler import handle_exit def template_flags(parser: ArgumentParser) -> None: @@ -26,6 +27,11 @@ def template_flags(parser: ArgumentParser) -> None: help="Define the output directory where the result will be saved.", ) parser.add_argument("--config", dest="config", help="Define the location of the config file.") + parser.add_argument( + "--mode", + dest="mode", + help="Define the harness generation strategy you want to use. Valid options are `simple`, `prank`, `actor`", + ) def template_command(args: Namespace) -> None: @@ -47,6 +53,8 @@ def template_command(args: Namespace) -> None: config["compilationPath"] = args.compilation_path if args.name: config["name"] = args.name + if args.mode: + config["mode"] = args.mode.lower() config["outputDir"] = output_dir CryticPrint().print_information("Running Slither...") @@ -58,3 +66,21 @@ def template_command(args: Namespace) -> None: generator = HarnessGenerator(config, slither, remappings) generator.generate_templates() + + +def check_configuration(config: dict) -> None: + """Checks the configuration""" + mandatory_configuration_fields = ["mode", "targets", "compilationPath"] + for field in mandatory_configuration_fields: + check_configuration_field_exists_and_non_empty(config, field) + + if config["mode"].lower() not in ("simple", "prank", "actor"): + handle_exit( + f"The selected mode {config['mode']} is not a valid harness generation strategy." + ) + + +def check_configuration_field_exists_and_non_empty(config: dict, field: str) -> None: + """Checks that the configuration dictionary contains a non-empty field""" + if field not in config or len(config[field]) == 0: + handle_exit(f"The template configuration field {field} is not configured.") diff --git a/fuzz_utils/template/HarnessGenerator.py b/fuzz_utils/template/HarnessGenerator.py index 475d8b5..69e69c5 100644 --- a/fuzz_utils/template/HarnessGenerator.py +++ b/fuzz_utils/template/HarnessGenerator.py @@ -101,12 +101,33 @@ def __init__( slither: Slither, remappings: dict, ) -> None: - if "actors" in config: - config["actors"] = check_and_populate_actor_fields(config["actors"], config["targets"]) - else: - CryticPrint().print_warning("Using default values for the Actor.") - config["actors"] = self.config["actors"] - config["actors"][0]["targets"] = config["targets"] + self.mode = config["mode"] + match config["mode"]: + case "actor": + if "actors" in config: + config["actors"] = check_and_populate_actor_fields( + config["actors"], config["targets"] + ) + else: + CryticPrint().print_warning("Using default values for the Actor.") + config["actors"] = self.config["actors"] + config["actors"][0]["targets"] = config["targets"] + case "simple": + config["actors"] = [] + case "prank": + if "actors" in config: + if not isinstance(config["actors"], list[str]) or len(config["actors"]) > 0: # type: ignore[misc] + CryticPrint().print_warning( + "Actors not defined. Using default 0xb4b3 and 0xb0b." + ) + config["actors"] = ["0xb4b3", "0xb0b"] + else: + CryticPrint().print_warning( + "Actors not defined. Using default 0xb4b3 and 0xb0b." + ) + config["actors"] = ["0xb4b3", "0xb0b"] + case _: + handle_exit(f"Invalid template mode {config['mode']} was provided.") for key, value in config.items(): if key in self.config and value: @@ -128,22 +149,33 @@ def generate_templates(self) -> None: CryticPrint().print_information( f"Generating the fuzzing Harness for contracts: {self.config['targets']}" ) - # Check if directories exists, if not, create them check_and_create_dirs(self.output_dir, ["utils", "actors", "harnesses", "attacks"]) - # Generate the Actors - actors: list[Actor] = self._generate_actors() - CryticPrint().print_success(" Actors generated!") + # Generate the Attacks attacks: list[Actor] = self._generate_attacks() CryticPrint().print_success(" Attacks generated!") + actors: list = [] + + # Generate actors and harnesses, depending on strategy + match self.mode: + case "actor": + # Generate the Actors + actors = self._generate_actors() + CryticPrint().print_success(" Actors generated!") + case "prank": + actors = self.config["actors"] + case _: + pass + # Generate the harness self._generate_harness(actors, attacks) + CryticPrint().print_success(" Harness generated!") CryticPrint().print_success(f"Files saved to {self.config['outputDir']}") - # pylint: disable=too-many-locals - def _generate_harness(self, actors: list[Actor], attacks: list[Actor]) -> None: + # pylint: disable=too-many-locals,too-many-statements,too-many-branches + def _generate_harness(self, actors: list, attacks: list[Actor]) -> None: CryticPrint().print_information(f"Generating {self.config['name']} Harness") # Generate inheritance and variables @@ -155,9 +187,12 @@ def _generate_harness(self, actors: list[Actor], attacks: list[Actor]) -> None: variables.append(f"{contract.name} {contract.name.lower()};") # Generate actor variables and imports - for actor in actors: - variables.append(f"Actor{actor.name}[] {actor.name}_actors;") - imports.append(f'import "{actor.path}";') + if self.mode == "actor": + for actor in actors: + variables.append(f"Actor{actor.name}[] {actor.name}_actors;") + imports.append(f'import "{actor.path}";') + elif self.mode == "prank": + variables.append("address[] pranked_actors;") # Generate attack variables and imports for attack in attacks: @@ -176,23 +211,33 @@ def _generate_harness(self, actors: list[Actor], attacks: list[Actor]) -> None: inputs_str: str = ", ".join(inputs) constructor += f" {contract.name.lower()} = new {contract.name}({inputs_str});\n" - for actor in actors: - constructor += " for(uint256 i; i < 3; i++) {\n" - constructor_arguments = "" - if actor.contract and hasattr(actor.contract.constructor, "parameters"): - constructor_arguments = ", ".join( - [f"address({x.name.strip('_')})" for x in actor.contract.constructor.parameters] + if self.mode == "actor": + for actor in actors: + constructor += " for(uint256 i; i < 3; i++) {\n" + constructor_arguments = "" + if actor.contract and hasattr(actor.contract.constructor, "parameters"): + constructor_arguments = ", ".join( + [ + f"address({x.name.strip('_')})" + for x in actor.contract.constructor.parameters + ] + ) + constructor += ( + f" {actor.name}_actors.push(new Actor{actor.name}({constructor_arguments}));\n" + + " }\n" ) - constructor += ( - f" {actor.name}_actors.push(new Actor{actor.name}({constructor_arguments}));\n" - + " }\n" - ) + elif self.mode == "prank": + for actor in actors: + constructor += f" pranked_actors.push(address({actor}));\n" for attack in attacks: constructor_arguments = "" if attack.contract and hasattr(attack.contract.constructor, "parameters"): constructor_arguments = ", ".join( - [f"address({x.name.strip('_')})" for x in actor.contract.constructor.parameters] + [ + f"address({x.name.strip('_')})" + for x in attack.contract.constructor.parameters + ] ) constructor += f" {attack.name.lower()}Attack = new {attack.name}({constructor_arguments});\n" constructor += " }\n" @@ -201,12 +246,26 @@ def _generate_harness(self, actors: list[Actor], attacks: list[Actor]) -> None: # Generate Functions functions: list[str] = [] - for actor in actors: - function_body = f" {actor.contract.name} selectedActor = {actor.name}_actors[clampBetween(actorIndex, 0, {actor.name}_actors.length - 1)];\n" - temp_list = self._generate_functions( - actor.contract, None, ["uint256 actorIndex"], function_body, "selectedActor" - ) - functions.extend(temp_list) + if self.mode == "actor": + for actor in actors: + function_body = f" {actor.contract.name} selectedActor = {actor.name}_actors[clampBetween(actorIndex, 0, {actor.name}_actors.length - 1)];\n" + temp_list = self._generate_functions( + actor.contract, None, ["uint256 actorIndex"], function_body, "selectedActor" + ) + functions.extend(temp_list) + else: + for contract in self.targets: + function_body = "" + appended_params = [] + if self.mode == "prank": + function_body = " address selectedActor = pranked_actors[clampBetween(actorIndex, 0, pranked_actors.length - 1)];\n" + function_body += " hevm.prank(selectedActor);\n" + appended_params.append("uint256 actorIndex") + + temp_list = self._generate_functions( + contract, None, appended_params, function_body, contract.name.lower() + ) + functions.extend(temp_list) for attack in attacks: temp_list = self._generate_functions( diff --git a/fuzz_utils/templates/default_config.py b/fuzz_utils/templates/default_config.py index d41416e..db41bb4 100644 --- a/fuzz_utils/templates/default_config.py +++ b/fuzz_utils/templates/default_config.py @@ -12,6 +12,7 @@ }, "template": { "name": "DefaultHarness", + "mode": "simple", "targets": [], "outputDir": "./test/fuzzing", "compilationPath": ".", diff --git a/fuzz_utils/templates/harness_templates.py b/fuzz_utils/templates/harness_templates.py index d828fee..9f83cb7 100644 --- a/fuzz_utils/templates/harness_templates.py +++ b/fuzz_utils/templates/harness_templates.py @@ -29,6 +29,7 @@ /// -------------------------------------------------------------------- import "{{remappings["properties"]}}util/PropertiesHelper.sol"; +import "{{remappings["properties"]}}util/Hevm.sol"; {% for import in target.imports -%} {{import}} {% endfor %} diff --git a/tests/test_harness.py b/tests/test_harness.py index c75dc4e..9b8876e 100644 --- a/tests/test_harness.py +++ b/tests/test_harness.py @@ -13,6 +13,7 @@ TEST_DATA_DIR = Path(__file__).resolve().parent / "test_data" default_config = { "name": "DefaultHarness", + "mode": "actor", "compilationPath": ".", "targets": [], "outputDir": "./test/fuzzing",