From e64de7f7b35cee88ba94127c1e77dd84a1585ee0 Mon Sep 17 00:00:00 2001 From: TheStarBoys Date: Sun, 24 Jul 2022 14:03:37 +0800 Subject: [PATCH] Add feature #1298 --- slither/core/declarations/contract.py | 9 +++++ slither/solc_parsing/declarations/contract.py | 2 + slither/tools/flattening/__main__.py | 7 ++++ slither/tools/flattening/flattening.py | 40 +++++++++++++++++++ 4 files changed, 58 insertions(+) diff --git a/slither/core/declarations/contract.py b/slither/core/declarations/contract.py index 2db1cd964..c70711502 100644 --- a/slither/core/declarations/contract.py +++ b/slither/core/declarations/contract.py @@ -78,6 +78,7 @@ class Contract(SourceMapping): # pylint: disable=too-many-public-methods self._using_for: Dict[Union[str, Type], List[Type]] = {} self._kind: Optional[str] = None self._is_interface: bool = False + self._is_library: bool = False self._signatures: Optional[List[str]] = None self._signatures_declared: Optional[List[str]] = None @@ -144,6 +145,14 @@ class Contract(SourceMapping): # pylint: disable=too-many-public-methods def is_interface(self, is_interface: bool): self._is_interface = is_interface + @property + def is_library(self) -> bool: + return self._is_library + + @is_library.setter + def is_library(self, is_library: bool): + self._is_library = is_library + # endregion ################################################################################### ################################################################################### diff --git a/slither/solc_parsing/declarations/contract.py b/slither/solc_parsing/declarations/contract.py index efc53f2aa..8823112fb 100644 --- a/slither/solc_parsing/declarations/contract.py +++ b/slither/solc_parsing/declarations/contract.py @@ -157,6 +157,8 @@ class ContractSolc(CallerContextExpression): if "contractKind" in attributes: if attributes["contractKind"] == "interface": self._contract.is_interface = True + elif attributes["contractKind"] == "library": + self._contract.is_library = True self._contract.kind = attributes["contractKind"] self._linearized_base_contracts = attributes["linearizedBaseContracts"] # self._contract.fullyImplemented = attributes["fullyImplemented"] diff --git a/slither/tools/flattening/__main__.py b/slither/tools/flattening/__main__.py index e0888f222..977b84896 100644 --- a/slither/tools/flattening/__main__.py +++ b/slither/tools/flattening/__main__.py @@ -79,6 +79,12 @@ def parse_args(): action="store_true", ) + group_patching.add_argument( + "--convert-library-to-internal", + help="Convert external or public functions to internal in library.", + action="store_true", + ) + group_patching.add_argument( "--remove-assert", help="Remove call to assert().", action="store_true" ) @@ -111,6 +117,7 @@ def main(): compilation_unit, external_to_public=args.convert_external, remove_assert=args.remove_assert, + convert_library_to_internal=args.convert_library_to_internal, private_to_internal=args.convert_private, export_path=args.dir, pragma_solidity=args.pragma_solidity, diff --git a/slither/tools/flattening/flattening.py b/slither/tools/flattening/flattening.py index f067bde6d..90b751afb 100644 --- a/slither/tools/flattening/flattening.py +++ b/slither/tools/flattening/flattening.py @@ -51,6 +51,7 @@ class Flattening: compilation_unit: SlitherCompilationUnit, external_to_public=False, remove_assert=False, + convert_library_to_internal=False, private_to_internal=False, export_path: Optional[str] = None, pragma_solidity: Optional[str] = None, @@ -61,6 +62,7 @@ class Flattening: self._external_to_public = external_to_public self._remove_assert = remove_assert self._use_abi_encoder_v2 = False + self._convert_library_to_internal = convert_library_to_internal self._private_to_internal = private_to_internal self._pragma_solidity = pragma_solidity @@ -146,6 +148,40 @@ class Flattening: ) ) + if self._convert_library_to_internal and contract.is_library: + for f in contract.functions_declared: + visibility = "" + if f.visibility == "external": + visibility = f.visibility + elif f.visibility == "public": + visibility = f.visibility + + if visibility != "": + attributes_start = ( + f.parameters_src().source_mapping["start"] + + f.parameters_src().source_mapping["length"] + ) + attributes_end = f.returns_src().source_mapping["start"] + attributes = content[attributes_start:attributes_end] + regex = ( + re.search(r"((\sexternal)\s+)|(\sexternal)$|(\)external)$", attributes) + if visibility == "external" + else re.search(r"((\spublic)\s+)|(\spublic)$|(\)public)$", attributes) + ) + if regex: + to_patch.append( + Patch( + attributes_start + regex.span()[0] + 1, + "external_to_internal" + if visibility == "external" + else "public_to_internal", + ) + ) + else: + raise SlitherException( + f"{visibility} keyword not found {f.name} {attributes}" + ) + if self._private_to_internal: for variable in contract.state_variables_declared: if variable.visibility == "private": @@ -186,6 +222,10 @@ class Flattening: index = index - start if patch_type == "public_to_external": content = content[:index] + "public" + content[index + len("external") :] + elif patch_type == "external_to_internal": + content = content[:index] + "internal" + content[index + len("external") :] + elif patch_type == "public_to_internal": + content = content[:index] + "internal" + content[index + len("public") :] elif patch_type == "private_to_internal": content = content[:index] + "internal" + content[index + len("private") :] elif patch_type == "calldata_to_memory":