crytic_compile.source_unit

  1"""
  2Module handling the source unit
  3Each source unit represents one file so may be associated with
  4One or more source units are associated with each compilation unit
  5"""
  6import re
  7from typing import Dict, List, Optional, Union, Tuple, TYPE_CHECKING
  8import cbor2
  9
 10from Crypto.Hash import keccak
 11
 12from crytic_compile.utils.naming import Filename
 13from crytic_compile.utils.natspec import Natspec
 14
 15if TYPE_CHECKING:
 16    from crytic_compile.compilation_unit import CompilationUnit
 17
 18
 19def get_library_candidate(filename: Filename, contract_name: str) -> List[str]:
 20    """
 21    Return candidate name for library linking. A candidate is a str that might be found in other bytecodes
 22
 23    Args:
 24        filename: filename of the contract
 25        contract_name: contract name
 26
 27    Returns:
 28        The list of candidates
 29    """
 30
 31    # Some platform use only the contract name
 32    # Some use fimename:contract_name
 33
 34    ret: List[str] = []
 35
 36    name_with_absolute_filename = filename.absolute + ":" + contract_name
 37    name_with_used_filename = filename.used + ":" + contract_name
 38
 39    # Only 36 char were used in the past
 40    # See https://docs.soliditylang.org/en/develop/using-the-compiler.html#library-linking
 41    names_candidates = [
 42        name_with_absolute_filename,
 43        name_with_absolute_filename[0:36],
 44        name_with_used_filename,
 45        name_with_used_filename[0:36],
 46    ]
 47
 48    # Solidity 0.4
 49    ret.append("__" + contract_name + "_" * (38 - len(contract_name)))
 50
 51    for name_candidate in names_candidates:
 52        # Solidity 0.4 with filename
 53        ret.append("__" + name_candidate + "_" * (38 - len(name_candidate)))
 54
 55        # Solidity 0.5
 56        sha3_result = keccak.new(digest_bits=256)
 57        sha3_result.update(name_candidate.encode("utf-8"))
 58        ret.append("__$" + sha3_result.hexdigest()[:34] + "$__")
 59
 60    return ret
 61
 62
 63# pylint: disable=too-many-instance-attributes,too-many-public-methods
 64class SourceUnit:
 65    """SourceUnit class"""
 66
 67    def __init__(self, compilation_unit: "CompilationUnit", filename: Filename):
 68
 69        self.filename = filename
 70        self.compilation_unit: "CompilationUnit" = compilation_unit
 71
 72        # ABI, bytecode and srcmap are indexed by contract_name
 73        self._abis: Dict = {}
 74        self._runtime_bytecodes: Dict = {}
 75        self._init_bytecodes: Dict = {}
 76        self._hashes: Dict = {}
 77        self._events: Dict = {}
 78        self._srcmaps: Dict[str, List[str]] = {}
 79        self._srcmaps_runtime: Dict[str, List[str]] = {}
 80        self.ast: Dict = {}
 81
 82        # Natspec
 83        self._natspec: Dict[str, Natspec] = {}
 84
 85        # Libraries used by the contract
 86        # contract_name -> (library, pattern)
 87        self._libraries: Dict[str, List[Tuple[str, str]]] = {}
 88
 89        # set containing all the contract names
 90        self._contracts_name: List[str] = []
 91
 92        # set containing all the contract name without the libraries
 93        self._contracts_name_without_libraries: Optional[List[str]] = None
 94
 95    # region ABI
 96    ###################################################################################
 97    ###################################################################################
 98
 99    @property
