From 414b33dbfbe4c0dfdf1e0c901d43882935299234 Mon Sep 17 00:00:00 2001
From: Pavol Rusnak <pavol@rusnak.io>
Date: Fri, 24 May 2019 15:52:57 +0200
Subject: [PATCH] core/eos: reduce code duplication

---
 core/src/apps/eos/get_public_key.py | 21 ++++++++-------------
 core/src/apps/eos/helpers.py        |  4 ++--
 core/src/apps/eos/sign_tx.py        |  2 +-
 3 files changed, 11 insertions(+), 16 deletions(-)

diff --git a/core/src/apps/eos/get_public_key.py b/core/src/apps/eos/get_public_key.py
index de72c0597..d86ce71ca 100755
--- a/core/src/apps/eos/get_public_key.py
+++ b/core/src/apps/eos/get_public_key.py
@@ -1,5 +1,4 @@
 from trezor import wire
-from trezor.crypto import base58
 from trezor.crypto.curve import secp256k1
 from trezor.crypto.hashlib import ripemd160
 from trezor.messages.EosGetPublicKey import EosGetPublicKey
@@ -7,23 +6,19 @@ from trezor.messages.EosPublicKey import EosPublicKey
 
 from apps.common import paths
 from apps.eos import CURVE
-from apps.eos.helpers import validate_full_path
+from apps.eos.helpers import base58_encode, validate_full_path
 from apps.eos.layout import require_get_public_key
 
 
-def _ripemd160_32(data: bytes) -> bytes:
-    return ripemd160(data).digest()[:4]
-
-
 def _public_key_to_wif(pub_key: bytes) -> str:
-    if len(pub_key) == 65:
-        head = 0x03 if pub_key[64] & 0x01 else 0x02
-        compresed_pub_key = bytes([head]) + pub_key[1:33]
-    elif len(pub_key) == 33:
-        compresed_pub_key = pub_key
+    if pub_key[0] == 0x04 and len(pub_key) == 65:
+        head = b"\x03" if pub_key[64] & 0x01 else b"\x02"
+        compressed_pub_key = head + pub_key[1:33]
+    elif pub_key[0] in [0x02, 0x03] and len(pub_key) == 33:
+        compressed_pub_key = pub_key
     else:
-        raise wire.DataError("invalid public key length")
-    return "EOS" + base58.encode_check(compresed_pub_key, _ripemd160_32)
+        raise wire.DataError("invalid public key")
+    return base58_encode("EOS", "", compressed_pub_key)
 
 
 def _get_public_key(node):
diff --git a/core/src/apps/eos/helpers.py b/core/src/apps/eos/helpers.py
index 9248e7e38..0fc6bac9b 100644
--- a/core/src/apps/eos/helpers.py
+++ b/core/src/apps/eos/helpers.py
@@ -7,9 +7,9 @@ from apps.common import HARDENED
 def base58_encode(prefix: str, sig_prefix: str, data: bytes) -> str:
     b58 = base58.encode(data + base58.ripemd160_32(data + sig_prefix.encode()))
     if sig_prefix:
-        return prefix + "_" + sig_prefix + "_" + b58
+        return prefix + sig_prefix + "_" + b58
     else:
-        return prefix + "_" + b58
+        return prefix + b58
 
 
 def eos_name_to_string(value) -> str:
diff --git a/core/src/apps/eos/sign_tx.py b/core/src/apps/eos/sign_tx.py
index 9c2b787ac..f959da686 100644
--- a/core/src/apps/eos/sign_tx.py
+++ b/core/src/apps/eos/sign_tx.py
@@ -36,7 +36,7 @@ async def sign_tx(ctx, msg: EosSignTx, keychain):
         node.private_key(), digest, True, secp256k1.CANONICAL_SIG_EOS
     )
 
-    return EosSignedTx(signature=base58_encode("SIG", "K1", signature))
+    return EosSignedTx(signature=base58_encode("SIG_", "K1", signature))
 
 
 async def _init(ctx, sha, msg):