[tests] decode public key get from the ledger

test/test_instructions.py

@@ -27,7 +27,7 @@
 from ragger.backend import BackendInterface
 from ragger.firmware import Firmware
 from utils.client import TezosClient, Version, Hwm, StatusCode
-from utils.account import Account
+from utils.account import Account, PublicKey
 from utils.helper import get_current_commit
 from utils.message import (
     Message,
@@ -554,9 +554,12 @@ def test_authorize_baking(account: Account, tezos_navigator: TezosNavigator) ->
     """Test the AUTHORIZE_BAKING instruction."""
     snap_path = Path(f"{account}")
 
-    public_key = tezos_navigator.authorize_baking(account, snap_path=snap_path)
+    data = tezos_navigator.authorize_baking(account, snap_path=snap_path)
 
-    account.check_public_key(public_key)
+    public_key = PublicKey.from_bytes(data, account.sig_scheme)
+
+    assert account.public_key == public_key, \
+        f"Expected public key {account.public_key} but got {public_key}"
 
     tezos_navigator.check_app_context(
         account,
@@ -633,9 +636,12 @@ def test_get_public_key_baking(account: Account, tezos_navigator: TezosNavigator
 
     tezos_navigator.authorize_baking(account)
 
-    public_key = tezos_navigator.authorize_baking(None, snap_path=Path(f"{account}"))
+    data = tezos_navigator.authorize_baking(None, snap_path=Path(f"{account}"))
+
+    public_key = PublicKey.from_bytes(data, account.sig_scheme)
 
-    account.check_public_key(public_key)
+    assert account.public_key == public_key, \
+        f"Expected public key {account.public_key} but got {public_key}"
 
 
 @pytest.mark.parametrize("account", ACCOUNTS)
@@ -644,7 +650,8 @@ def test_get_public_key_silent(account: Account, client: TezosClient) -> None:
 
     public_key = client.get_public_key_silent(account)
 
-    account.check_public_key(public_key)
+    assert account.public_key == public_key, \
+        f"Expected public key {account.public_key} but got {public_key}"
 
 
 @pytest.mark.parametrize("account", ACCOUNTS)
@@ -653,7 +660,8 @@ def test_get_public_key_prompt(account: Account, tezos_navigator: TezosNavigator
 
     public_key = tezos_navigator.get_public_key_prompt(account, snap_path=Path(f"{account}"))
 
-    account.check_public_key(public_key)
+    assert account.public_key == public_key, \
+        f"Expected public key {account.public_key} but got {public_key}"
 
 
 def test_reset_app_context(tezos_navigator: TezosNavigator) -> None:
@@ -688,7 +696,8 @@ def test_setup_app_context(account: Account, tezos_navigator: TezosNavigator) ->
         snap_path=snap_path
     )
 
-    account.check_public_key(public_key)
+    assert account.public_key == public_key, \
+        f"Expected public key {account.public_key} but got {public_key}"
 
     tezos_navigator.check_app_context(
         account,

test/utils/account.py

@@ -24,6 +24,7 @@
 import fastecdsa
 
 import pytezos
+from pytezos.crypto.encoding import base58_encode
 from bip_utils.bip.bip32.bip32_path import Bip32Path, Bip32PathParser
 from bip_utils.bip.bip32.bip32_key_data import Bip32KeyIndex
 from utils.helper import BytesReader
@@ -133,6 +134,56 @@ def from_bytes(cls, data: bytes, sig_scheme: SigScheme) -> 'Signature':
             return Signature(data)
         return Signature.from_tlv(data)
 
+class PublicKey:
+    """Set of functions over public key management"""
+
+    class CompressionKind(IntEnum):
+        """Bytes compression kind"""
+        EVEN         = 0x02
+        ODD          = 0x03
+        UNCOMPRESSED = 0x04
+
+        def __bytes__(self) -> bytes:
+            return bytes([self])
+
+    @staticmethod
+    def from_bytes(data: bytes, sig_scheme: SigScheme) -> str:
+        """Convert a public key from bytes to string"""
+        # `data` should be:
+        # kind + pk
+        # pk length = 32 for compressed, 64 for uncompressed
+        kind = data[0]
+        data = data[1:]
+
+        # Ed25519
+        if sig_scheme in [
+                SigScheme.ED25519,
+                SigScheme.BIP32_ED25519
+        ]:
+            assert kind == PublicKey.CompressionKind.EVEN, \
+                f"Wrong Ed25519 public key compression kind: {kind}"
+            assert len(data) == 32, \
+                f"Wrong Ed25519 public key length: {len(data)}"
+            return base58_encode(data, b'edpk').decode()
+
+        # Secp256
+        if sig_scheme in [
+                SigScheme.SECP256K1,
+                SigScheme.SECP256R1
+        ]:
+            assert kind == PublicKey.CompressionKind.UNCOMPRESSED, \
+                f"Wrong Secp256 public key compression kind: {kind}"
+            assert len(data) == 2 * 32, \
+                f"Wrong Secp256 public key length: {len(data)}"
+            kind = PublicKey.CompressionKind.ODD if data[-1] & 1 else \
+                PublicKey.CompressionKind.EVEN
+            prefix = b'sppk' if sig_scheme == SigScheme.SECP256K1 \
+                else b'p2pk'
+            data = bytes(kind) + data[:32]
+            return base58_encode(data, prefix).decode()
+
+        assert False, f"Wrong signature type: {sig_scheme}"
+
 class Account:
     """Class representing account."""
 
@@ -198,62 +249,6 @@ def sign_prehashed_message(self, prehashed_message: bytes) -> bytes:
             return r.to_bytes(32, 'big') + s.to_bytes(32, 'big')
         raise ValueError(f"Account do not have a right signature type: {self.sig_scheme}")
 
-    @property
-    def base58_decoded(self) -> bytes:
-        """base58_decoded of the account."""
-
-        # Get the public_key without prefix
-        public_key = base58.b58decode_check(self.public_key)
-
-        if self.sig_scheme in [
-                SigScheme.ED25519,
-                SigScheme.BIP32_ED25519
-        ]:
-            prefix = bytes.fromhex("0d0f25d9") # edpk(54)
-        elif self.sig_scheme == SigScheme.SECP256K1:
-            prefix = bytes.fromhex("03fee256") # sppk(55)
-        elif self.sig_scheme == SigScheme.SECP256R1:
-            prefix = bytes.fromhex("03b28b7f") # p2pk(55)
-        else:
-            raise ValueError(f"Account do not have a right signature type: {self.sig_scheme}")
-        assert public_key.startswith(prefix), \
-            "Expected prefix {prefix.hex()} but got {public_key.hex()}"
-
-        public_key = public_key[len(prefix):]
-
-        if self.sig_scheme in [
-                SigScheme.SECP256K1,
-                SigScheme.SECP256R1
-        ]:
-            assert public_key[0] in [0x02, 0x03], \
-                "Expected a prefix kind of 0x02 or 0x03 but got {public_key[0]}"
-            public_key = public_key[1:]
-
-        return public_key
-
-    def check_public_key(self, data: bytes) -> None:
-        """Check that the data correspond to the account."""
-
-        # `data` should be:
-        # length + kind + pk
-        # kind : 02=odd, 03=even, 04=uncompressed
-        # pk length = 32 for compressed, 64 for uncompressed
-        assert len(data) - 1 == data[0], \
-                "Expected a length of {data[0]} but got {len(data) - 1}"
-        if data[1] == 0x04: # public key uncompressed
-            assert data[0] == 1 + 32 + 32, \
-                "Expected a length of 1 + 32 + 32 but got {data[0]}"
-        elif data[1] in [0x02, 0x03]: # public key even or odd (compressed)
-            assert data[0] == 1 + 32, \
-                "Expected a length of 1 + 32 but got {data[0]}"
-        else:
-            raise ValueError(f"Expected a prefix kind of 0x02, 0x03 or 0x04 but got {data[1]}")
-        data = data[2:2+32]
-
-        public_key = self.base58_decoded
-        assert data == public_key, \
-            f"Expected public key {public_key.hex()} but got {data.hex()}"
-
     def check_signature(self,
                         signature: Union[bytes, Signature],
                         message: Union[str, bytes]):

test/utils/client.py

@@ -23,7 +23,7 @@
 from ragger.backend import BackendInterface
 from ragger.error import ExceptionRAPDU
 from pytezos.michelson import forge
-from utils.account import Account, SigScheme, BipPath, Signature
+from utils.account import Account, BipPath, PublicKey, Signature, SigScheme
 from utils.helper import BytesReader
 from utils.message import Message
 
@@ -237,11 +237,16 @@ def authorize_baking(self, account: Optional[Account]) -> bytes:
             sig_scheme=account.sig_scheme
             payload=bytes(account.path)
 
-        return self._exchange(
+        data = self._exchange(
             ins=Ins.AUTHORIZE_BAKING,
             sig_scheme=sig_scheme,
             payload=payload)
 
+        length, data = data[0], data[1:]
+        assert length == len(data), f"Wrong data size, {length} != {len(data)}"
+
+        return data
+
     def deauthorize(self) -> None:
         """Send the DEAUTHORIZE instruction."""
         data = self._exchange(ins=Ins.DEAUTHORIZE)
@@ -259,20 +264,30 @@ def get_auth_key_with_curve(self) -> Tuple[SigScheme, BipPath]:
         sig_scheme = SigScheme(data[0])
         return sig_scheme, BipPath.from_bytes(data[1:])
 
-    def get_public_key_silent(self, account: Account) -> bytes:
+    def get_public_key_silent(self, account: Account) -> str:
         """Send the GET_PUBLIC_KEY instruction."""
-        return self._exchange(
+        data = self._exchange(
             ins=Ins.GET_PUBLIC_KEY,
             sig_scheme=account.sig_scheme,
             payload=bytes(account.path))
 
-    def get_public_key_prompt(self, account: Account) -> bytes:
+        length, data = data[0], data[1:]
+        assert length == len(data), f"Wrong data size, {length} != {len(data)}"
+
+        return PublicKey.from_bytes(data, account.sig_scheme)
+
+    def get_public_key_prompt(self, account: Account) -> str:
         """Send the PROMPT_PUBLIC_KEY instruction."""
-        return  self._exchange(
+        data = self._exchange(
             ins=Ins.PROMPT_PUBLIC_KEY,
             sig_scheme=account.sig_scheme,
             payload=bytes(account.path))
 
+        length, data = data[0], data[1:]
+        assert length == len(data), f"Wrong data size, {length} != {len(data)}"
+
+        return PublicKey.from_bytes(data, account.sig_scheme)
+
     def reset_app_context(self, reset_level: int) -> None:
         """Send the RESET instruction."""
         reset_level_raw = reset_level.to_bytes(4, byteorder='big')
@@ -286,7 +301,7 @@ def setup_app_context(self,
                           account: Account,
                           main_chain_id: str,
                           main_hwm: Hwm,
-                          test_hwm: Hwm) -> bytes:
+                          test_hwm: Hwm) -> str:
         """Send the SETUP instruction."""
 
         data: bytes = b''
@@ -295,11 +310,16 @@ def setup_app_context(self,
         data += bytes(test_hwm)
         data += bytes(account.path)
 
-        return self._exchange(
+        data = self._exchange(
             ins=Ins.SETUP,
             sig_scheme=account.sig_scheme,
             payload=data)
 
+        length, data = data[0], data[1:]
+        assert length == len(data), f"Wrong data size, {length} != {len(data)}"
+
+        return PublicKey.from_bytes(data, account.sig_scheme)
+
     def get_main_hwm(self) -> Hwm:
         """Send the QUERY_MAIN_HWM instruction."""
         return Hwm.from_bytes(self._exchange(ins=Ins.QUERY_MAIN_HWM))

test/utils/navigator.py

@@ -406,7 +406,7 @@ def authorize_baking(self,
     def get_public_key_prompt(self,
                               account: Account,
                               navigate: Optional[Callable] = None,
-                              **kwargs) -> bytes:
+                              **kwargs) -> str:
         """Send a get public key request and navigate until accept"""
         if navigate is None:
             navigate = self.accept_key_navigate
@@ -473,7 +473,7 @@ def setup_app_context(self,
                           main_hwm: Hwm,
                           test_hwm: Hwm,
                           navigate: Optional[Callable] = None,
-                          **kwargs) -> bytes:
+                          **kwargs) -> str:
         """Send a setup request and navigate until accept"""
         if navigate is None:
             navigate = self.accept_setup_navigate