100    def abis(self) -> Dict:
101        """Return the ABIs
102
103        Returns:
104            Dict: ABIs (solc/vyper format) (contract name -> ABI)
105        """
106        return self._abis
107
108    def abi(self, name: str) -> Dict:
109        """Get the ABI from a contract
110
111        Args:
112            name (str): Contract name
113
114        Returns:
115            Dict: ABI (solc/vyper format)
116        """
117        return self._abis.get(name, None)
118
119    # endregion
120    ###################################################################################
121    ###################################################################################
122    # region Bytecode
123    ###################################################################################
124    ###################################################################################
125
126    @property
127    def bytecodes_runtime(self) -> Dict[str, str]:
128        """Return the runtime bytecodes
129
130        Returns:
131            Dict[str, str]: contract => runtime bytecode
132        """
133        return self._runtime_bytecodes
134
135    @bytecodes_runtime.setter
136    def bytecodes_runtime(self, bytecodes: Dict[str, str]) -> None:
137        """Set the bytecodes runtime
138
139        Args:
140            bytecodes (Dict[str, str]): New bytecodes runtime
141        """
142        self._runtime_bytecodes = bytecodes
143
144    @property
145    def bytecodes_init(self) -> Dict[str, str]:
146        """Return the init bytecodes
147
148        Returns:
149            Dict[str, str]: contract => init bytecode
150        """
151        return self._init_bytecodes
152
153    @bytecodes_init.setter
154    def bytecodes_init(self, bytecodes: Dict[str, str]) -> None:
155        """Set the bytecodes init
156
157        Args:
158            bytecodes (Dict[str, str]): New bytecodes init
159        """
160        self._init_bytecodes = bytecodes
161
162    def bytecode_runtime(self, name: str, libraries: Optional[Dict[str, int]] = None) -> str:
163        """Return the runtime bytecode of the contract.
164        If library is provided, patch the bytecode
165
166        Args:
167            name (str): contract name
168            libraries (Optional[Dict[str, str]], optional): lib_name => address. Defaults to None.
169
170        Returns:
171            str: runtime bytecode
172        """
173        runtime = self._runtime_bytecodes.get(name, None)
174        return self._update_bytecode_with_libraries(runtime, libraries)
175
176    def bytecode_init(self, name: str, libraries: Optional[Dict[str, int]] = None) -> str:
177        """Return the init bytecode of the contract.
178        If library is provided, patch the bytecode
179
180        Args:
181            name (str): contract name
182            libraries (Optional[Dict[str, int]], optional): lib_name => address. Defaults to None.
183
184        Returns:
185            str: init bytecode
186        """
187        init = self._init_bytecodes.get(name, None)
188        return self._update_bytecode_with_libraries(init, libraries)
189
190    # endregion
191    ###################################################################################
192    ###################################################################################
193    # region Source mapping
194    ###################################################################################
195    ###################################################################################
196
197    @property
198    def srcmaps_init(self) -> Dict[str, List[str]]:
199        """Return the srcmaps init
200
201        Returns:
202            Dict[str, List[str]]: Srcmaps init (solc/vyper format)
203        """
204        return self._srcmaps
205
206    @property
207    def srcmaps_runtime(self) -> Dict[str, List[str]]:
208        """Return the srcmaps runtime
209
210        Returns:
211            Dict[str, List[str]]: Srcmaps runtime (solc/vyper format)
212        """
213        return self._srcmaps_runtime
214
215    def srcmap_init(self, name: str) -> List[str]:
216        """Return the srcmap init of a contract
217
218        Args:
219            name (str): name of the contract
220
221        Returns:
222            List[str]: Srcmap init (solc/vyper format)
223        """
224        return self._srcmaps.get(name, [])
225
226    def srcmap_runtime(self, name: str) -> List[str]:
227        """Return the srcmap runtime of a contract
228
229        Args:
230            name (str): name of the contract
231
232        Returns:
233            List[str]: Srcmap runtime (solc/vyper format)
234        """
235        return self._srcmaps_runtime.get(name, [])
236
237    # endregion
238    ###################################################################################
239    ###################################################################################
240    # region Libraries
241    ###################################################################################
242    ###################################################################################
243
244    @property
245    def libraries(self) -> Dict[str, List[Tuple[str, str]]]:
246        """Return the libraries used
247
248        Returns:
249            Dict[str, List[Tuple[str, str]]]:  (contract_name -> [(library, pattern))])
250        """
251        return self._libraries
252
253    def _convert_libraries_names(self, libraries: Dict[str, int]) -> Dict[str, int]:
254        """Convert the libraries names
255        The name in the argument can be the library name, or filename:library_name
256        The returned dict contains all the names possible with the different solc versions
257
258        Args:
259            libraries (Dict[str, int]): lib_name => address
260
261        Returns:
262            Dict[str, int]: lib_name => address
263        """
264        new_names = {}
265        for (lib, addr) in libraries.items():
266            # Prior solidity 0.5
267            # libraries were on the format __filename:contract_name_____
268            # From solidity 0.5,
269            # libraries are on the format __$keccak(filename:contract_name)[34]$__
270            # https://solidity.readthedocs.io/en/v0.5.7/050-breaking-changes.html#command-line-and-json-interfaces
271
272            lib_4 = "__" + lib + "_" * (38 - len(lib))
273
274            sha3_result = keccak.new(digest_bits=256)
275            sha3_result.update(lib.encode("utf-8"))
276            lib_5 = "__$" + sha3_result.hexdigest()[:34] + "$__"
277
278            new_names[lib] = addr
279            new_names[lib_4] = addr
280            new_names[lib_5] = addr
281
282            for lib_filename, contract_names in self.compilation_unit.filename_to_contracts.items():
283                for contract_name in contract_names:
284                    if contract_name != lib:
285                        continue
286
287                    for candidate in get_library_candidate(lib_filename, lib):
288                        new_names[candidate] = addr
289
290        return new_names
291
292    def _library_name_lookup(
293        self, lib_name: str, original_contract: str
294    ) -> Optional[Tuple[str, str]]:
295        """Do a lookup on a library name to its name used in contracts
296        The library can be:
297        - the original contract name
298        - __X__ following Solidity 0.4 format
299        - __$..$__ following Solidity 0.5 format
300
301        Args:
302            lib_name (str): library name
303            original_contract (str): original contract name
304
305        Returns:
306            Optional[Tuple[str, str]]: contract_name, library_name
307        """
308
309        for filename, contract_names in self.compilation_unit.filename_to_contracts.items():
310            for name in contract_names:
311                if name == lib_name:
312                    return name, name
313
314                for candidate in get_library_candidate(filename, name):
315                    if candidate == lib_name:
316                        return name, candidate
317
318        # handle specific case of collision for Solidity <0.4
319        # We can only detect that the second contract is meant to be the library
320        # if there is only two contracts in the codebase
321        if len(self._contracts_name) == 2:
322            return next(
323                (
324                    (c, "__" + c + "_" * (38 - len(c)))
325                    for c in self._contracts_name
326                    if c != original_contract
327                ),
328                None,
329            )
330
331        return None
332
333    def libraries_names(self, name: str) -> List[str]:
334        """Return the names of the libraries used by the contract
335
336        Args:
337            name (str): contract name
338
339        Returns:
340            List[str]: libraries used
341        """
342
343        if name not in self._libraries:
344            init = re.findall(r"__.{36}__", self.bytecode_init(name))
345            runtime = re.findall(r"__.{36}__", self.bytecode_runtime(name))
346            libraires = [self._library_name_lookup(x, name) for x in set(init + runtime)]
347            self._libraries[name] = [lib for lib in libraires if lib]
348        return [name for (name, _) in self._libraries[name]]
349
350    def libraries_names_and_patterns(self, name: str) -> List[Tuple[str, str]]:
351        """Return the names and the patterns of the libraries used by the contract
352
353        Args:
354            name (str): contract name
355
356        Returns:
357            List[Tuple[str, str]]: (lib_name, pattern)
358        """
359
360        if name not in self._libraries:
361            init = re.findall(r"__.{36}__", self.bytecode_init(name))
362            runtime = re.findall(r"__.{36}__", self.bytecode_runtime(name))
363            libraires = [self._library_name_lookup(x, name) for x in set(init + runtime)]
364            self._libraries[name] = [lib for lib in libraires if lib]
365        return self._libraries[name]
366
367    def _update_bytecode_with_libraries(
368        self, bytecode: str, libraries: Union[None, Dict[str, int]]
369    ) -> str:
370        """Update the bytecode with the libraries address
371
372        Args:
373            bytecode (str): bytecode to patch
374            libraries (Union[None, Dict[str, int]]): pattern => address
375
376        Returns:
377            str: Patched bytecode
378        """
379        if libraries:
380            libraries = self._convert_libraries_names(libraries)
381            for library_found in re.findall(r"__.{36}__", bytecode):
382                if library_found in libraries:
383                    bytecode = re.sub(
384                        re.escape(library_found),
385                        f"{libraries[library_found]:0>40x}",
386                        bytecode,
387                    )
388        return bytecode
389
390    # endregion
391    ###################################################################################
392    ###################################################################################
393    # region Natspec
394    ###################################################################################
395    ###################################################################################
396
397    @property
398    def natspec(self) -> Dict[str, Natspec]:
399        """Return the natspec of the contracts
400
401        Returns:
402            Dict[str, Natspec]: Contract name -> Natspec
403        """
404        return self._natspec
405
406    # endregion
407    ###################################################################################
408    ###################################################################################
409    # region Contract Names
410    ###################################################################################
411    ###################################################################################
412
413    @property
414    def contracts_names(self) -> List[str]:
415        """Return the contracts names
416
417        Returns:
418            List[str]: List of the contracts names
419        """
420        return self._contracts_name
421
422    @contracts_names.setter
423    def contracts_names(self, names: List[str]) -> None:
424        """Set the contract names
425
426        Args:
427            names (List[str]): New contracts names
428        """
429        self._contracts_name = names
430
431    def add_contract_name(self, name: str) -> None:
432        """Add name to contracts_names, if not already present
433
434        Args:
435            name (str): Name to add to the list
436        """
437        if name not in self.contracts_names:
438            self.contracts_names.append(name)
439
440    @property
441    def contracts_names_without_libraries(self) -> List[str]:
442        """Return the contracts names without the librairies
443
444        Returns:
445            List[str]: List of contracts
446        """
447        if self._contracts_name_without_libraries is None:
448            libraries: List[str] = []
449            for contract_name in self._contracts_name:
450                libraries += self.libraries_names(contract_name)
451            self._contracts_name_without_libraries = [
452                l for l in self._contracts_name if l not in set(libraries)
453            ]
454        return self._contracts_name_without_libraries
455
456    # endregion
457    ###################################################################################
458    ###################################################################################
459    # region Hashes
460    ###################################################################################
461    ###################################################################################
462
463    def hashes(self, name: str) -> Dict[str, int]:
464        """Return the hashes of the functions
465
466        Args:
467            name (str): contract name
468
469        Returns:
470            Dict[str, int]: (function name => signature)
471        """
472        if not name in self._hashes:
473            self._compute_hashes(name)
474        return self._hashes[name]
475
476    def _compute_hashes(self, name: str) -> None:
477        """Compute the function hashes
478
479        Args:
480            name (str): contract name
481        """
482        self._hashes[name] = {}
483        for sig in self.abi(name):
484            if "type" in sig:
485                if sig["type"] == "function":
486                    sig_name = sig["name"]
487                    arguments = ",".join([x["type"] for x in sig["inputs"]])
488                    sig = f"{sig_name}({arguments})"
489                    sha3_result = keccak.new(digest_bits=256)
490                    sha3_result.update(sig.encode("utf-8"))
491                    self._hashes[name][sig] = int("0x" + sha3_result.hexdigest()[:8], 16)
492
493    # endregion
494    ###################################################################################
495    ###################################################################################
496    # region Events
497    ###################################################################################
498    ###################################################################################
499
500    def events_topics(self, name: str) -> Dict[str, Tuple[int, List[bool]]]:
501        """Return the topics of the contract's events
502
503        Args:
504            name (str): contract name
505
506        Returns:
507            Dict[str, Tuple[int, List[bool]]]: event signature => topic hash, [is_indexed for each parameter]
508        """
509        if not name in self._events:
510            self._compute_topics_events(name)
511        return self._events[name]
512
513    def _compute_topics_events(self, name: str) -> None:
514        """Compute the topics of the contract's events
515
516        Args:
517            name (str): contract name
518        """
519        self._events[name] = {}
520        for sig in self.abi(name):
521            if "type" in sig:
522                if sig["type"] == "event":
523                    sig_name = sig["name"]
524                    arguments = ",".join([x["type"] for x in sig["inputs"]])
525                    indexes = [x.get("indexed", False) for x in sig["inputs"]]
526                    sig = f"{sig_name}({arguments})"
527                    sha3_result = keccak.new(digest_bits=256)
528                    sha3_result.update(sig.encode("utf-8"))
529
530                    self._events[name][sig] = (int("0x" + sha3_result.hexdigest()[:8], 16), indexes)
531
532    # endregion
533    ###################################################################################
534    ###################################################################################
535    # region Metadata
536    ###################################################################################
537    ###################################################################################
538
539    def metadata_of(self, name: str) -> Dict[str, Union[str, bool]]:
540        """Return the parsed metadata of a contract by name
541
542        Args:
543            name (str): contract name
544
545        Raises:
546            ValueError: If no contract/library with that name exists
547
548        Returns:
549            Dict[str, Union[str, bool]]: fielname => value
550        """
551        # the metadata is at the end of the runtime(!) bytecode
552        try:
553            bytecode = self._runtime_bytecodes[name]
554            print("runtime bytecode", bytecode)
555        except:
556            raise ValueError(  # pylint: disable=raise-missing-from
557                f"contract {name} does not exist"
558            )
559
560        # the last two bytes contain the length of the preceding metadata.
561        metadata_length = int(f"0x{bytecode[-4:]}", base=16)
562        # extract the metadata
563        metadata = bytecode[-(metadata_length * 2 + 4) :]
564        metadata_decoded = cbor2.loads(bytearray.fromhex(metadata))
565
566        for k, v in metadata_decoded.items():
567            if len(v) == 1:
568                metadata_decoded[k] = bool(v)
569            elif k == "solc":
570                metadata_decoded[k] = ".".join([str(d) for d in v])
571            else:
572                # there might be nested items or other unforeseen errors
573                try:
574                    metadata_decoded[k] = v.hex()
575                except:  # pylint: disable=bare-except
576                    pass
577
578        return metadata_decoded
579
580    def remove_metadata(self) -> None:
581        """Remove init bytecode
582        See
583        http://solidity.readthedocs.io/en/v0.4.24/metadata.html#encoding-of-the-metadata-hash-in-the-bytecode
584        """
585        # the metadata is at the end of the runtime(!) bytecode of each contract
586        for (key, bytecode) in self._runtime_bytecodes.items():
587            if not bytecode or bytecode == "0x":
588                continue
589            # the last two bytes contain the length of the preceding metadata.
590            metadata_length = int(f"0x{bytecode[-4:]}", base=16)
591            # store the metadata here so we can remove it from the init bytecode later on
592            metadata = bytecode[-(metadata_length * 2 + 4) :]
593            # remove the metadata from the runtime bytecode, '+ 4' for the two length-indication bytes at the end
594            self._runtime_bytecodes[key] = bytecode[0 : -(metadata_length * 2 + 4)]
595            # remove the metadata from the init bytecode
596            self._init_bytecodes[key] = self._init_bytecodes[key].replace(metadata, "")
597
598    # endregion
599    ###################################################################################
600    ###################################################################################