etherscan.py

"""
Etherscan platform.
"""

import json
import logging
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

from crytic_compile.compilation_unit import CompilationUnit
from crytic_compile.compiler.compiler import CompilerVersion
from crytic_compile.platform import solc_standard_json
from crytic_compile.platform.abstract_platform import AbstractPlatform
from crytic_compile.platform.exceptions import InvalidCompilation
from crytic_compile.platform.types import Type
from crytic_compile.utils.naming import Filename

# Cycle dependency

if TYPE_CHECKING:
    from crytic_compile import CryticCompile

LOGGER = logging.getLogger("CryticCompile")


ETHERSCAN_BASE = "https://api%s/api?module=contract&action=getsourcecode&address=%s"

ETHERSCAN_BASE_BYTECODE = "https://%s/address/%s#code"

SUPPORTED_NETWORK = {
    # Key, (prefix_base, perfix_bytecode)
    "mainet:": (".etherscan.io", "etherscan.io"),
    "ropsten:": ("-ropsten.etherscan.io", "ropsten.etherscan.io"),
    "kovan:": ("-kovan.etherscan.io", "kovan.etherscan.io"),
    "rinkeby:": ("-rinkeby.etherscan.io", "rinkeby.etherscan.io"),
    "goerli:": ("-goerli.etherscan.io", "goerli.etherscan.io"),
    "tobalaba:": ("-tobalaba.etherscan.io", "tobalaba.etherscan.io"),
    "bsc:": (".bscscan.com", "bscscan.com"),
    "testnet.bsc:": ("-testnet.bscscan.com", "testnet.bscscan.com"),
    "arbi:": (".arbiscan.io", "arbiscan.io"),
    "testnet.arbi:": ("-testnet.arbiscan.io", "testnet.arbiscan.io"),
    "poly:": (".polygonscan.com", "polygonscan.com"),
    "avax:": (".snowtrace.io", "snowtrace.io"),
    "testnet.avax:": ("-testnet.snowtrace.io", "testnet.snowtrace.io"),
    "ftm:": (".ftmscan.com", "ftmscan.com"),
}


def _handle_bytecode(crytic_compile: "CryticCompile", target: str, result_b: bytes) -> None:
    """Parse the bytecode and populate CryticCompile info

    Args:
        crytic_compile (CryticCompile): Associate CryticCompile object
        target (str): path to the target
        result_b (bytes): text containing the bytecode
    """

    # There is no direct API to get the bytecode from etherscan
    # The page changes from time to time, we use for now a simple parsing, it will not be robust
    begin = """Search Algorithm">\nSimilar Contracts</button>\n"""
    begin += """<div id="dividcode">\n<pre class=\'wordwrap\' style=\'height: 15pc;\'>0x"""
    result = result_b.decode("utf8")
    # Removing everything before the begin string
    result = result[result.find(begin) + len(begin) :]
    bytecode = result[: result.find("<")]

    contract_name = f"Contract_{target}"

    contract_filename = Filename(absolute="", relative="", short="", used="")

    compilation_unit = CompilationUnit(crytic_compile, str(target))

    compilation_unit.contracts_names.add(contract_name)
    compilation_unit.filename_to_contracts[contract_filename].add(contract_name)
    compilation_unit.abis[contract_name] = {}
    compilation_unit.bytecodes_init[contract_name] = bytecode
    compilation_unit.bytecodes_runtime[contract_name] = ""
    compilation_unit.srcmaps_init[contract_name] = []
    compilation_unit.srcmaps_runtime[contract_name] = []

    compilation_unit.compiler_version = CompilerVersion(
        compiler="unknown", version="", optimized=False
    )

    crytic_compile.bytecode_only = True


# def _etherscan_single_file():


def _handle_single_file(
    source_code: str, addr: str, prefix: Optional[str], contract_name: str, export_dir: str
) -> str:
    """Handle a result with a single file

    Args:
        source_code (str): source code
        addr (str): contract address
        prefix (Optional[str]): used to separate different chains
        contract_name (str): contract name
        export_dir (str): directory where the code will be saved

    Returns:
        str: filename containing the source code
    """
    if prefix:
        filename = os.path.join(export_dir, f"{addr}{prefix}-{contract_name}.sol")
    else:
        filename = os.path.join(export_dir, f"{addr}-{contract_name}.sol")

    with open(filename, "w", encoding="utf8") as file_desc:
        file_desc.write(source_code)

    return filename


