From 7d8cafd5d926c1d34523b4ff89c3ddbde431d29a Mon Sep 17 00:00:00 2001
From: doggie <3859395+fubuloubu@users.noreply.github.com>
Date: Wed, 5 Jun 2024 18:32:47 -0400
Subject: [PATCH] refactor: use new `AccountAPI.sign_raw_msghash` method

---
 ape_safe/client/__init__.py | 17 ++++-------------
 ape_safe/client/base.py     |  6 +++---
 2 files changed, 7 insertions(+), 16 deletions(-)

diff --git a/ape_safe/client/__init__.py b/ape_safe/client/__init__.py
index c693191..5c09201 100644
--- a/ape_safe/client/__init__.py
+++ b/ape_safe/client/__init__.py
@@ -200,27 +200,18 @@ def get_delegates(self) -> Dict[AddressType, List[AddressType]]:
         return delegates
 
     def add_delegate(self, delegate: AddressType, label: str, delegator: AccountAPI):
-        # TODO: Replace this by adding raw hash signing into supported account plugins
-        #       See: https://github.com/ApeWorX/ape/issues/1962
-        if not isinstance(delegator, KeyfileAccount):
-            raise ActionNotPerformedError("Need access to private key for this method.")
-
-        logger.warning("Need to unlock account to add a delegate.")
-        delegator.unlock()  # NOTE: Ensures we have the key handy
-
         msg_hash = self.create_delegate_message(delegate)
+
         # NOTE: This is required as Safe API uses an antiquated .signHash method
-        sig = EthAccount.signHash(
-            msg_hash,
-            delegator._KeyfileAccount__cached_key,  # type: ignore[attr-defined]
-        )
+        if not (sig := delegator.sign_raw_msghash(msg_hash)):
+            raise ActionNotPerformedError("Did not sign delegation")
 
         payload = {
             "safe": self.address,
             "delegate": delegate,
             "delegator": delegator.address,
             "label": label,
-            "signature": sig.signature.hex(),
+            "signature": sig.encode_rsv().hex(),
         }
         self._post("delegates", json=payload)
 
diff --git a/ape_safe/client/base.py b/ape_safe/client/base.py
index 2a7bbe1..83ecfd3 100644
--- a/ape_safe/client/base.py
+++ b/ape_safe/client/base.py
@@ -7,7 +7,7 @@
 import requests
 import urllib3
 from ape.api import AccountAPI
-from ape.types import AddressType, MessageSignature
+from ape.types import AddressType, HexBytes, MessageSignature
 from eth_utils import keccak
 from requests import Response
 from requests.adapters import HTTPAdapter
@@ -113,11 +113,11 @@ def get_transactions(
 
             yield txn
 
-    def create_delegate_message(self, delegate: AddressType) -> bytes:
+    def create_delegate_message(self, delegate: AddressType) -> HexBytes:
         # NOTE: referencing https://github.com/safe-global/safe-eth-py/blob/
         # a0a5771622f143ee6301cfc381c5ed50832ff482/gnosis/safe/api/transaction_service_api.py#L34
         totp = int(time.time()) // 3600
-        return keccak(text=(delegate + str(totp)))
+        return HexBytes(keccak(text=(delegate + str(totp))))
 
     @abstractmethod
     def get_delegates(self) -> dict[AddressType, list[AddressType]]: ...