From a93e745532185ba44d6ee8a2831461559572beb8 Mon Sep 17 00:00:00 2001
From: Roman Zeyde <me@romanzey.de>
Date: Sat, 10 Mar 2018 18:44:26 +0200
Subject: [PATCH] apps/wallet/{sign_identity, ecdh}: show protocol as part of
 the header

---
 src/apps/wallet/ecdh.py          | 10 ++++++----
 src/apps/wallet/sign_identity.py | 25 ++++++++++++++++++++-----
 2 files changed, 26 insertions(+), 9 deletions(-)

diff --git a/src/apps/wallet/ecdh.py b/src/apps/wallet/ecdh.py
index 6ed4b6070..a69fc40ed 100644
--- a/src/apps/wallet/ecdh.py
+++ b/src/apps/wallet/ecdh.py
@@ -4,7 +4,7 @@ from trezor.messages.ECDHSessionKey import ECDHSessionKey
 from ustruct import pack, unpack
 from trezor.utils import chunks
 from apps.common.confirm import require_confirm
-from apps.wallet.sign_identity import serialize_identity
+from apps.wallet.sign_identity import serialize_identity, serialize_identity_without_proto
 from trezor.ui.text import Text
 
 from ..common import seed
@@ -16,7 +16,7 @@ async def get_ecdh_session_key(ctx, msg):
 
     identity = serialize_identity(msg.identity)
 
-    await require_confirm_ecdh_session_key(ctx, identity)
+    await require_confirm_ecdh_session_key(ctx, msg.identity)
 
     address_n = get_ecdh_path(identity, msg.identity.index or 0)
     node = await seed.derive_node(ctx, address_n, msg.ecdsa_curve_name)
@@ -28,8 +28,10 @@ async def get_ecdh_session_key(ctx, msg):
 
 
 async def require_confirm_ecdh_session_key(ctx, identity):
-    lines = chunks(identity, 18)
-    content = Text('Decrypt', ui.ICON_DEFAULT, ui.MONO, *lines, max_lines=5)
+    lines = chunks(serialize_identity_without_proto(identity), 18)
+    proto = identity.proto.upper() if identity.proto else 'identity'
+    header = 'Decrypt %s' % (proto,)
+    content = Text(header, ui.ICON_DEFAULT, ui.MONO, *lines, max_lines=5)
     await require_confirm(ctx, content)
 
 
diff --git a/src/apps/wallet/sign_identity.py b/src/apps/wallet/sign_identity.py
index b915d550a..7904d4192 100644
--- a/src/apps/wallet/sign_identity.py
+++ b/src/apps/wallet/sign_identity.py
@@ -1,5 +1,6 @@
 from trezor import ui
 from trezor.crypto.hashlib import sha256
+from trezor.messages.IdentityType import IdentityType
 from trezor.messages.SignedIdentity import SignedIdentity
 from ustruct import pack, unpack
 from trezor.utils import chunks
@@ -15,7 +16,7 @@ async def sign_identity(ctx, msg):
 
     identity = serialize_identity(msg.identity)
 
-    await require_confirm_sign_identity(ctx, identity, msg.challenge_visual)
+    await require_confirm_sign_identity(ctx, msg.identity, msg.challenge_visual)
 
     address_n = get_identity_path(identity, msg.identity.index or 0)
     node = await seed.derive_node(ctx, address_n, msg.ecdsa_curve_name)
@@ -44,10 +45,16 @@ async def sign_identity(ctx, msg):
 
 
 async def require_confirm_sign_identity(ctx, identity, challenge_visual):
-    lines = chunks(identity, 18)
-    content = Text('Sign identity', ui.ICON_DEFAULT,
-                   challenge_visual,
-                   ui.MONO, *lines, max_lines=5)
+    lines = []
+    if challenge_visual:
+        lines.append(challenge_visual)
+
+    lines.append(ui.MONO)
+    lines.extend(chunks(serialize_identity_without_proto(identity), 18))
+
+    proto = identity.proto.upper() if identity.proto else 'identity'
+    header = 'Sign %s' % (proto,)
+    content = Text(header, ui.ICON_DEFAULT, *lines, max_lines=5)
     await require_confirm(ctx, content)
 
 
@@ -66,6 +73,14 @@ def serialize_identity(identity):
     return s
 
 
+def serialize_identity_without_proto(identity):
+    proto = identity.proto
+    identity.proto = None  # simplify serialized identity string
+    s = serialize_identity(identity)
+    identity.proto = proto
+    return s
+
+
 def get_identity_path(identity: str, index: int):
     identity_hash = sha256(pack('<I', index) + identity).digest()