From ebb516036132ad80c39e9852b93d2824a16a697a Mon Sep 17 00:00:00 2001 From: Alexander Remie Date: Thu, 17 Mar 2022 17:21:52 +0100 Subject: [PATCH 1/4] add etherscan function to ignore/remove contracts which are not part of a deployed contract --- crytic_compile/platform/etherscan.py | 92 ++++++++++++++++++++++++++++ 1 file changed, 92 insertions(+) diff --git a/crytic_compile/platform/etherscan.py b/crytic_compile/platform/etherscan.py index e83b206a..ee7a60ce 100644 --- a/crytic_compile/platform/etherscan.py +++ b/crytic_compile/platform/etherscan.py @@ -7,6 +7,8 @@ import os import re import urllib.request +import glob +import shutil from json.decoder import JSONDecodeError from pathlib import Path from typing import TYPE_CHECKING, Dict, List, Union, Tuple, Optional @@ -330,6 +332,8 @@ def compile(self, crytic_compile: "CryticCompile", **kwargs: str) -> None: solc_standard_json.standalone_compile(filenames, compilation_unit, working_dir=working_dir) + _remove_unused_contracts(compilation_unit, export_dir) + @staticmethod def is_supported(target: str, **kwargs: str) -> bool: """Check if the target is a etherscan project @@ -390,3 +394,91 @@ def _relative_to_short(relative: Path) -> Path: Path: Translated path """ return relative + + +def _remove_unused_contracts(compilation_unit: CompilationUnit, export_dir: str) -> None: + """ + Removes unused contracts from the compilation unit and file system + + Args: + compilation_unit (CompilationUnit): compilation unit to populate + export_dir (str): export dir + + Returns: + + """ + if len(list(compilation_unit.asts.keys())) == 1: + # there is only 1 file + return + + # for etherscan this will be the value the etherscan api returns in 'ContractName' + root_contract_name = compilation_unit.unique_id + + # find the root file path by it's name + # and also get the base path used by all paths (the keys under 'asts') + root_file_path = None + base_path = None + for file_path, file_ast in compilation_unit.asts.items(): + if root_file_path is not None: # already found target contract + break + for node in file_ast['nodes']: + if node['nodeType'] == 'ContractDefinition' and node['name'] == root_contract_name: + root_file_path = file_path + base_path = file_path.replace(file_ast['absolutePath'], '') + break + + if root_file_path is None: + # we could not find a contract with that name in any of the files + return + + # Starting with the root contract, fetch all dependencies (and their dependencies, etc.) + files_to_include = [] + files_to_check = [root_file_path] + while any(files_to_check): + target_file_path = files_to_check.pop() + for node in compilation_unit.asts[target_file_path]['nodes']: + if node['nodeType'] == 'ImportDirective': + import_path = os.path.join(base_path, node['absolutePath']) + if import_path not in files_to_check and import_path not in files_to_include: + files_to_check.append(import_path) + files_to_include.append(target_file_path) + + if len(list(compilation_unit.asts.keys())) == len(files_to_include): + # all of the files need to be included + return + + # Remove all of the unused files from the compilation unit + included_contractnames = set() + for target_file_path in files_to_include: + for node in compilation_unit.asts[target_file_path]['nodes']: + if node['nodeType'] == 'ContractDefinition': + included_contractnames.add(node['name']) + + for contractname in list(compilation_unit.contracts_names): + if contractname not in included_contractnames: + compilation_unit.contracts_names.remove(contractname) + del compilation_unit.abis[contractname] + del compilation_unit.natspec[contractname] + del compilation_unit.bytecodes_init[contractname] + del compilation_unit.bytecodes_runtime[contractname] + del compilation_unit.srcmaps_init[contractname] + del compilation_unit.srcmaps_runtime[contractname] + + for contract_filename in list(compilation_unit.filename_to_contracts.keys()): + if contract_filename.absolute not in files_to_include: + del compilation_unit.filename_to_contracts[contract_filename] + + for fileobj in list(compilation_unit.crytic_compile.filenames): + if fileobj.absolute not in files_to_include: + compilation_unit.crytic_compile.filenames.remove(fileobj) + compilation_unit.filenames.remove(fileobj) + del compilation_unit.asts[fileobj.absolute] + + # remove all unused files + for filename in glob.iglob(os.path.join(export_dir, '**/**'), recursive=True): + if os.path.isfile(filename) and filename not in files_to_include: + os.remove(filename) + # remove all folders which are now empty + for filename in glob.iglob(os.path.join(export_dir, '**/**'), recursive=True): + if os.path.isdir(filename) and len(os.listdir(filename)) == 0: + shutil.rmtree(filename) \ No newline at end of file From 67517a221b374e4812eb7be1810388f47e6c3192 Mon Sep 17 00:00:00 2001 From: Alexander Remie Date: Thu, 17 Mar 2022 17:32:00 +0100 Subject: [PATCH 2/4] linting fixes --- crytic_compile/platform/etherscan.py | 29 ++++++++++++++-------------- 1 file changed, 15 insertions(+), 14 deletions(-) diff --git a/crytic_compile/platform/etherscan.py b/crytic_compile/platform/etherscan.py index ee7a60ce..71408894 100644 --- a/crytic_compile/platform/etherscan.py +++ b/crytic_compile/platform/etherscan.py @@ -396,6 +396,7 @@ def _relative_to_short(relative: Path) -> Path: return relative +# pylint: disable=too-many-locals,too-many-branches def _remove_unused_contracts(compilation_unit: CompilationUnit, export_dir: str) -> None: """ Removes unused contracts from the compilation unit and file system @@ -417,14 +418,14 @@ def _remove_unused_contracts(compilation_unit: CompilationUnit, export_dir: str) # find the root file path by it's name # and also get the base path used by all paths (the keys under 'asts') root_file_path = None - base_path = None + base_path = "" for file_path, file_ast in compilation_unit.asts.items(): - if root_file_path is not None: # already found target contract + if root_file_path is not None: # already found target contract break - for node in file_ast['nodes']: - if node['nodeType'] == 'ContractDefinition' and node['name'] == root_contract_name: + for node in file_ast["nodes"]: + if node["nodeType"] == "ContractDefinition" and node["name"] == root_contract_name: root_file_path = file_path - base_path = file_path.replace(file_ast['absolutePath'], '') + base_path = file_path.replace(file_ast["absolutePath"], "") break if root_file_path is None: @@ -436,9 +437,9 @@ def _remove_unused_contracts(compilation_unit: CompilationUnit, export_dir: str) files_to_check = [root_file_path] while any(files_to_check): target_file_path = files_to_check.pop() - for node in compilation_unit.asts[target_file_path]['nodes']: - if node['nodeType'] == 'ImportDirective': - import_path = os.path.join(base_path, node['absolutePath']) + for node in compilation_unit.asts[target_file_path]["nodes"]: + if node["nodeType"] == "ImportDirective": + import_path = os.path.join(base_path, node["absolutePath"]) if import_path not in files_to_check and import_path not in files_to_include: files_to_check.append(import_path) files_to_include.append(target_file_path) @@ -450,9 +451,9 @@ def _remove_unused_contracts(compilation_unit: CompilationUnit, export_dir: str) # Remove all of the unused files from the compilation unit included_contractnames = set() for target_file_path in files_to_include: - for node in compilation_unit.asts[target_file_path]['nodes']: - if node['nodeType'] == 'ContractDefinition': - included_contractnames.add(node['name']) + for node in compilation_unit.asts[target_file_path]["nodes"]: + if node["nodeType"] == "ContractDefinition": + included_contractnames.add(node["name"]) for contractname in list(compilation_unit.contracts_names): if contractname not in included_contractnames: @@ -475,10 +476,10 @@ def _remove_unused_contracts(compilation_unit: CompilationUnit, export_dir: str) del compilation_unit.asts[fileobj.absolute] # remove all unused files - for filename in glob.iglob(os.path.join(export_dir, '**/**'), recursive=True): + for filename in glob.iglob(os.path.join(export_dir, "**/**"), recursive=True): if os.path.isfile(filename) and filename not in files_to_include: os.remove(filename) # remove all folders which are now empty - for filename in glob.iglob(os.path.join(export_dir, '**/**'), recursive=True): + for filename in glob.iglob(os.path.join(export_dir, "**/**"), recursive=True): if os.path.isdir(filename) and len(os.listdir(filename)) == 0: - shutil.rmtree(filename) \ No newline at end of file + shutil.rmtree(filename) From 6ee79fe63755c321dd6f9bd7d79ffdb42ab12f00 Mon Sep 17 00:00:00 2001 From: Alexander Remie Date: Thu, 17 Mar 2022 17:37:00 +0100 Subject: [PATCH 3/4] improve comment --- crytic_compile/platform/etherscan.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crytic_compile/platform/etherscan.py b/crytic_compile/platform/etherscan.py index 71408894..207b342c 100644 --- a/crytic_compile/platform/etherscan.py +++ b/crytic_compile/platform/etherscan.py @@ -415,7 +415,7 @@ def _remove_unused_contracts(compilation_unit: CompilationUnit, export_dir: str) # for etherscan this will be the value the etherscan api returns in 'ContractName' root_contract_name = compilation_unit.unique_id - # find the root file path by it's name + # find the root file path according to a contract with the correct name being defined in it # and also get the base path used by all paths (the keys under 'asts') root_file_path = None base_path = "" From a9489f31f2979a1bb0a531c2c0b9663efddb1b74 Mon Sep 17 00:00:00 2001 From: Alexander Remie Date: Fri, 18 Mar 2022 15:46:19 +0100 Subject: [PATCH 4/4] add --etherscan-target-only flag and do not remove files/dirs from filesystem --- crytic_compile/cryticparser/cryticparser.py | 8 ++++++++ crytic_compile/cryticparser/defaults.py | 1 + crytic_compile/platform/etherscan.py | 21 ++++++--------------- 3 files changed, 15 insertions(+), 15 deletions(-) diff --git a/crytic_compile/cryticparser/cryticparser.py b/crytic_compile/cryticparser/cryticparser.py index 19df89e7..f525338d 100755 --- a/crytic_compile/cryticparser/cryticparser.py +++ b/crytic_compile/cryticparser/cryticparser.py @@ -304,6 +304,14 @@ def _init_etherscan(parser: ArgumentParser) -> None: default=DEFAULTS_FLAG_IN_CONFIG["etherscan_only_bytecode"], ) + group_etherscan.add_argument( + "--etherscan-target-only", + help="Etherscan only include target contract.", + action="store_true", + dest="etherscan_target_only", + default=DEFAULTS_FLAG_IN_CONFIG["etherscan_target_only"], + ) + group_etherscan.add_argument( "--etherscan-apikey", help="Etherscan API key.", diff --git a/crytic_compile/cryticparser/defaults.py b/crytic_compile/cryticparser/defaults.py index b8f99677..b196a392 100755 --- a/crytic_compile/cryticparser/defaults.py +++ b/crytic_compile/cryticparser/defaults.py @@ -30,6 +30,7 @@ "etherlime_compile_arguments": None, "etherscan_only_source_code": False, "etherscan_only_bytecode": False, + "etherscan_target_only": False, "etherscan_api_key": None, "etherscan_export_directory": "etherscan-contracts", "waffle_ignore_compile": False, diff --git a/crytic_compile/platform/etherscan.py b/crytic_compile/platform/etherscan.py index 207b342c..da1fdcb4 100644 --- a/crytic_compile/platform/etherscan.py +++ b/crytic_compile/platform/etherscan.py @@ -7,8 +7,6 @@ import os import re import urllib.request -import glob -import shutil from json.decoder import JSONDecodeError from pathlib import Path from typing import TYPE_CHECKING, Dict, List, Union, Tuple, Optional @@ -204,6 +202,8 @@ def compile(self, crytic_compile: "CryticCompile", **kwargs: str) -> None: only_source = kwargs.get("etherscan_only_source_code", False) only_bytecode = kwargs.get("etherscan_only_bytecode", False) + target_only = kwargs.get("etherscan_target_only", False) + etherscan_api_key = kwargs.get("etherscan_api_key", None) arbiscan_api_key = kwargs.get("arbiscan_api_key", None) polygonscan_api_key = kwargs.get("polygonscan_api_key", None) @@ -332,7 +332,8 @@ def compile(self, crytic_compile: "CryticCompile", **kwargs: str) -> None: solc_standard_json.standalone_compile(filenames, compilation_unit, working_dir=working_dir) - _remove_unused_contracts(compilation_unit, export_dir) + if target_only: + _remove_unused_contracts(compilation_unit) @staticmethod def is_supported(target: str, **kwargs: str) -> bool: @@ -397,13 +398,12 @@ def _relative_to_short(relative: Path) -> Path: # pylint: disable=too-many-locals,too-many-branches -def _remove_unused_contracts(compilation_unit: CompilationUnit, export_dir: str) -> None: +def _remove_unused_contracts(compilation_unit: CompilationUnit) -> None: """ - Removes unused contracts from the compilation unit and file system + Removes unused contracts from the compilation unit Args: compilation_unit (CompilationUnit): compilation unit to populate - export_dir (str): export dir Returns: @@ -474,12 +474,3 @@ def _remove_unused_contracts(compilation_unit: CompilationUnit, export_dir: str) compilation_unit.crytic_compile.filenames.remove(fileobj) compilation_unit.filenames.remove(fileobj) del compilation_unit.asts[fileobj.absolute] - - # remove all unused files - for filename in glob.iglob(os.path.join(export_dir, "**/**"), recursive=True): - if os.path.isfile(filename) and filename not in files_to_include: - os.remove(filename) - # remove all folders which are now empty - for filename in glob.iglob(os.path.join(export_dir, "**/**"), recursive=True): - if os.path.isdir(filename) and len(os.listdir(filename)) == 0: - shutil.rmtree(filename)