refactor!: accept any msg type via AccountAPI.sign_message [APE-130…

…5] (#1614) * feat: implement new sign_message fn accepting Any * fix: add missing return statement * fix: Update src/ape/api/accounts.py Co-authored-by: El De-dog-lo <[email protected]> * fix: add comment explaining instance check/casting * feat: support EIP712Message * feat: Return None for unsupported type, log warning * fix: print message before signing (string) * fix: lint * feat: better display of signable message * fix: decode salt bytes before printing * fix: add test signing string * feat: add support for signing ints directly * fix: check for EIP712Message can be elif * fix: add check for SignableMessage * fix: display salt as hex * fix: address pr feedback better handling of cases for different types in sign_message * fix: shouldn't be an f string * chore: isort * fix: keyfile accounts now properly handling signablemessage --------- Co-authored-by: El De-dog-lo <[email protected]> Co-authored-by: Juliya Smith <[email protected]>
ApeWorX · Dec 8, 2023 · d033f85 · d033f85
1 parent 271a272
commit d033f85
Show file tree

Hide file tree

Showing 4 changed files with 104 additions and 19 deletions.
diff --git a/src/ape/api/accounts.py b/src/ape/api/accounts.py
@@ -1,9 +1,12 @@
 from pathlib import Path
-from typing import TYPE_CHECKING, Iterator, List, Optional, Type, Union
+from typing import TYPE_CHECKING, Any, Iterator, List, Optional, Type, Union
 
 import click
+from eip712.messages import EIP712Message
 from eip712.messages import SignableMessage as EIP712SignableMessage
 from eth_account import Account
+from eth_account.messages import encode_defunct
+from hexbytes import HexBytes
 
 from ape.api.address import BaseAddress
 from ape.api.transactions import ReceiptAPI, TransactionAPI
@@ -54,18 +57,21 @@ def alias(self) -> Optional[str]:
         return None
 
     @abstractmethod
-    def sign_message(self, msg: SignableMessage) -> Optional[MessageSignature]:
+    def sign_message(self, msg: Any, **signer_options) -> Optional[MessageSignature]:
         """
         Sign a message.
 
         Args:
-          msg (:class:`~ape.types.signatures.SignableMessage`): The message to sign.
+          msg (Any): The message to sign. Account plugins can handle various types of messages.
+            For example, :class:`~ape_accounts.accouns.KeyfileAccount` can handle
+            :class:`~ape.types.signatures.SignableMessage`, str, int, and bytes.
             See these
             `docs <https://eth-account.readthedocs.io/en/stable/eth_account.html#eth_account.messages.SignableMessage>`__  # noqa: E501
-            for more type information on this type.
+            for more type information on the ``SignableMessage`` type.
+          **signer_options: Additional kwargs given to the signer to modify the signing operation.
 
         Returns:
-          :class:`~ape.types.signatures.MessageSignature` (optional): The signed message.
+          :class:`~ape.types.signatures.MessageSignature` (optional): The signature corresponding to the message.
         """
 
     @abstractmethod
@@ -258,7 +264,7 @@ def declare(self, contract: "ContractContainer", *args, **kwargs) -> ReceiptAPI:
 
     def check_signature(
         self,
-        data: Union[SignableMessage, TransactionAPI],
+        data: Union[SignableMessage, TransactionAPI, str, EIP712Message, int],
         signature: Optional[MessageSignature] = None,  # TransactionAPI doesn't need it
     ) -> bool:
         """
@@ -273,6 +279,12 @@ def check_signature(
         Returns:
             bool: ``True`` if the data was signed by this account. ``False`` otherwise.
         """
+        if isinstance(data, str):
+            data = encode_defunct(text=data)
+        elif isinstance(data, int):
+            data = encode_defunct(hexstr=HexBytes(data).hex())
+        elif isinstance(data, EIP712Message):
+            data = data.signable_message
         if isinstance(data, (SignableMessage, EIP712SignableMessage)):
             if signature:
                 return self.address == Account.recover_message(data, vrs=signature)
@@ -494,7 +506,7 @@ class ImpersonatedAccount(AccountAPI):
     def address(self) -> AddressType:
         return self.raw_address
 
-    def sign_message(self, msg: SignableMessage) -> Optional[MessageSignature]:
+    def sign_message(self, msg: Any, **signer_options) -> Optional[MessageSignature]:
         raise NotImplementedError("This account cannot sign messages")
 
     def sign_transaction(self, txn: TransactionAPI, **kwargs) -> Optional[TransactionAPI]:

diff --git a/src/ape_accounts/accounts.py b/src/ape_accounts/accounts.py
@@ -1,10 +1,12 @@
 import json
 from os import environ
 from pathlib import Path
-from typing import Dict, Iterator, Optional
+from typing import Any, Dict, Iterator, Optional
 
 import click
+from eip712.messages import EIP712Message
 from eth_account import Account as EthAccount
+from eth_account.messages import encode_defunct
 from eth_keys import keys  # type: ignore
 from eth_utils import to_bytes
 from ethpm_types import HexBytes
@@ -155,8 +157,50 @@ def delete(self):
         self.__decrypt_keyfile(passphrase)
         self.keyfile_path.unlink()
 
-    def sign_message(self, msg: SignableMessage) -> Optional[MessageSignature]:
-        user_approves = self.__autosign or click.confirm(f"{msg}\n\nSign: ")
+    def sign_message(self, msg: Any, **signer_options) -> Optional[MessageSignature]:
+        user_approves = False
+
+        if isinstance(msg, str):
+            user_approves = self.__autosign or click.confirm(f"Message: {msg}\n\nSign: ")
+            msg = encode_defunct(text=msg)
+        elif isinstance(msg, int):
+            user_approves = self.__autosign or click.confirm(f"Message: {msg}\n\nSign: ")
+            msg = encode_defunct(hexstr=HexBytes(msg).hex())
+        elif isinstance(msg, bytes):
+            user_approves = self.__autosign or click.confirm(f"Message: {msg.hex()}\n\nSign: ")
+            msg = encode_defunct(primitive=msg)
+        elif isinstance(msg, EIP712Message):
+            # Display message data to user
+            display_msg = "Signing EIP712 Message\n"
+
+            # Domain Data
+            display_msg += "Domain\n"
+            if msg._name_:
+                display_msg += f"\tName: {msg._name_}\n"
+            if msg._version_:
+                display_msg += f"\tVersion: {msg._version_}\n"
+            if msg._chainId_:
+                display_msg += f"\tChain ID: {msg._chainId_}\n"
+            if msg._verifyingContract_:
+                display_msg += f"\tContract: {msg._verifyingContract_}\n"
+            if msg._salt_:
+                display_msg += f"\tSalt: 0x{msg._salt_.hex()}\n"
+
+            # Message Data
+            display_msg += "Message\n"
+            for field, value in msg._body_["message"].items():
+                display_msg += f"\t{field}: {value}\n"
+
+            user_approves = self.__autosign or click.confirm(f"{display_msg}\nSign: ")
+
+            # Convert EIP712Message to SignableMessage for handling below
+            msg = msg.signable_message
+        elif isinstance(msg, SignableMessage):
+            user_approves = self.__autosign or click.confirm(f"{msg}\n\nSign: ")
+        else:
+            logger.warning("Unsupported message type, (type=%r, msg=%r)", type(msg), msg)
+            return None
+
         if not user_approves:
             return None
 

diff --git a/src/ape_test/accounts.py b/src/ape_test/accounts.py
@@ -1,8 +1,9 @@
-from typing import Iterator, List, Optional
+from typing import Any, Iterator, List, Optional
 
 from eth_account import Account as EthAccount
-from eth_account.messages import SignableMessage
+from eth_account.messages import SignableMessage, encode_defunct
 from eth_utils import to_bytes
+from hexbytes import HexBytes
 
 from ape.api import TestAccountAPI, TestAccountContainerAPI, TransactionAPI
 from ape.types import AddressType, MessageSignature, TransactionSignature
@@ -101,13 +102,23 @@ def alias(self) -> str:
     def address(self) -> AddressType:
         return self.network_manager.ethereum.decode_address(self.address_str)
 
-    def sign_message(self, msg: SignableMessage) -> Optional[MessageSignature]:
-        signed_msg = EthAccount.sign_message(msg, self.private_key)
-        return MessageSignature(
-            v=signed_msg.v,
-            r=to_bytes(signed_msg.r),
-            s=to_bytes(signed_msg.s),
-        )
+    def sign_message(self, msg: Any, **signer_options) -> Optional[MessageSignature]:
+        # Convert str and int to SignableMessage if needed
+        if isinstance(msg, str):
+            msg = encode_defunct(text=msg)
+        elif isinstance(msg, int):
+            msg = HexBytes(msg).hex()
+            msg = encode_defunct(hexstr=msg)
+
+        # Process SignableMessage
+        if isinstance(msg, SignableMessage):
+            signed_msg = EthAccount.sign_message(msg, self.private_key)
+            return MessageSignature(
+                v=signed_msg.v,
+                r=to_bytes(signed_msg.r),
+                s=to_bytes(signed_msg.s),
+            )
+        return None
 
     def sign_transaction(self, txn: TransactionAPI, **kwargs) -> Optional[TransactionAPI]:
         # Signs anything that's given to it

diff --git a/tests/functional/test_accounts.py b/tests/functional/test_accounts.py
@@ -53,6 +53,24 @@ def test_sign_message(signer, message):
     assert signer.check_signature(message, signature)
 
 
+def test_sign_string(signer):
+    message = "Hello Apes!"
+    signature = signer.sign_message(message)
+    assert signer.check_signature(message, signature)
+
+
+def test_sign_int(signer):
+    message = 4
+    signature = signer.sign_message(message)
+    assert signer.check_signature(message, signature)
+
+
+def test_sign_message_unsupported_type_returns_none(signer):
+    message = 1234.123
+    signature = signer.sign_message(message)
+    assert signature is None
+
+
 def test_recover_signer(signer, message):
     signature = signer.sign_message(message)
     assert recover_signer(message, signature) == signer