From 8b252bff80e170226e7354856e03156299d958e5 Mon Sep 17 00:00:00 2001
From: Adrian Matejov <adrian.matejov@simplestaking.com>
Date: Tue, 11 Sep 2018 13:32:15 +0200
Subject: [PATCH] tezos: fixed nitpicks

Signed-off-by: Adrian Matejov <adrian.matejov@simplestaking.com>
---
 src/apps/tezos/get_address.py    |  12 ++--
 src/apps/tezos/get_public_key.py |   8 +--
 src/apps/tezos/helpers.py        |  13 +---
 src/apps/tezos/sign_tx.py        | 107 +++++++++++++++++--------------
 tests/test_apps.tezos.address.py |  19 +++---
 tests/test_apps.tezos.encode.py  |  85 +++++++++++++++---------
 6 files changed, 134 insertions(+), 110 deletions(-)

diff --git a/src/apps/tezos/get_address.py b/src/apps/tezos/get_address.py
index 927fec202b..77049fe419 100644
--- a/src/apps/tezos/get_address.py
+++ b/src/apps/tezos/get_address.py
@@ -1,20 +1,22 @@
 from trezor.crypto import hashlib
-from trezor.crypto.curve import ed25519
 from trezor.messages.TezosAddress import TezosAddress
 
 from apps.common import seed
 from apps.common.layout import show_address, show_qr
-from apps.tezos.helpers import TEZOS_CURVE, b58cencode, tezos_get_address_prefix
+from apps.tezos.helpers import (
+    TEZOS_CURVE,
+    TEZOS_ED25519_ADDRESS_PREFIX,
+    base58_encode_check,
+)
 
 
 async def get_address(ctx, msg):
     address_n = msg.address_n or ()
     node = await seed.derive_node(ctx, address_n, TEZOS_CURVE)
 
-    sk = node.private_key()
-    pk = ed25519.publickey(sk)
+    pk = seed.remove_ed25519_prefix(node.public_key())
     pkh = hashlib.blake2b(pk, outlen=20).digest()
-    address = b58cencode(pkh, prefix=tezos_get_address_prefix(0))
+    address = base58_encode_check(pkh, prefix=TEZOS_ED25519_ADDRESS_PREFIX)
 
     if msg.show_display:
         while True:
diff --git a/src/apps/tezos/get_public_key.py b/src/apps/tezos/get_public_key.py
index dd8759be82..964944ce01 100644
--- a/src/apps/tezos/get_public_key.py
+++ b/src/apps/tezos/get_public_key.py
@@ -1,5 +1,4 @@
 from trezor import ui
-from trezor.crypto.curve import ed25519
 from trezor.messages import ButtonRequestType
 from trezor.messages.TezosPublicKey import TezosPublicKey
 from trezor.ui.text import Text
@@ -7,16 +6,15 @@ from trezor.utils import chunks
 
 from apps.common import seed
 from apps.common.confirm import require_confirm
-from apps.tezos.helpers import TEZOS_CURVE, TEZOS_PUBLICKEY_PREFIX, b58cencode
+from apps.tezos.helpers import TEZOS_CURVE, TEZOS_PUBLICKEY_PREFIX, base58_encode_check
 
 
 async def get_public_key(ctx, msg):
     address_n = msg.address_n or ()
     node = await seed.derive_node(ctx, address_n, TEZOS_CURVE)
 
-    sk = node.private_key()
-    pk = ed25519.publickey(sk)
-    pk_prefixed = b58cencode(pk, prefix=TEZOS_PUBLICKEY_PREFIX)
+    pk = seed.remove_ed25519_prefix(node.public_key())
+    pk_prefixed = base58_encode_check(pk, prefix=TEZOS_PUBLICKEY_PREFIX)
 
     if msg.show_display:
         await _show_tezos_pubkey(ctx, pk_prefixed)
diff --git a/src/apps/tezos/helpers.py b/src/apps/tezos/helpers.py
index 8ad0cd2280..4a82a6aa1b 100644
--- a/src/apps/tezos/helpers.py
+++ b/src/apps/tezos/helpers.py
@@ -1,10 +1,10 @@
 from micropython import const
 
-from trezor import wire
 from trezor.crypto import base58
 
 TEZOS_CURVE = "ed25519"
 TEZOS_AMOUNT_DIVISIBILITY = const(6)
+TEZOS_ED25519_ADDRESS_PREFIX = "tz1"
 TEZOS_ORIGINATED_ADDRESS_PREFIX = "KT1"
 TEZOS_PUBLICKEY_PREFIX = "edpk"
 TEZOS_SIGNATURE_PREFIX = "edsig"
