From 6e79da8df137da91e582fa569d9591a21c397ab7 Mon Sep 17 00:00:00 2001
From: Pavol Rusnak <stick@gk2.sk>
Date: Fri, 18 Nov 2016 14:59:36 +0100
Subject: [PATCH] apps.wallet: add CipherKeyValue

---
 src/apps/wallet/__init__.py              | 10 +++++-
 src/apps/wallet/layout_cipherkeyvalue.py | 41 ++++++++++++++++++++++++
 2 files changed, 50 insertions(+), 1 deletion(-)
 create mode 100644 src/apps/wallet/layout_cipherkeyvalue.py

diff --git a/src/apps/wallet/__init__.py b/src/apps/wallet/__init__.py
index 571aaf3ff1..9cdacec8e8 100644
--- a/src/apps/wallet/__init__.py
+++ b/src/apps/wallet/__init__.py
@@ -3,7 +3,8 @@ from trezor.utils import unimport
 from trezor.messages.wire_types import \
     GetPublicKey, GetAddress, SignTx, EstimateTxSize, \
     SignMessage, VerifyMessage, \
-    SignIdentity
+    SignIdentity, \
+    CipherKeyValue
 
 
 @unimport
@@ -51,6 +52,12 @@ def dispatch_SignIdentity(*args, **kwargs):
     return layout_sign_identity(*args, **kwargs)
 
 
+@unimport
+def dispatch_CipherKeyValue(*args, **kwargs):
+    from .layout_cipherkeyvalue import layout_cipherkeyvalue
+    return layout_cipherkeyvalue(*args, **kwargs)
+
+
 def boot():
     register_type(GetPublicKey, protobuf_handler, dispatch_GetPublicKey)
     register_type(GetAddress, protobuf_handler, dispatch_GetAddress)
@@ -59,3 +66,4 @@ def boot():
     register_type(SignMessage, protobuf_handler, dispatch_SignMessage)
     register_type(VerifyMessage, protobuf_handler, dispatch_VerifyMessage)
     register_type(SignIdentity, protobuf_handler, dispatch_SignIdentity)
+    register_type(CipherKeyValue, protobuf_handler, dispatch_CipherKeyValue)
diff --git a/src/apps/wallet/layout_cipherkeyvalue.py b/src/apps/wallet/layout_cipherkeyvalue.py
new file mode 100644
index 0000000000..771fd33de8
--- /dev/null
+++ b/src/apps/wallet/layout_cipherkeyvalue.py
@@ -0,0 +1,41 @@
+from trezor import wire, ui
+from trezor.utils import unimport
+
+
+@unimport
+async def layout_cipherkeyvalue(msg, session_id):
+    from trezor.messages.CipheredKeyValue import CipheredKeyValue
+    from ..common.seed import get_node
+    from trezor.crypto.hashlib import sha512
+    from trezor.crypto import hmac
+    from trezor.crypto.aes import AES_CBC_Encrypt, AES_CBC_Decrypt
+
+    if len(msg.value) % 16 > 0:
+        raise ValueError('Value length must be a multiple of 16')
+
+    ui.display.clear()
+    ui.display.text(10, 30, 'CipherKeyValue',
+                    ui.BOLD, ui.LIGHT_GREEN, ui.BLACK)
+    ui.display.text(10, 60, msg.key, ui.MONO, ui.WHITE, ui.BLACK)
+
+    node = await get_node(session_id, msg.address_n)
+    seckey = node.private_key()
+
+    data = msg.key
+    data += 'E1' if msg.ask_on_encrypt else 'E0'
+    data += 'D1' if msg.ask_on_decrypt else 'D0'
+    data = hmac.new(seckey, data, sha512).digest()
+    key = data[:32]
+    if msg.iv and len(msg.iv) == 16:
+        iv = msg.iv
+    else:
+        iv = data[32:48]
+
+    if msg.encrypt:
+        aes = AES_CBC_Encrypt(key=key, iv=iv)
+    else:
+        aes = AES_CBC_Decrypt(key=key, iv=iv)
+
+    value = aes.update(msg.value)
+
+    return CipheredKeyValue(value=value)