From 9674a58db036962a2b2913819914a81cc22ab9b3 Mon Sep 17 00:00:00 2001
From: Tomas Susanka <tsusanka@gmail.com>
Date: Wed, 21 Mar 2018 14:30:28 +0100
Subject: [PATCH] nem: encrypted payload

---
 src/apps/nem/helpers.py     |  4 ++++
 src/apps/nem/signing.py     | 40 +++++++++++++++++++++++++++----------
 src/apps/nem/transaction.py |  1 +
 3 files changed, 34 insertions(+), 11 deletions(-)

diff --git a/src/apps/nem/helpers.py b/src/apps/nem/helpers.py
index 39d48aba5b..091498e6fe 100644
--- a/src/apps/nem/helpers.py
+++ b/src/apps/nem/helpers.py
@@ -15,6 +15,10 @@ NEM_TRANSACTION_TYPE_PROVISION_NAMESPACE = const(0x2001)
 NEM_TRANSACTION_TYPE_MOSAIC_CREATION = const(0x4001)
 NEM_TRANSACTION_TYPE_MOSAIC_SUPPLY_CHANGE = const(0x4002)
 
+NEM_SALT_SIZE = const(32)
+AES_BLOCK_SIZE = const(16)
+NEM_HASH_ALG = 'keccak'
+
 
 def nem_validate_network(network):
     if network in (NEM_NETWORK_MAINNET, NEM_NETWORK_TESTNET, NEM_NETWORK_MIJIN):
diff --git a/src/apps/nem/signing.py b/src/apps/nem/signing.py
index 2f07ef5a2e..4973a57a39 100644
--- a/src/apps/nem/signing.py
+++ b/src/apps/nem/signing.py
@@ -1,17 +1,25 @@
 from apps.nem.transaction import *
 from apps.nem.layout import *
+from apps.nem import helpers
+from apps.common import seed
 from trezor.messages.NEMSignTx import NEMSignTx
 from trezor.messages.NEMSignedTx import NEMSignedTx
+from trezor.crypto.curve import ed25519
+from trezor.crypto import random
 
 
 async def nem_sign_tx(ctx, msg: NEMSignTx):
-    from ..common import seed
-    from trezor.crypto.curve import ed25519
-
-    # if len(msg.transfer.public_key):
-        # todo encrypt
 
     node = await seed.derive_node(ctx, msg.transaction.address_n, NEM_CURVE)
+
+    payload = msg.transfer.payload
+    encrypted = False
+    if msg.transfer.public_key is not None:
+        if payload is None:
+            raise ValueError("Public key provided but no payload to encrypt")
+        payload = _nem_encrypt(node, msg.transfer.public_key, msg.transfer.payload)
+        encrypted = True
+
     # 0x01 prefix is not part of the actual public key, hence removed
     public_key = node.public_key()[1:]
 
@@ -23,21 +31,31 @@ async def nem_sign_tx(ctx, msg: NEMSignTx):
         msg.transaction.deadline,
         msg.transfer.recipient,
         msg.transfer.amount,
-        msg.transfer.payload,  # todo might require encryption
-        msg.transfer.public_key is not None,
+        payload,
+        encrypted,
         len(msg.transfer.mosaics)
     )
 
     for mosaic in msg.transfer.mosaics:
         nem_transaction_write_mosaic(tx, mosaic.namespace, mosaic.mosaic, mosaic.quantity)
 
-    await require_confirm_action(ctx)
-    await require_confirm_fee(ctx, msg.transfer.amount, msg.transaction.fee)
-    await require_confirm_tx(ctx, msg.transfer.recipient, msg.transfer.amount)
+    if payload:  # confirm unencrypted
+        # todo encrypted vs unencrypted
+        await require_confirm_action(ctx)  # todo
 
-    signature = ed25519.sign(node.private_key(), tx, 'keccak')
+    await require_confirm_fee(ctx, msg.transfer.amount, msg.transaction.fee)  # todo
+    await require_confirm_tx(ctx, msg.transfer.recipient, msg.transfer.amount)  # todo
+
+    signature = ed25519.sign(node.private_key(), tx, helpers.NEM_HASH_ALG)
 
     resp = NEMSignedTx()
     resp.data = tx
     resp.signature = signature
     return resp
+
+
+def _nem_encrypt(node, public_key: bytes, payload: bytes) -> bytes:
+    salt = random.bytes(helpers.NEM_SALT_SIZE)
+    iv = random.bytes(helpers.AES_BLOCK_SIZE)
+    encrypted = node.nem_encrypt(public_key, iv, salt, payload)
+    return iv + salt + encrypted
diff --git a/src/apps/nem/transaction.py b/src/apps/nem/transaction.py
index 90cac69f75..7fc44c236a 100644
--- a/src/apps/nem/transaction.py
+++ b/src/apps/nem/transaction.py
@@ -1,6 +1,7 @@
 
 from .helpers import *
 from .writers import *
+from ubinascii import hexlify
 
 
 def nem_transaction_create_transfer(network: int, timestamp: int, signer_public_key: bytes, fee: int, deadline: int,