@@ -23,21 +23,14 @@ TEZOS_PREFIX_BYTES = {
 }
 
 
-def tezos_get_address_prefix(tag):
-    prefixes = ["tz1", "tz2", "tz3"]
-    if 0 <= tag < len(prefixes):
-        return prefixes[tag]
-    raise wire.DataError("Invalid tag in address hash")
-
-
-def b58cencode(payload, prefix=None):
+def base58_encode_check(payload, prefix=None):
     result = payload
     if prefix is not None:
         result = bytes(TEZOS_PREFIX_BYTES[prefix]) + payload
     return base58.encode_check(result)
 
 
-def b58cdecode(enc, prefix=None):
+def base58_decode_check(enc, prefix=None):
     decoded = base58.decode_check(enc)
     if prefix is not None:
         decoded = decoded[len(TEZOS_PREFIX_BYTES[prefix]) :]
diff --git a/src/apps/tezos/sign_tx.py b/src/apps/tezos/sign_tx.py
index b1f4c51d5c..134dffb033 100644
--- a/src/apps/tezos/sign_tx.py
+++ b/src/apps/tezos/sign_tx.py
@@ -1,5 +1,3 @@
-import ustruct
-
 from trezor import wire
 from trezor.crypto import hashlib
 from trezor.crypto.curve import ed25519
@@ -7,12 +5,12 @@ from trezor.messages import TezosContractType
 from trezor.messages.TezosSignedTx import TezosSignedTx
 
 from apps.common import seed
+from apps.common.writers import write_bytes, write_uint8
 from apps.tezos.helpers import (
     TEZOS_CURVE,
     TEZOS_ORIGINATED_ADDRESS_PREFIX,
     TEZOS_SIGNATURE_PREFIX,
-    b58cencode,
-    tezos_get_address_prefix,
+    base58_encode_check,
 )
 from apps.tezos.layout import *
 
@@ -50,8 +48,12 @@ async def sign_tx(ctx, msg):
     else:
         raise wire.DataError("Invalid operation")
 
-    opbytes = _get_operation_bytes(msg)
+    w = bytearray()
+    _get_operation_bytes(w, msg)
 
+    opbytes = bytes(w)
+
+    # watermark 0x03 is prefix for transactions, delegations, originations, reveals...
     watermark = bytes([3])
     wm_opbytes = watermark + opbytes
     wm_opbytes_hash = hashlib.blake2b(wm_opbytes, outlen=32).digest()
@@ -60,9 +62,9 @@ async def sign_tx(ctx, msg):
 
     sig_op_contents = opbytes + signature
     sig_op_contents_hash = hashlib.blake2b(sig_op_contents, outlen=32).digest()
-    ophash = b58cencode(sig_op_contents_hash, prefix="o")
+    ophash = base58_encode_check(sig_op_contents_hash, prefix="o")
 
