Add support for displaying yul internal files (#1447)

mythril/mythx/__init__.py

@@ -50,7 +50,7 @@ def analyze(contracts: List[SolidityContract], analysis_mode: str = "quick") ->
 
         try:
             main_source = contract.input_file
-            for solidity_file in contract.solidity_files:
+            for solidity_file in contract.solc_indices.values():
                 source_codes[solidity_file.filename] = solidity_file.data
             for filename in contract.solc_json["sources"].keys():
                 sources[filename] = {}

mythril/solidity/soliditycontract.py

@@ -56,21 +56,23 @@ def get_contracts_from_file(input_file, solc_settings_json=None, solc_binary="so
     )
 
     try:
-        for contract_name in data["contracts"][input_file].keys():
-            if len(
-                data["contracts"][input_file][contract_name]["evm"]["deployedBytecode"][
-                    "object"
-                ]
-            ):
-                yield SolidityContract(
-                    input_file=input_file,
-                    name=contract_name,
-                    solc_settings_json=solc_settings_json,
-                    solc_binary=solc_binary,
-                )
+        contract_names = data["contracts"][input_file].keys()
     except KeyError:
         raise NoContractFoundError
 
+    for contract_name in contract_names:
+        if len(
+            data["contracts"][input_file][contract_name]["evm"]["deployedBytecode"][
+                "object"
+            ]
+        ):
+            yield SolidityContract(
+                input_file=input_file,
+                name=contract_name,
+                solc_settings_json=solc_settings_json,
+                solc_binary=solc_binary,
+            )
+
 
 class SolidityContract(EVMContract):
     """Representation of a Solidity contract."""
@@ -82,20 +84,10 @@ class SolidityContract(EVMContract):
             input_file, solc_settings_json=solc_settings_json, solc_binary=solc_binary
         )
 
-        self.solidity_files = []
+        self.solc_indices = self.get_solc_indices(data)
         self.solc_json = data
         self.input_file = input_file
 
-        for filename, contract in data["sources"].items():
-            with open(filename, "r", encoding="utf-8") as file:
-                code = file.read()
-                full_contract_src_maps = self.get_full_contract_src_maps(
-                    contract["ast"]
-                )
-                self.solidity_files.append(
-                    SolidityFile(filename, code, full_contract_src_maps)
-                )
-
         has_contract = False
 
         # If a contract name has been specified, find the bytecode of that specific contract
@@ -138,6 +130,39 @@ class SolidityContract(EVMContract):
 
         super().__init__(code, creation_code, name=name)
 
+    @staticmethod
+    def get_solc_indices(data: Dict) -> Dict:
+        """
+        Returns solc file indices
+        """
+        indices = {}
+        has_sources = True
+        for contract_data in data["contracts"].values():
+            for source_data in contract_data.values():
+                if "generatedSources" not in source_data["evm"]["deployedBytecode"]:
+                    has_sources = False
+                    break
+                sources = source_data["evm"]["deployedBytecode"]["generatedSources"]
+                for source in sources:
+                    full_contract_src_maps = SolidityContract.get_full_contract_src_maps(
+                        source["ast"]
+                    )
+                    indices[source["id"]] = SolidityFile(
+                        source["name"], source["contents"], full_contract_src_maps
+                    )
+            if has_sources is False:
+                break
+        for source in data["sources"].values():
+            full_contract_src_maps = SolidityContract.get_full_contract_src_maps(
+                source["ast"]
+            )
+            with open(source["ast"]["absolutePath"]) as f:
+                code = f.read()
+                indices[source["id"]] = SolidityFile(
+                    source["ast"]["absolutePath"], code, full_contract_src_maps
+                )
+        return indices
+
     @staticmethod
     def get_full_contract_src_maps(ast: Dict) -> Set[str]:
         """
@@ -146,9 +171,14 @@ class SolidityContract(EVMContract):
         :return: The source maps
         """
         source_maps = set()
-        for child in ast["nodes"]:
-            if child.get("contractKind"):
+        if ast["nodeType"] == "SourceUnit":
+            for child in ast["nodes"]:
+                if child.get("contractKind"):
+                    source_maps.add(child["src"])
+        elif ast["nodeType"] == "YulBlock":
+            for child in ast["statements"]:
                 source_maps.add(child["src"])
+
         return source_maps
 
     def get_source_info(self, address, constructor=False):
@@ -162,7 +192,7 @@ class SolidityContract(EVMContract):
         mappings = self.constructor_mappings if constructor else self.mappings
         index = helper.get_instruction_index(disassembly.instruction_list, address)
 
-        solidity_file = self.solidity_files[mappings[index].solidity_file_idx]
+        solidity_file = self.solc_indices[mappings[index].solidity_file_idx]
         filename = solidity_file.filename
 
         offset = mappings[index].offset
@@ -183,16 +213,12 @@ class SolidityContract(EVMContract):
         :return: True if the code is internally generated, else false
         """
 
-        """Handle internal compiler files
-        The second condition will be placed till 
-        https://github.com/ethereum/solidity/issues/10300 is dealt
-        """
-        if file_index == -1 or file_index >= len(self.solidity_files):
+        if file_index == -1:
             return True
         # Handle the common code src map for the entire code.
         if (
             "{}:{}:{}".format(offset, length, file_index)
-            in self.solidity_files[file_index].full_contract_src_maps
+            in self.solc_indices[file_index].full_contract_src_maps
         ):
             return True
 
@@ -224,7 +250,7 @@ class SolidityContract(EVMContract):
                 lineno = None
             else:
                 lineno = (
-                    self.solidity_files[idx]
+                    self.solc_indices[idx]
                     .data.encode("utf-8")[0:offset]
                     .count("\n".encode("utf-8"))
                     + 1

mythril/support/source_support.py

@@ -29,7 +29,9 @@ class Source:
             self.source_type = "solidity-file"
             self.source_format = "text"
             for contract in contracts:
-                self.source_list += [file.filename for file in contract.solidity_files]
+                self.source_list += [
+                    file.filename for file in contract.solc_indices.values()
+                ]
                 self._source_hash.append(contract.bytecode_hash)
                 self._source_hash.append(contract.creation_bytecode_hash)
         elif isinstance(contracts[0], EVMContract):