Merge pull request #40 from crytic/template-modes

Add `template` modes

Author: tuturu-tech

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**
 

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

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(

fuzz_utils/templates/default_config.py

@@ -12,6 +12,7 @@
     },
     "template": {
         "name": "DefaultHarness",
+        "mode": "simple",
         "targets": [],
         "outputDir": "./test/fuzzing",
         "compilationPath": ".",

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 %}

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",