-    sig_prefixed = b58cencode(signature, prefix=TEZOS_SIGNATURE_PREFIX)
+    sig_prefixed = base58_encode_check(signature, prefix=TEZOS_SIGNATURE_PREFIX)
 
     return TezosSignedTx(
         signature=sig_prefixed, sig_op_contents=sig_op_contents, operation_hash=ophash
@@ -70,8 +72,12 @@ async def sign_tx(ctx, msg):
 
 
 def _get_address_by_tag(address_hash):
+    prefixes = ["tz1", "tz2", "tz3"]
     tag = int(address_hash[0])
-    return b58cencode(address_hash[1:], prefix=tezos_get_address_prefix(tag))
+
+    if 0 <= tag < len(prefixes):
+        return base58_encode_check(address_hash[1:], prefix=prefixes[tag])
+    raise wire.DataError("Invalid tag in address hash")
 
 
 def _get_address_from_contract(address):
@@ -79,77 +85,80 @@ def _get_address_from_contract(address):
         return _get_address_by_tag(address.hash)
 
     elif address.tag == TezosContractType.Originated:
-        return b58cencode(address.hash[:-1], prefix=TEZOS_ORIGINATED_ADDRESS_PREFIX)
+        return base58_encode_check(
+            address.hash[:-1], prefix=TEZOS_ORIGINATED_ADDRESS_PREFIX
+        )
 
     raise wire.DataError("Invalid contract type")
 
 
-def _get_operation_bytes(msg):
-    result = msg.branch
+def _get_operation_bytes(w: bytearray, msg):
+    write_bytes(w, msg.branch)
 
     # when the account sends first operation in lifetime,
     # we need to reveal its publickey
     if msg.reveal is not None:
-        result += _encode_common(msg.reveal, "reveal")
-        result += msg.reveal.public_key
+        _encode_common(w, msg.reveal, "reveal")
+        write_bytes(w, msg.reveal.public_key)
 
     # transaction operation
     if msg.transaction is not None:
-        result += _encode_common(msg.transaction, "transaction")
-        result += _encode_zarith(msg.transaction.amount)
-        result += _encode_contract_id(msg.transaction.destination)
-        result += _encode_data_with_bool_prefix(msg.transaction.parameters)
+        _encode_common(w, msg.transaction, "transaction")
+        _encode_zarith(w, msg.transaction.amount)
+        _encode_contract_id(w, msg.transaction.destination)
+        _encode_data_with_bool_prefix(w, msg.transaction.parameters)
     # origination operation
     elif msg.origination is not None:
-        result += _encode_common(msg.origination, "origination")
-        result += msg.origination.manager_pubkey
-        result += _encode_zarith(msg.origination.balance)
-        result += _encode_bool(msg.origination.spendable)
-        result += _encode_bool(msg.origination.delegatable)
-        result += _encode_data_with_bool_prefix(msg.origination.delegate)
-        result += _encode_data_with_bool_prefix(msg.origination.script)
+        _encode_common(w, msg.origination, "origination")
+        write_bytes(w, msg.origination.manager_pubkey)
+        _encode_zarith(w, msg.origination.balance)
+        _encode_bool(w, msg.origination.spendable)
+        _encode_bool(w, msg.origination.delegatable)
+        _encode_data_with_bool_prefix(w, msg.origination.delegate)
+        _encode_data_with_bool_prefix(w, msg.origination.script)
     # delegation operation
     elif msg.delegation is not None:
-        result += _encode_common(msg.delegation, "delegation")
-        result += _encode_data_with_bool_prefix(msg.delegation.delegate)
-
-    return bytes(result)
+        _encode_common(w, msg.delegation, "delegation")
+        _encode_data_with_bool_prefix(w, msg.delegation.delegate)
 
 
-def _encode_common(operation, str_operation):
+def _encode_common(w: bytearray, operation, str_operation):
     operation_tags = {"reveal": 7, "transaction": 8, "origination": 9, "delegation": 10}
-    result = ustruct.pack("<b", operation_tags[str_operation])
-    result += _encode_contract_id(operation.source)
-    result += _encode_zarith(operation.fee)
-    result += _encode_zarith(operation.counter)
-    result += _encode_zarith(operation.gas_limit)
-    result += _encode_zarith(operation.storage_limit)
-    return result
+    write_uint8(w, operation_tags[str_operation])
+    _encode_contract_id(w, operation.source)
+    _encode_zarith(w, operation.fee)
+    _encode_zarith(w, operation.counter)
+    _encode_zarith(w, operation.gas_limit)
+    _encode_zarith(w, operation.storage_limit)
 
 
-def _encode_contract_id(contract_id):
-    return ustruct.pack("<b", contract_id.tag) + contract_id.hash
+def _encode_contract_id(w: bytearray, contract_id):
+    write_uint8(w, contract_id.tag)
+    write_bytes(w, contract_id.hash)
 
 
-def _encode_bool(boolean):
-    return ustruct.pack("<b", 255) if boolean else ustruct.pack("<b", 0)
+def _encode_bool(w: bytearray, boolean):
+    if boolean:
+        write_uint8(w, 255)
+    else:
+        write_uint8(w, 0)
 
 
-def _encode_data_with_bool_prefix(data):
-    return _encode_bool(True) + data if data is not None else _encode_bool(False)
+def _encode_data_with_bool_prefix(w: bytearray, data):
+    if data:
+        _encode_bool(w, True)
+        write_bytes(w, data)
+    else:
+        _encode_bool(w, False)
 
 
-def _encode_zarith(num):
-    result = bytes()
-
+def _encode_zarith(w: bytearray, num):
     while True:
         byte = num & 127
         num = num >> 7
 
         if num == 0:
-            result += ustruct.pack("<b", byte)
+            write_uint8(w, byte)
             break
 
-        result += ustruct.pack("<b", 128 | byte)
-
-    return result
+        write_uint8(w, 128 | byte)
diff --git a/tests/test_apps.tezos.address.py b/tests/test_apps.tezos.address.py
index 3d2bd8c38b..0640db4ee2 100644
--- a/tests/test_apps.tezos.address.py
+++ b/tests/test_apps.tezos.address.py
@@ -1,42 +1,43 @@
-from common import *
 from ubinascii import unhexlify
+
+from common import *
 from trezor.messages import TezosContractType
 from trezor.messages.TezosContractID import TezosContractID
+
 from apps.tezos.sign_tx import _get_address_from_contract
 
 
 class TestTezosAddress(unittest.TestCase):
-
     def test_get_address_from_contract(self):
         contracts = [
             TezosContractID(
                 tag=TezosContractType.Implicit,
-                hash=unhexlify("0090ec585b4d5fa39b20213e46b232cc57a4cfab4b")
+                hash=unhexlify("0090ec585b4d5fa39b20213e46b232cc57a4cfab4b"),
             ),
             TezosContractID(
                 tag=TezosContractType.Implicit,
-                hash=unhexlify("017dfb3fef44082eca8cd3eccebd77db44633ffc9e")
+                hash=unhexlify("017dfb3fef44082eca8cd3eccebd77db44633ffc9e"),
             ),
             TezosContractID(
                 tag=TezosContractType.Implicit,
-                hash=unhexlify("02c1fc1b7e503825068ff4fe2f8916f98af981eab1")
+                hash=unhexlify("02c1fc1b7e503825068ff4fe2f8916f98af981eab1"),
             ),
             TezosContractID(
                 tag=TezosContractType.Originated,
-                hash=unhexlify("65671dedc69669f066f45d586a2ecdeddacc95af00")
-            )
+                hash=unhexlify("65671dedc69669f066f45d586a2ecdeddacc95af00"),
+            ),
         ]
 
         outputs = [
             "tz1YrK8Hqt6GAPYRHAaeJmhETYyPSQCHTrkj",
             "tz2KoN7TFjhp96V2XikqYSGyDmVVUHXvkzko",
             "tz3e1k3QzCwEbRZrfHCwT3Npvw1rezmMQArY",
-            "KT1HpwLq2AjZgEQspiSnYmdtaHy4NgXw6BDC"
+            "KT1HpwLq2AjZgEQspiSnYmdtaHy4NgXw6BDC",
         ]
 
         for i, contract in enumerate(contracts):
             self.assertEqual(_get_address_from_contract(contract), outputs[i])
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     unittest.main()
diff --git a/tests/test_apps.tezos.encode.py b/tests/test_apps.tezos.encode.py
index 7553f41c30..ee875d3757 100644
--- a/tests/test_apps.tezos.encode.py
+++ b/tests/test_apps.tezos.encode.py
@@ -1,79 +1,100 @@
-from common import *
 from ubinascii import unhexlify
 