def _handle_multiple_files(
    dict_source_code: Dict, addr: str, prefix: Optional[str], contract_name: str, export_dir: str
) -> Tuple[List[str], str]:
    """Handle a result with a multiple files. Generate multiple Solidity files

    Args:
        dict_source_code (Dict): dict result from etherscan
        addr (str): contract address
        prefix (Optional[str]): used to separate different chains
        contract_name (str): contract name
        export_dir (str): directory where the code will be saved

    Returns:
        Tuple[List[str], str]: filesnames, directory, where target_filename is the main file
    """
    if prefix:
        directory = os.path.join(export_dir, f"{addr}{prefix}-{contract_name}")
    else:
        directory = os.path.join(export_dir, f"{addr}-{contract_name}")

    if "sources" in dict_source_code:
        # etherscan might return an object with a sources prop, which contains an object with contract names as keys
        source_codes = dict_source_code["sources"]
    else:
        # or etherscan might return an object with contract names as keys
        source_codes = dict_source_code

    filtered_paths: List[str] = []
    for filename, source_code in source_codes.items():
        path_filename = Path(filename)
        if "contracts" in path_filename.parts and not filename.startswith("@"):
            path_filename = Path(*path_filename.parts[path_filename.parts.index("contracts") :])

        filtered_paths.append(str(path_filename))
        path_filename = Path(directory, path_filename)

        if not os.path.exists(path_filename.parent):
            os.makedirs(path_filename.parent)
        with open(path_filename, "w", encoding="utf8") as file_desc:
            file_desc.write(source_code["content"])

    return list(filtered_paths), directory


