From cd9eb3dffe13a36ff03488eb2ee6ec5834d33989 Mon Sep 17 00:00:00 2001 From: tuturu-tech Date: Fri, 29 Mar 2024 16:58:22 +0100 Subject: [PATCH 1/5] added mode command, updated generator with first iteration of modes --- fuzz_utils/parsing/commands/template.py | 26 +++++ fuzz_utils/template/HarnessGenerator.py | 143 +++++++++++++++++++++--- 2 files changed, 155 insertions(+), 14 deletions(-) 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..eb71fae 100644 --- a/fuzz_utils/template/HarnessGenerator.py +++ b/fuzz_utils/template/HarnessGenerator.py @@ -3,6 +3,7 @@ import os import copy from dataclasses import dataclass +from eth_utils import to_checksum_address from slither import Slither from slither.core.declarations.contract import Contract @@ -101,12 +102,21 @@ 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": + pass + case "prank": + pass + 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 +138,127 @@ 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!") - # Generate the harness - self._generate_harness(actors, attacks) + + # Generate actors and harnesses, depending on strategy + match self.mode: + case "actor": + # Generate the Actors + actors: list[Actor] = self._generate_actors() + CryticPrint().print_success(" Actors generated!") + + # Generate the harness + self._generate_harness_with_actors(actors, attacks) + case "simple": + # Generate the harness + self._generate_harness_simple_or_prank([], attacks) + case "prank": + # Generate the harness + self._generate_harness_simple_or_prank(self.config["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: + def _generate_harness_simple_or_prank(self, actors: list, attacks: list[Actor]) -> None: + CryticPrint().print_information(f"Generating {self.config['name']} Harness") + + # Generate inheritance and variables + imports: list[str] = [] + variables: list[str] = [] + + for contract in self.targets: + imports.append(f'import "{contract.source_mapping.filename.relative}";') + variables.append(f"{contract.name} {contract.name.lower()};") + + # Generate actor variables and imports + if self.mode == "prank": + variables.append("address[] pranked_actors;") + + # Generate attack variables and imports + for attack in attacks: + variables.append(f"Attack{attack.name} {attack.name.lower()}Attack;") + imports.append(f'import "{attack.path}";') + + # Generate constructor with contract, actor, and attack deployment + constructor = "constructor() {\n" + for contract in self.targets: + inputs: list[str] = [] + if contract.constructor: + constructor_parameters = contract.constructor.parameters + for param in constructor_parameters: + constructor += f" {param.type} {param.name};\n" + inputs.append(param.name) + inputs_str: str = ", ".join(inputs) + constructor += f" {contract.name.lower()} = new {contract.name}({inputs_str});\n" + + if self.mode == "prank": + for actor in actors: + constructor += " for(uint256 i; i < 3; i++) {\n" + constructor += ( + f" pranked_actors.push(address({actor}));\n" + + " }\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 attack.contract.constructor.parameters] + ) + constructor += f" {attack.name.lower()}Attack = new {attack.name}({constructor_arguments});\n" + constructor += " }\n" + # Generate dependencies + dependencies: str = "PropertiesAsserts" + + # Generate Functions + functions: list[str] = [] + 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( + attack.contract, None, [], None, f"{attack.name.lower()}Attack" + ) + functions.extend(temp_list) + + # Generate harness class + harness = Harness( + name=self.config["name"], + constructor=constructor, + dependencies=dependencies, + content="", + path="", + targets=self.targets, + actors=actors, + imports=imports, + variables=variables, + functions=functions, + ) + + content, path = self._render_template( + templates["HARNESS"], "harnesses", self.config["name"], harness + ) + harness.set_content(content) + harness.set_path(path) + + # pylint: disable=too-many-locals + def _generate_harness_with_actors(self, actors: list[Actor], attacks: list[Actor]) -> None: CryticPrint().print_information(f"Generating {self.config['name']} Harness") # Generate inheritance and variables @@ -192,7 +307,7 @@ def _generate_harness(self, actors: list[Actor], attacks: list[Actor]) -> None: 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" From a377b5be832f95f63ccb028104b508cd17a20a62 Mon Sep 17 00:00:00 2001 From: tuturu-tech Date: Fri, 29 Mar 2024 16:59:38 +0100 Subject: [PATCH 2/5] formatting --- fuzz_utils/template/HarnessGenerator.py | 22 ++++++++++++++-------- 1 file changed, 14 insertions(+), 8 deletions(-) diff --git a/fuzz_utils/template/HarnessGenerator.py b/fuzz_utils/template/HarnessGenerator.py index eb71fae..c1cafd4 100644 --- a/fuzz_utils/template/HarnessGenerator.py +++ b/fuzz_utils/template/HarnessGenerator.py @@ -3,7 +3,6 @@ import os import copy from dataclasses import dataclass -from eth_utils import to_checksum_address from slither import Slither from slither.core.declarations.contract import Contract @@ -106,7 +105,9 @@ def __init__( match config["mode"]: case "actor": if "actors" in config: - config["actors"] = check_and_populate_actor_fields(config["actors"], config["targets"]) + 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"] @@ -164,7 +165,7 @@ def generate_templates(self) -> None: CryticPrint().print_success(" Harness generated!") CryticPrint().print_success(f"Files saved to {self.config['outputDir']}") - # pylint: disable=too-many-locals + # pylint: disable=too-many-locals,too-many-branches def _generate_harness_simple_or_prank(self, actors: list, attacks: list[Actor]) -> None: CryticPrint().print_information(f"Generating {self.config['name']} Harness") @@ -175,7 +176,7 @@ def _generate_harness_simple_or_prank(self, actors: list, attacks: list[Actor]) for contract in self.targets: imports.append(f'import "{contract.source_mapping.filename.relative}";') variables.append(f"{contract.name} {contract.name.lower()};") - + # Generate actor variables and imports if self.mode == "prank": variables.append("address[] pranked_actors;") @@ -201,15 +202,17 @@ def _generate_harness_simple_or_prank(self, actors: list, attacks: list[Actor]) for actor in actors: constructor += " for(uint256 i; i < 3; i++) {\n" constructor += ( - f" pranked_actors.push(address({actor}));\n" - + " }\n" + f" pranked_actors.push(address({actor}));\n" + " }\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 attack.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" @@ -307,7 +310,10 @@ def _generate_harness_with_actors(self, actors: list[Actor], attacks: list[Actor constructor_arguments = "" if attack.contract and hasattr(attack.contract.constructor, "parameters"): constructor_arguments = ", ".join( - [f"address({x.name.strip('_')})" for x in attack.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" From 3c31c7c012a8e9b5cb8e2c6fb4e734b0b3935da2 Mon Sep 17 00:00:00 2001 From: tuturu-tech Date: Fri, 29 Mar 2024 17:28:39 +0100 Subject: [PATCH 3/5] modify default config, add simple, prank, and actor mode to harness generation --- fuzz_utils/template/HarnessGenerator.py | 180 +++++++--------------- fuzz_utils/templates/default_config.py | 1 + fuzz_utils/templates/harness_templates.py | 1 + 3 files changed, 61 insertions(+), 121 deletions(-) diff --git a/fuzz_utils/template/HarnessGenerator.py b/fuzz_utils/template/HarnessGenerator.py index c1cafd4..69e69c5 100644 --- a/fuzz_utils/template/HarnessGenerator.py +++ b/fuzz_utils/template/HarnessGenerator.py @@ -113,9 +113,19 @@ def __init__( config["actors"] = self.config["actors"] config["actors"][0]["targets"] = config["targets"] case "simple": - pass + config["actors"] = [] case "prank": - pass + 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.") @@ -145,28 +155,27 @@ def generate_templates(self) -> None: # 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: list[Actor] = self._generate_actors() + actors = self._generate_actors() CryticPrint().print_success(" Actors generated!") - - # Generate the harness - self._generate_harness_with_actors(actors, attacks) - case "simple": - # Generate the harness - self._generate_harness_simple_or_prank([], attacks) case "prank": - # Generate the harness - self._generate_harness_simple_or_prank(self.config["actors"], attacks) + 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,too-many-branches - def _generate_harness_simple_or_prank(self, actors: list, 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 @@ -178,7 +187,11 @@ def _generate_harness_simple_or_prank(self, actors: list, attacks: list[Actor]) variables.append(f"{contract.name} {contract.name.lower()};") # Generate actor variables and imports - if self.mode == "prank": + 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 @@ -198,12 +211,24 @@ def _generate_harness_simple_or_prank(self, actors: list, attacks: list[Actor]) inputs_str: str = ", ".join(inputs) constructor += f" {contract.name.lower()} = new {contract.name}({inputs_str});\n" - if self.mode == "prank": + 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" pranked_actors.push(address({actor}));\n" + " }\n" + 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 = "" @@ -221,113 +246,26 @@ def _generate_harness_simple_or_prank(self, actors: list, attacks: list[Actor]) # Generate Functions functions: list[str] = [] - 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( - attack.contract, None, [], None, f"{attack.name.lower()}Attack" - ) - functions.extend(temp_list) - - # Generate harness class - harness = Harness( - name=self.config["name"], - constructor=constructor, - dependencies=dependencies, - content="", - path="", - targets=self.targets, - actors=actors, - imports=imports, - variables=variables, - functions=functions, - ) - - content, path = self._render_template( - templates["HARNESS"], "harnesses", self.config["name"], harness - ) - harness.set_content(content) - harness.set_path(path) - - # pylint: disable=too-many-locals - def _generate_harness_with_actors(self, actors: list[Actor], attacks: list[Actor]) -> None: - CryticPrint().print_information(f"Generating {self.config['name']} Harness") - - # Generate inheritance and variables - imports: list[str] = [] - variables: list[str] = [] - - for contract in self.targets: - imports.append(f'import "{contract.source_mapping.filename.relative}";') - 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}";') - - # Generate attack variables and imports - for attack in attacks: - variables.append(f"Attack{attack.name} {attack.name.lower()}Attack;") - imports.append(f'import "{attack.path}";') - - # Generate constructor with contract, actor, and attack deployment - constructor = "constructor() {\n" - for contract in self.targets: - inputs: list[str] = [] - if contract.constructor: - constructor_parameters = contract.constructor.parameters - for param in constructor_parameters: - constructor += f" {param.type} {param.name};\n" - inputs.append(param.name) - 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: + 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" ) - constructor += ( - f" {actor.name}_actors.push(new Actor{actor.name}({constructor_arguments}));\n" - + " }\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 attack.contract.constructor.parameters - ] + 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() ) - constructor += f" {attack.name.lower()}Attack = new {attack.name}({constructor_arguments});\n" - constructor += " }\n" - # Generate dependencies - dependencies: str = "PropertiesAsserts" - - # 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) + 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 %} From bc4c364387b09ce4259af4e17b0bfe5d3538faf2 Mon Sep 17 00:00:00 2001 From: tuturu-tech Date: Fri, 29 Mar 2024 17:38:44 +0100 Subject: [PATCH 4/5] add mode explanation to readme --- README.md | 7 +++++++ 1 file changed, 7 insertions(+) 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** From 18c0fa90f5bf5eb2c252ac86fe6b2a62d71455dd Mon Sep 17 00:00:00 2001 From: tuturu-tech Date: Fri, 29 Mar 2024 17:48:16 +0100 Subject: [PATCH 5/5] fix harness tests --- tests/test_harness.py | 1 + 1 file changed, 1 insertion(+) 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",