+from common import *
 from trezor.messages import TezosContractType
 from trezor.messages.TezosContractID import TezosContractID
 
-from apps.tezos.helpers import b58cencode, b58cdecode
+from apps.tezos.helpers import base58_decode_check, base58_encode_check
 from apps.tezos.sign_tx import (
-    _encode_zarith,
-    _encode_data_with_bool_prefix,
     _encode_bool,
-    _encode_contract_id
+    _encode_contract_id,
+    _encode_data_with_bool_prefix,
+    _encode_zarith,
 )
 
 
 class TestTezosEncoding(unittest.TestCase):
-
     def test_tezos_encode_zarith(self):
         inputs = [2000000, 159066, 200, 60000, 157000000, 0]
         outputs = ["80897a", "dada09", "c801", "e0d403", "c0c2ee4a", "00"]
 
         for i, o in zip(inputs, outputs):
-            self.assertEqual(_encode_zarith(i), unhexlify(o))
+            w = bytearray()
+            _encode_zarith(w, i)
+            self.assertEqual(bytes(w), unhexlify(o))
 
     def test_tezos_encode_data_with_bool_prefix(self):
-        self.assertEqual(_encode_data_with_bool_prefix(None), bytes([0]))
+        w = bytearray()
+        _encode_data_with_bool_prefix(w, None)
+        self.assertEqual(bytes(w), bytes([0]))
 
         data = "afffeb1dc3c0"
-        self.assertEqual(_encode_data_with_bool_prefix(unhexlify(data)),
-                         unhexlify("ff" + data))
+        w = bytearray()
+        _encode_data_with_bool_prefix(w, unhexlify(data))
+        self.assertEqual(bytes(w), unhexlify("ff" + data))
 
     def test_tezos_encode_bool(self):