class Etherscan(AbstractPlatform):
    """
    Etherscan platform
    """

    NAME = "Etherscan"
    PROJECT_URL = "https://etherscan.io/"
    TYPE = Type.ETHERSCAN

    # pylint: disable=too-many-locals,too-many-branches,too-many-statements
    def compile(self, crytic_compile: "CryticCompile", **kwargs: str) -> None:
        """Run the compilation

        Args:
            crytic_compile (CryticCompile): Associated CryticCompile object
            **kwargs: optional arguments. Used "solc", "etherscan_only_source_code", "etherscan_only_bytecode",
                "etherscan_api_key", "export_dir"

        Raises:
            InvalidCompilation: if etherscan returned an error, or its results were not correctly parsed
        """

        target = self._target

        if target.startswith(tuple(SUPPORTED_NETWORK)):
            prefix: Union[None, str] = SUPPORTED_NETWORK[target[: target.find(":") + 1]][0]
            prefix_bytecode = SUPPORTED_NETWORK[target[: target.find(":") + 1]][1]
            addr = target[target.find(":") + 1 :]
            etherscan_url = ETHERSCAN_BASE % (prefix, addr)
            etherscan_bytecode_url = ETHERSCAN_BASE_BYTECODE % (prefix_bytecode, addr)

        else:
            etherscan_url = ETHERSCAN_BASE % (".etherscan.io", target)
            etherscan_bytecode_url = ETHERSCAN_BASE_BYTECODE % ("etherscan.io", target)
            addr = target
            prefix = None

        only_source = kwargs.get("etherscan_only_source_code", False)
        only_bytecode = kwargs.get("etherscan_only_bytecode", 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)
        avax_api_key = kwargs.get("avax_api_key", None)
        ftmscan_api_key = kwargs.get("ftmscan_api_key", None)
        bscan_api_key = kwargs.get("bscan_api_key", None)

        export_dir = kwargs.get("export_dir", "crytic-export")
        export_dir = os.path.join(
            export_dir, kwargs.get("etherscan_export_dir", "etherscan-contracts")
        )

        if etherscan_api_key and "etherscan" in etherscan_url:
            etherscan_url += f"&apikey={etherscan_api_key}"
            etherscan_bytecode_url += f"&apikey={etherscan_api_key}"
        if arbiscan_api_key and "arbiscan" in etherscan_url:
            etherscan_url += f"&apikey={arbiscan_api_key}"
            etherscan_bytecode_url += f"&apikey={arbiscan_api_key}"
        if polygonscan_api_key and "polygonscan" in etherscan_url:
            etherscan_url += f"&apikey={polygonscan_api_key}"
            etherscan_bytecode_url += f"&apikey={polygonscan_api_key}"
        if avax_api_key and "snowtrace" in etherscan_url:
            etherscan_url += f"&apikey={avax_api_key}"
            etherscan_bytecode_url += f"&apikey={avax_api_key}"
        if ftmscan_api_key and "ftmscan" in etherscan_url:
            etherscan_url += f"&apikey={ftmscan_api_key}"
            etherscan_bytecode_url += f"&apikey={ftmscan_api_key}"
        if bscan_api_key and "bscscan" in etherscan_url:
            etherscan_url += f"&apikey={bscan_api_key}"
            etherscan_bytecode_url += f"&apikey={bscan_api_key}"

        source_code: str = ""
        result: Dict[str, Union[bool, str, int]] = {}
        contract_name: str = ""

        if not only_bytecode:
            with urllib.request.urlopen(etherscan_url) as response:
                html = response.read()

            info = json.loads(html)

            if "result" in info and info["result"] == "Max rate limit reached":
                LOGGER.error("Etherscan API rate limit exceeded")
                raise InvalidCompilation("Etherscan api rate limit exceeded")

            if "message" not in info:
                LOGGER.error("Incorrect etherscan request")
                raise InvalidCompilation("Incorrect etherscan request " + etherscan_url)

            if not info["message"].startswith("OK"):
                LOGGER.error("Contract has no public source code")
                raise InvalidCompilation("Contract has no public source code: " + etherscan_url)

            if "result" not in info:
                LOGGER.error("Contract has no public source code")
                raise InvalidCompilation("Contract has no public source code: " + etherscan_url)

            result = info["result"][0]
            # Assert to help mypy
            assert isinstance(result["SourceCode"], str)
            assert isinstance(result["ContractName"], str)
            source_code = result["SourceCode"]
            contract_name = result["ContractName"]

        if source_code == "" and not only_source:
            LOGGER.info("Source code not available, try to fetch the bytecode only")

            req = urllib.request.Request(
                etherscan_bytecode_url, headers={"User-Agent": "Mozilla/5.0"}
            )
            with urllib.request.urlopen(req) as response:
                html = response.read()

            _handle_bytecode(crytic_compile, target, html)
            return

        if source_code == "":
            LOGGER.error("Contract has no public source code")
            raise InvalidCompilation("Contract has no public source code: " + etherscan_url)

        if not os.path.exists(export_dir):
            os.makedirs(export_dir)

        # Assert to help mypy
        assert isinstance(result["CompilerVersion"], str)

        compiler_version = re.findall(
            r"\d+\.\d+\.\d+", _convert_version(result["CompilerVersion"])
        )[0]

        optimization_used: bool = result["OptimizationUsed"] == "1"

        optimize_runs = None
        if optimization_used:
            optimize_runs = int(result["Runs"])

        working_dir: Optional[str] = None

        try:
            # etherscan might return an object with two curly braces, {{ content }}
            dict_source_code = json.loads(source_code[1:-1])
            filenames, working_dir = _handle_multiple_files(
                dict_source_code, addr, prefix, contract_name, export_dir
            )
        except JSONDecodeError:
            try:
                # or etherscan might return an object with single curly braces, { content }
                dict_source_code = json.loads(source_code)
                filenames, working_dir = _handle_multiple_files(
                    dict_source_code, addr, prefix, contract_name, export_dir
                )
            except JSONDecodeError:
                filenames = [
                    _handle_single_file(source_code, addr, prefix, contract_name, export_dir)
                ]

        compilation_unit = CompilationUnit(crytic_compile, contract_name)

        compilation_unit.compiler_version = CompilerVersion(
            compiler=kwargs.get("solc", "solc"),
            version=compiler_version,
            optimized=optimization_used,
            optimize_runs=optimize_runs,
        )
        compilation_unit.compiler_version.look_for_installed_version()

        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

        Args:
            target (str): path to the target
            **kwargs: optional arguments. Used "etherscan_ignore"

        Returns:
            bool: True if the target is a etherscan project
        """
        etherscan_ignore = kwargs.get("etherscan_ignore", False)
        if etherscan_ignore:
            return False
        if target.startswith(tuple(SUPPORTED_NETWORK)):
            target = target[target.find(":") + 1 :]
        return bool(re.match(r"^\s*0x[a-zA-Z0-9]{40}\s*$", target))

    def is_dependency(self, _path: str) -> bool:
        """Check if the path is a dependency

        Args:
            _path (str): path to the target

        Returns:
            bool: True if the target is a dependency
        """
        return False

    def _guessed_tests(self) -> List[str]:
        """Guess the potential unit tests commands

        Returns:
            List[str]: The guessed unit tests commands
        """
        return []


def _convert_version(version: str) -> str:
    """Convert the compiler version

    Args:
        version (str): original version

    Returns:
        str: converted version
    """
    return version[1 : version.find("+")]


def _relative_to_short(relative: Path) -> Path:
    """Translate relative path to short (do nothing for etherscan)

    Args:
        relative (Path): path to the target

    Returns:
        Path: Translated 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

    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 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 = ""
    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)