refactor: public key class (#169)

alexbarnsley · web-flow · commit c8181aec0bdd · 2025-04-23T11:28:36.000+02:00
diff --git a/crypto/identity/public_key.py b/crypto/identity/public_key.py
@@ -1,19 +1,25 @@
-from binascii import hexlify, unhexlify
-
 from coincurve import PublicKey as PubKey
 
 from crypto.identity.private_key import PrivateKey
 
 class PublicKey(object):
+    public_key: str
+
     def __init__(self, public_key: str):
-        self.public_key = PubKey(unhexlify(public_key.encode()))
+        self.public_key = public_key
 
     @classmethod
-    def from_passphrase(cls, passphrase: str) -> str:
+    def from_passphrase(cls, passphrase: str) -> 'PublicKey':
         private_key = PrivateKey.from_passphrase(passphrase)
 
-        return private_key.public_key
+        return cls(private_key.public_key)
 
     @classmethod
-    def from_hex(cls, public_key):
+    def from_hex(cls, public_key) -> 'PublicKey':
         return cls(public_key)
+
+    @classmethod
+    def recover(cls, message: bytes, signature: bytes) -> 'PublicKey':
+        recovered_public_key = PubKey.from_signature_and_message(signature, message, hasher=None)
+
+        return cls(recovered_public_key.format().hex())
diff --git a/crypto/transactions/builder/abstract_transaction_builder.py b/crypto/transactions/builder/abstract_transaction_builder.py
diff --git a/crypto/transactions/builder/evm_call_builder.py b/crypto/transactions/builder/evm_call_builder.py
@@ -1,4 +1,4 @@
-from crypto.transactions.builder.base import AbstractTransactionBuilder
+from crypto.transactions.builder.abstract_transaction_builder import AbstractTransactionBuilder
 from crypto.transactions.types.evm_call import EvmCall
 from crypto.utils.transaction_utils import TransactionUtils
 
diff --git a/crypto/transactions/builder/multipayment_builder.py b/crypto/transactions/builder/multipayment_builder.py
@@ -1,5 +1,5 @@
 from crypto.enums.contract_addresses import ContractAddresses
-from crypto.transactions.builder.base import AbstractTransactionBuilder
+from crypto.transactions.builder.abstract_transaction_builder import AbstractTransactionBuilder
 from crypto.transactions.types.multipayment import Multipayment
 
 class MultipaymentBuilder(AbstractTransactionBuilder):
diff --git a/crypto/transactions/builder/transfer_builder.py b/crypto/transactions/builder/transfer_builder.py
@@ -1,4 +1,4 @@
-from crypto.transactions.builder.base import AbstractTransactionBuilder
+from crypto.transactions.builder.abstract_transaction_builder import AbstractTransactionBuilder
 from crypto.transactions.types.transfer import Transfer
 
 
diff --git a/crypto/transactions/builder/unvote_builder.py b/crypto/transactions/builder/unvote_builder.py
@@ -1,5 +1,5 @@
 from crypto.enums.contract_addresses import ContractAddresses
-from crypto.transactions.builder.base import AbstractTransactionBuilder
+from crypto.transactions.builder.abstract_transaction_builder import AbstractTransactionBuilder
 from crypto.transactions.types.unvote import Unvote
 
 class UnvoteBuilder(AbstractTransactionBuilder):
diff --git a/crypto/transactions/builder/username_registration_builder.py b/crypto/transactions/builder/username_registration_builder.py
@@ -1,7 +1,7 @@
 import re
 from crypto.enums.contract_addresses import ContractAddresses
 from crypto.exceptions import InvalidUsernameException
-from crypto.transactions.builder.base import AbstractTransactionBuilder
+from crypto.transactions.builder.abstract_transaction_builder import AbstractTransactionBuilder
 from crypto.transactions.types.username_registration import UsernameRegistration
 
 class UsernameRegistrationBuilder(AbstractTransactionBuilder):
diff --git a/crypto/transactions/builder/username_resignation_builder.py b/crypto/transactions/builder/username_resignation_builder.py
@@ -1,5 +1,5 @@
 from crypto.enums.contract_addresses import ContractAddresses
-from crypto.transactions.builder.base import AbstractTransactionBuilder
+from crypto.transactions.builder.abstract_transaction_builder import AbstractTransactionBuilder
 from crypto.transactions.types.username_resignation import UsernameResignation
 
 class UsernameResignationBuilder(AbstractTransactionBuilder):
diff --git a/crypto/transactions/builder/validator_registration_builder.py b/crypto/transactions/builder/validator_registration_builder.py
@@ -1,5 +1,5 @@
 from crypto.enums.contract_addresses import ContractAddresses
-from crypto.transactions.builder.base import AbstractTransactionBuilder
+from crypto.transactions.builder.abstract_transaction_builder import AbstractTransactionBuilder
 from crypto.transactions.types.validator_registration import ValidatorRegistration
 
 class ValidatorRegistrationBuilder(AbstractTransactionBuilder):
diff --git a/crypto/transactions/builder/validator_resignation_builder.py b/crypto/transactions/builder/validator_resignation_builder.py
@@ -1,5 +1,5 @@
 from crypto.enums.contract_addresses import ContractAddresses
-from crypto.transactions.builder.base import AbstractTransactionBuilder
+from crypto.transactions.builder.abstract_transaction_builder import AbstractTransactionBuilder
 from crypto.transactions.types.validator_resignation import ValidatorResignation
 
 class ValidatorResignationBuilder(AbstractTransactionBuilder):
diff --git a/crypto/transactions/builder/vote_builder.py b/crypto/transactions/builder/vote_builder.py
@@ -1,5 +1,5 @@
 from crypto.enums.contract_addresses import ContractAddresses
-from crypto.transactions.builder.base import AbstractTransactionBuilder
+from crypto.transactions.builder.abstract_transaction_builder import AbstractTransactionBuilder
 from crypto.transactions.types.vote import Vote
 
 class VoteBuilder(AbstractTransactionBuilder):
diff --git a/crypto/transactions/types/abstract_transaction.py b/crypto/transactions/types/abstract_transaction.py
@@ -5,8 +5,8 @@
 from crypto.enums.contract_abi_type import ContractAbiType
 from crypto.identity.address import Address
 from crypto.identity.private_key import PrivateKey
+from crypto.identity.public_key import PublicKey
 from crypto.utils.transaction_utils import TransactionUtils
-from coincurve import PublicKey
 
 class AbstractTransaction:
     def __init__(self, data: dict):
@@ -35,19 +35,14 @@ def sign(self, private_key: PrivateKey):
 
         return self
 
-    def get_public_key(self, compact_signature, hash_):
-        public_key = PublicKey.from_signature_and_message(compact_signature, hash_, hasher=None)
-
-        return public_key
-
     def recover_sender(self):
         signature_with_recid = self._get_signature()
         if not signature_with_recid:
             return False
 
         hash_ = bytes.fromhex(self.hash(skip_signature=True))
-        public_key = self.get_public_key(signature_with_recid, hash_)
-        self.data['senderPublicKey'] = public_key.format().hex()
+        public_key = self.__recover_public_key(signature_with_recid, hash_)
+        self.data['senderPublicKey'] = public_key.public_key
         self.data['senderAddress'] = Address.from_public_key(self.data['senderPublicKey'])
 
     def verify(self) -> bool:
@@ -56,14 +51,12 @@ def verify(self) -> bool:
             return False
 
         hash_ = bytes.fromhex(self.hash(skip_signature=True))
-        recovered_public_key = self.get_public_key(signature_with_recid, hash_)
+        recovered_public_key = self.__recover_public_key(signature_with_recid, hash_)
         sender_public_key_hex = self.data.get('senderPublicKey')
         if not sender_public_key_hex:
             return False
 
-        sender_public_key_bytes = bytes.fromhex(sender_public_key_hex)
-
-        return recovered_public_key.format() == sender_public_key_bytes
+        return recovered_public_key.public_key == sender_public_key_hex
 
     def serialize(self, skip_signature: bool = False) -> bytes:
         from crypto.transactions.serializer import Serializer
@@ -94,3 +87,6 @@ def _decode_payload(data: dict, abi_type: ContractAbiType = ContractAbiType.CONS
         from crypto.transactions.deserializer import Deserializer
 
         return Deserializer.decode_payload(data, abi_type)
+
+    def __recover_public_key(self, compact_signature, hash_):
+        return PublicKey.recover(hash_, compact_signature)
diff --git a/crypto/utils/message.py b/crypto/utils/message.py
@@ -3,9 +3,9 @@
 from typing import Union
 
 from Cryptodome.Hash import keccak
-from coincurve import PublicKey
 
 from crypto.identity.private_key import PrivateKey
+from crypto.identity.public_key import PublicKey
 
 class Message(object):
     public_key: bytes
@@ -79,9 +79,9 @@ def verify(self):
 
         signature = signature_r + signature_s + bytes([signature_v - 27])
 
-        public_key = PublicKey.from_signature_and_message(signature, message_hash, hasher=None)
+        public_key = PublicKey.recover(message_hash, signature)
 
-        return public_key.format() == unhexlify(self.public_key)
+        return public_key.public_key == unhexlify(self.public_key).hex()
 
     def to_dict(self):
         """Return a dictionary of the message
diff --git a/tests/identity/test_public_key.py b/tests/identity/test_public_key.py
@@ -1,16 +1,13 @@
-from binascii import hexlify
 from crypto.identity.public_key import PublicKey
 
 
 def test_public_key_from_passphrase(identity):
     public_key = PublicKey.from_passphrase(identity['passphrase'])
-    assert public_key == identity['data']['public_key']
+    assert public_key.public_key == identity['data']['public_key']
 
 
 def test_public_key_from_hex(identity):
     public_key = PublicKey.from_hex(identity['data']['public_key'])
     assert isinstance(public_key, PublicKey)
 
-    public_key_hex = hexlify(public_key.public_key.format()).decode()
-
-    assert public_key_hex == identity['data']['public_key']
+    assert public_key.public_key == identity['data']['public_key']

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-from crypto.transactions.builder.base import AbstractTransactionBuilder`
	`1`	`+from crypto.transactions.builder.abstract_transaction_builder import AbstractTransactionBuilder`
`2`	`2`	`from crypto.transactions.types.evm_call import EvmCall`
`3`	`3`	`from crypto.utils.transaction_utils import TransactionUtils`
`4`	`4`