-        self.assertEqual(_encode_bool(True), bytes([255]))
-        self.assertEqual(_encode_bool(False), bytes([0]))
+        w = bytearray()
+        _encode_bool(w, True)
+        self.assertEqual(bytes(w), bytes([255]))
+
+        w = bytearray()
+        _encode_bool(w, False)
+        self.assertEqual(bytes(w), bytes([0]))
 
     def test_tezos_encode_contract_id(self):
         implicit = TezosContractID(
             tag=TezosContractType.Implicit,
-            hash=unhexlify("00101368afffeb1dc3c089facbbe23f5c30b787ce9")
+            hash=unhexlify("00101368afffeb1dc3c089facbbe23f5c30b787ce9"),
+        )
+        w = bytearray()
+        _encode_contract_id(w, implicit)
+        self.assertEqual(
+            bytes(w), unhexlify("0000101368afffeb1dc3c089facbbe23f5c30b787ce9")
         )
-        self.assertEqual(_encode_contract_id(implicit),
-                         unhexlify("0000101368afffeb1dc3c089facbbe23f5c30b787ce9"))
 
         originated = TezosContractID(
             tag=TezosContractType.Originated,
-            hash=unhexlify("65671dedc69669f066f45d586a2ecdeddacc95af00")
+            hash=unhexlify("65671dedc69669f066f45d586a2ecdeddacc95af00"),
+        )
+        w = bytearray()
+        _encode_contract_id(w, originated)
+        self.assertEqual(
+            bytes(w), unhexlify("0165671dedc69669f066f45d586a2ecdeddacc95af00")
         )
-        self.assertEqual(_encode_contract_id(originated),
-                         unhexlify("0165671dedc69669f066f45d586a2ecdeddacc95af00"))
 
-    def test_tezos_b58cencode(self):
+    def test_tezos_base58_encode_check(self):
         pkh = unhexlify("101368afffeb1dc3c089facbbe23f5c30b787ce9")
 
-        self.assertEqual(b58cencode(pkh, prefix="tz1"),
-                         "tz1M72kkAJrntPtayM4yU4CCwQPLSdpEgRrn")
-        self.assertEqual(b58cencode(pkh, prefix="tz2"),
-                         "tz29nEixktH9p9XTFX7p8hATUyeLxXEz96KR")
-        self.assertEqual(b58cencode(pkh, prefix="tz3"),
-                         "tz3Mo3gHekQhCmykfnC58ecqJLXrjMKzkF2Q")
-        self.assertEqual(b58cencode(pkh), "2U14dJ6ED97bBHDZTQWA6umVL8SAVefXj")
+        self.assertEqual(
+            base58_encode_check(pkh, prefix="tz1"),
+            "tz1M72kkAJrntPtayM4yU4CCwQPLSdpEgRrn",
+        )
+        self.assertEqual(
+            base58_encode_check(pkh, prefix="tz2"),
+            "tz29nEixktH9p9XTFX7p8hATUyeLxXEz96KR",
+        )
+        self.assertEqual(
+            base58_encode_check(pkh, prefix="tz3"),
+            "tz3Mo3gHekQhCmykfnC58ecqJLXrjMKzkF2Q",
+        )
+        self.assertEqual(base58_encode_check(pkh), "2U14dJ6ED97bBHDZTQWA6umVL8SAVefXj")
 
-    def test_tezos_b58cdecode(self):
+    def test_tezos_base58_decode_check(self):
         pkh = unhexlify("101368afffeb1dc3c089facbbe23f5c30b787ce9")
 
         address = "tz1M72kkAJrntPtayM4yU4CCwQPLSdpEgRrn"
-        self.assertEqual(b58cdecode(address, prefix="tz1"), pkh)
+        self.assertEqual(base58_decode_check(address, prefix="tz1"), pkh)
 
         address = "tz29nEixktH9p9XTFX7p8hATUyeLxXEz96KR"
-        self.assertEqual(b58cdecode(address, prefix="tz2"), pkh)
+        self.assertEqual(base58_decode_check(address, prefix="tz2"), pkh)
 
         address = "tz3Mo3gHekQhCmykfnC58ecqJLXrjMKzkF2Q"
-        self.assertEqual(b58cdecode(address, prefix="tz3"), pkh)
+        self.assertEqual(base58_decode_check(address, prefix="tz3"), pkh)
 
         address = "2U14dJ6ED97bBHDZTQWA6umVL8SAVefXj"
-        self.assertEqual(b58cdecode(address), pkh)
+        self.assertEqual(base58_decode_check(address), pkh)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     unittest.main()