frame accounts with brownie

from typing import Any, Dict, List, Optional, Union
import socket
import logging

from brownie import accounts as brownie_accounts, rpc
from brownie._config import CONFIG
from brownie.convert import to_address, to_int
from brownie.network import web3 as brownie_web3
from brownie.network.account import _PrivateKeyAccount, Account
from eip712.messages import EIP712Message, _hash_eip191_message
from eth_account.datastructures import SignedMessage
from eth_account.messages import encode_defunct
from eth_keys.datatypes import Signature
from eth_typing.encoding import Primitives, HexStr

from eth_typing import URI
from eth_utils.applicators import apply_formatters_to_dict
from hexbytes import HexBytes
from web3 import WebsocketProvider
from web3.providers.base import JSONBaseProvider
from web3.providers.websocket import DEFAULT_WEBSOCKET_TIMEOUT, WebsocketProvider

from flashprofits.web3_helpers import get_chain_id

logger = logging.getLogger(__name__)


def hostname_resolves(hostname):
    try:
        socket.gethostbyname(hostname)
        return 1
    except socket.error:
        return 0


def docker_or_localhost():
    # TODO: check an env var instead of doing specific checks here?
    if hostname_resolves("host.docker.internal"):
        # we are inside docker for desktop
        return "host.docker.internal"

    return "127.0.0.1"


class ExternalWeb3Provider:
    def __init__(self, display_name: str, web3_provider: JSONBaseProvider) -> None:
        response = web3_provider.make_request("eth_chainId", [])

        if "error" in response:
            raise ValueError(response["error"]["message"])

        external_chain_id = to_int(response["result"])

        if brownie_web3.isConnected():
            while external_chain_id != get_chain_id():
                response = web3_provider.make_request("wallet_addEthereumChain", [])

                raise NotImplementedError
        else:
            logger.debug(
                "brownie is not connected. could not do a safety check on %s's chain id (%s)",
                display_name,
                external_chain_id,
            )

        self.display_name = display_name
        self._inner = web3_provider

    def brownie_request_accounts(self, expected_account=None) -> None:
        while True:
            if expected_account:
                logger.info("Waiting for account %s from %s...", expected_account, self.display_name)
            else:
                logger.info("Waiting for accounts from %s...", self.display_name)

            response = self._inner.make_request("eth_requestAccounts", [])

            if "error" in response:
                raise ValueError(response["error"]["message"])

            for address in response["result"]:
                if to_address(address) not in brownie_accounts._accounts:
                    brownie_accounts._accounts.append(ExternalWeb3Account(address, self))

            external_accounts = [i for i in brownie_accounts._accounts if getattr(i, "_external_web3_provider") == self]

            if expected_account and expected_account not in brownie_accounts._accounts:
                logger.error("Missing %s: %r", expected_account, external_accounts)
                continue

            if not external_accounts:
                logger.error("No external accounts from %s were connected", self.display_name)
                continue

            break
        return external_accounts

    def brownie_disconnect_accounts(self) -> None:
        """
        Disconnect from the External Provider.

        Removes all of this provider's `ExternalWeb3Account` objects from brownie's accounts container.
        """
        self._accounts = [i for i in brownie_accounts._accounts if getattr(i, "_external_web3_provider") != self]

    def make_request(self, *args, **kwargs) -> Any:
        return self._inner.make_request(*args, **kwargs)


class ExternalWeb3Account(_PrivateKeyAccount):
    """
    Class for interacting with an Ethereum account where signing is handled by an external web3 provider.

    If brownie is connected to a forked network, an unlocked account on the forked rpc is used for tx signatures.
    """

    def __init__(self, address: str, provider: ExternalWeb3Provider) -> None:
        self._external_web3_provider = provider

        super().__init__(address)

    def __repr__(self) -> str:
        display_name = self._external_web3_provider.display_name
        return f"<{display_name}Provider '{self.address}'>"

    def sign_defunct_message(self, message: str) -> SignedMessage:
        """Signs an `EIP-191` using this account's private key."""
        return self.sign_eip191_message(text=message)

    def sign_eip191_message(
        self,
        primitive: Optional[Primitives] = None,
        hexstr: Optional[HexStr] = None,
        text: Optional[str] = None,
    ) -> SignedMessage:
        """Signs an `EIP-191` using this account's private key.

        Args:
            message: An text

        Returns:
            An eth_account `SignedMessage` instance.
        """
        signable = encode_defunct(primitive=primitive, hexstr=hexstr, text=text)
        messageHash = HexBytes(_hash_eip191_message(signable))

        response = self._external_web3_provider.make_request("personal_sign", [signable, self.address])

        if "error" in response:
            raise ValueError(response["error"]["message"])

        signature = HexBytes(response["result"])

        (v, r, s) = Signature(signature_bytes=signature).vrs

        return SignedMessage(
            messageHash=messageHash,
            r=r,
            s=s,
            v=v,
            signature=signature,
        )

    def sign_message(self, message: EIP712Message) -> SignedMessage:
        """Signs an `EIP712Message` using this account's private key.

        Args:
            message: An `EIP712Message` instance.

        Returns:
            An eth_account `SignedMessage` instance.
        """
        response = self._external_web3_provider.make_request("eth_signTypedData_v4", [self.address, message.body_data()])

        if "error" in response:
            raise ValueError(response["error"]["message"])

        signature = HexBytes(response["result"])

        (v, r, s) = Signature(signature_bytes=signature).vrs

        return SignedMessage(
            messageHash=HexBytes(message.body()),
            r=r,
            s=s,
            v=v,
            signature=signature,
        )

    def _transact(self, tx: Dict, allow_revert: bool) -> None:
        if allow_revert is None:
            allow_revert = CONFIG.network_type == "development"

        if not allow_revert:
            self._check_for_revert(tx)

        formatters = {
            "nonce": brownie_web3.toHex,
            "value": brownie_web3.toHex,
            "chainId": brownie_web3.toHex,
            "data": brownie_web3.toHex,
            "from": to_address,
        }
        if "to" in tx:
            formatters["to"] = to_address

        tx["chainId"] = brownie_web3.chain_id
        tx = apply_formatters_to_dict(formatters, tx)

        response = self._external_web3_provider.make_request("personal_signTransaction", [tx])
        if "error" in response:
            raise ValueError(response["error"]["message"])

        # TODO: use brownie or frame's rpc? Maybe let the user decide with a kwarg on FrameProvider?
        return brownie_web3.eth.send_raw_transaction(response["result"]["raw"])


class FrameProvider(ExternalWeb3Provider):
    """
    Web3 provider connected to a local [Frame](https://frame.sh/).

    If run from inside a docker desktop container, the docker host's Frame will be used.
    """

    def __init__(
        self,
        origin: Optional[str] = None,
        websocket_kwargs: Optional[Any] = None,
        websocket_timeout: int = DEFAULT_WEBSOCKET_TIMEOUT,
    ) -> None:
        host = docker_or_localhost()

        if websocket_kwargs is None:
            websocket_kwargs = {}

        if origin:
            websocket_kwargs["origin"] = origin

        ws_provider = WebsocketProvider(
            endpoint_uri=f"ws://{host}:1248",
            websocket_kwargs=websocket_kwargs,
            websocket_timeout=websocket_timeout,
        )

        super().__init__("Frame", ws_provider)