fix: forgot to update removal too

ApeWorX · Jun 5, 2024 · 00b4bfa · 00b4bfa
1 parent 7d8cafd
commit 00b4bfa
Showing 1 changed file with 5 additions and 14 deletions.
diff --git a/ape_safe/client/__init__.py b/ape_safe/client/__init__.py
@@ -204,7 +204,7 @@ def add_delegate(self, delegate: AddressType, label: str, delegator: AccountAPI)
 
         # NOTE: This is required as Safe API uses an antiquated .signHash method
         if not (sig := delegator.sign_raw_msghash(msg_hash)):
-            raise ActionNotPerformedError("Did not sign delegation")
+            raise ActionNotPerformedError("Did not sign delegate approval")
 
         payload = {
             "safe": self.address,
@@ -216,24 +216,15 @@ def add_delegate(self, delegate: AddressType, label: str, delegator: AccountAPI)
         self._post("delegates", json=payload)
 
     def remove_delegate(self, delegate: AddressType, 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 delegate removal")
 
         payload = {
             "delegator": delegator.address,
-            "signature": sig.signature.hex(),
+            "signature": sig.encode_rsv().hex(),
         }
         self._delete(f"delegates/{delegate}", json=payload)