-
Notifications
You must be signed in to change notification settings - Fork 5
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #27 from crytic/template-generation
Template generation
- Loading branch information
Showing
32 changed files
with
1,845 additions
and
212 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,3 +1,9 @@ | ||
[submodule "tests/test_data/lib/forge-std"] | ||
path = tests/test_data/lib/forge-std | ||
url = https://github.com/foundry-rs/forge-std | ||
[submodule "tests/test_data/lib/properties"] | ||
path = tests/test_data/lib/properties | ||
url = https://github.com/crytic/properties | ||
[submodule "tests/test_data/lib/solmate"] | ||
path = tests/test_data/lib/solmate | ||
url = https://github.com/transmissions11/solmate |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,99 @@ | ||
"""The FoundryTest class that handles generation of unit tests from call sequences""" | ||
import os | ||
import sys | ||
import json | ||
from typing import Any | ||
import jinja2 | ||
|
||
from slither import Slither | ||
from slither.core.declarations.contract import Contract | ||
from fuzz_utils.utils.crytic_print import CryticPrint | ||
|
||
from fuzz_utils.generate.fuzzers.Medusa import Medusa | ||
from fuzz_utils.generate.fuzzers.Echidna import Echidna | ||
from fuzz_utils.templates.foundry_templates import templates | ||
|
||
|
||
class FoundryTest: # pylint: disable=too-many-instance-attributes | ||
""" | ||
Handles the generation of Foundry test files | ||
""" | ||
|
||
def __init__( | ||
self, | ||
config: dict, | ||
slither: Slither, | ||
fuzzer: Echidna | Medusa, | ||
) -> None: | ||
self.inheritance_path = config["inheritancePath"] | ||
self.target_name = config["targetContract"] | ||
self.corpus_path = config["corpusDir"] | ||
self.test_dir = config["testsDir"] | ||
self.all_sequences = config["allSequences"] | ||
self.slither = slither | ||
self.target = self.get_target_contract() | ||
self.fuzzer = fuzzer | ||
|
||
def get_target_contract(self) -> Contract: | ||
"""Gets the Slither Contract object for the specified contract file""" | ||
contracts = self.slither.get_contract_from_name(self.target_name) | ||
# Loop in case slither fetches multiple contracts for some reason (e.g., similar names?) | ||
for contract in contracts: | ||
if contract.name == self.target_name: | ||
return contract | ||
|
||
# TODO throw error if no contract found | ||
sys.exit(-1) | ||
|
||
def create_poc(self) -> str: | ||
"""Takes in a directory path to the echidna reproducers and generates a test file""" | ||
|
||
file_list: list[dict[str, Any]] = [] | ||
tests_list = [] | ||
dir_list = [] | ||
if self.all_sequences: | ||
dir_list = self.fuzzer.corpus_dirs | ||
else: | ||
dir_list = [self.fuzzer.reproducer_dir] | ||
|
||
# 1. Iterate over each directory and reproducer file (open it) | ||
for directory in dir_list: | ||
for entry in os.listdir(directory): | ||
full_path = os.path.join(directory, entry) | ||
|
||
if os.path.isfile(full_path): | ||
try: | ||
with open(full_path, "r", encoding="utf-8") as file: | ||
file_list.append({"path": full_path, "content": json.load(file)}) | ||
except Exception: # pylint: disable=broad-except | ||
print(f"Fail on {full_path}") | ||
|
||
# 2. Parse each reproducer file and add each test function to the functions list | ||
for idx, file_obj in enumerate(file_list): | ||
try: | ||
tests_list.append( | ||
self.fuzzer.parse_reproducer(file_obj["path"], file_obj["content"], idx) | ||
) | ||
except Exception: # pylint: disable=broad-except | ||
print(f"Parsing fail on {file_obj['content']}: index: {idx}") | ||
|
||
# 4. Generate the test file | ||
template = jinja2.Template(templates["CONTRACT"]) | ||
write_path = f"{self.test_dir}{self.target_name}" | ||
inheritance_path = f"{self.inheritance_path}{self.target_name}" | ||
|
||
# 5. Save the test file | ||
test_file_str = template.render( | ||
file_path=f"{inheritance_path}.sol", | ||
target_name=self.target_name, | ||
amount=0, | ||
tests=tests_list, | ||
fuzzer=self.fuzzer.name, | ||
) | ||
with open(f"{write_path}_{self.fuzzer.name}_Test.t.sol", "w", encoding="utf-8") as outfile: | ||
outfile.write(test_file_str) | ||
CryticPrint().print_success( | ||
f"Generated a test file in {write_path}_{self.fuzzer.name}_Test.t.sol" | ||
) | ||
|
||
return test_file_str |
File renamed without changes.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
File renamed without changes.
Empty file.
Oops, something went wrong.