From 6aef64d367d7a2c705d84b77501cf854d3794696 Mon Sep 17 00:00:00 2001
From: Tomas Susanka <tsusanka@gmail.com>
Date: Wed, 4 Apr 2018 21:06:37 +0200
Subject: [PATCH] nem: compute address from pubkey

---
 .../modtrezorcrypto/modtrezorcrypto-nem.h     | 19 ++++++++++
 mocks/generated/trezorcrypto.py               |  6 +++
 tests/test_apps.nem.address.py                | 37 +++++++++++++++++++
 tests/test_apps.nem.hdnode.py                 |  2 +-
 4 files changed, 63 insertions(+), 1 deletion(-)
 create mode 100644 tests/test_apps.nem.address.py

diff --git a/embed/extmod/modtrezorcrypto/modtrezorcrypto-nem.h b/embed/extmod/modtrezorcrypto/modtrezorcrypto-nem.h
index 67e35a3b7d..d3dfb7dd23 100644
--- a/embed/extmod/modtrezorcrypto/modtrezorcrypto-nem.h
+++ b/embed/extmod/modtrezorcrypto/modtrezorcrypto-nem.h
@@ -34,9 +34,28 @@ STATIC mp_obj_t mod_trezorcrypto_nem_validate_address(mp_obj_t address, mp_obj_t
 }
 STATIC MP_DEFINE_CONST_FUN_OBJ_2(mod_trezorcrypto_nem_validate_address_obj, mod_trezorcrypto_nem_validate_address);
 
+/// def compute_address(public_key: bytes, network: int) -> str:
+///     '''
+///     Compute a NEM address from a public key
+///     '''
+STATIC mp_obj_t mod_trezorcrypto_nem_compute_address(mp_obj_t public_key, mp_obj_t network) {
+
+    mp_buffer_info_t p;
+    mp_get_buffer_raise(public_key, &p, MP_BUFFER_READ);
+    uint32_t n = mp_obj_get_int_truncated(network);
+
+    char address[NEM_ADDRESS_SIZE + 1];
+    if (!nem_get_address(p.buf, n, address)) {
+        mp_raise_ValueError("Failed to compute a NEM address from a provided public key");
+    }
+    return mp_obj_new_str(address, strlen(address), false);
+}
+STATIC MP_DEFINE_CONST_FUN_OBJ_2(mod_trezorcrypto_nem_compute_address_obj, mod_trezorcrypto_nem_compute_address);
+
 // objects definition
 STATIC const mp_rom_map_elem_t mod_trezorcrypto_nem_globals_table[] = {
     { MP_ROM_QSTR(MP_QSTR_validate_address), MP_ROM_PTR(&mod_trezorcrypto_nem_validate_address_obj) },
+    { MP_ROM_QSTR(MP_QSTR_compute_address), MP_ROM_PTR(&mod_trezorcrypto_nem_compute_address_obj) },
 };
 STATIC MP_DEFINE_CONST_DICT(mod_trezorcrypto_nem_globals, mod_trezorcrypto_nem_globals_table);
 
diff --git a/mocks/generated/trezorcrypto.py b/mocks/generated/trezorcrypto.py
index 6023d51ccd..62d89b923b 100644
--- a/mocks/generated/trezorcrypto.py
+++ b/mocks/generated/trezorcrypto.py
@@ -320,6 +320,12 @@ def validate_address(address: str, network: int) -> bool:
     Validate a NEM address
     '''
 
+# extmod/modtrezorcrypto/modtrezorcrypto-nem.h
+def compute_address(public_key: bytes, network: int) -> str:
+    '''
+    Compute a NEM address from a public key
+    '''
+
 # extmod/modtrezorcrypto/modtrezorcrypto-nist256p1.h
 def generate_secret() -> bytes:
     '''
diff --git a/tests/test_apps.nem.address.py b/tests/test_apps.nem.address.py
new file mode 100644
index 0000000000..dab6faa8e6
--- /dev/null
+++ b/tests/test_apps.nem.address.py
@@ -0,0 +1,37 @@
+from common import *
+from ubinascii import unhexlify
+from trezor.crypto import nem
+from apps.nem.helpers import NEM_NETWORK_MAINNET, NEM_NETWORK_TESTNET
+
+
+class TestNemAddress(unittest.TestCase):
+
+    def test_addresses(self):
+        pubkey = unhexlify('c5f54ba980fcbb657dbaaa42700539b207873e134d2375efeab5f1ab52f87844')
+        address = nem.compute_address(pubkey, NEM_NETWORK_MAINNET)
+        self.assertEqual(address, 'NDD2CT6LQLIYQ56KIXI3ENTM6EK3D44P5JFXJ4R4')
+
+        pubkey = unhexlify('114171230ad6f8522a000cdc73fbc5c733b30bb71f2b146ccbdf34499f79a810')
+        address = nem.compute_address(pubkey, NEM_NETWORK_MAINNET)
+        self.assertEqual(address, 'NCUKWDY3J3THKQHAKOK5ALF6ANJQABZHCH7VN6DP')
+
+    def test_validate_address(self):
+        validity = nem.validate_address('NDD2CT6LQLIYQ56KIXI3ENTM6EK3D44P5JFXJ4R4', NEM_NETWORK_MAINNET)
+        self.assertTrue(validity)
+
+        validity = nem.validate_address('NCUKWDY3J3THKQHAKOK5ALF6ANJQABZHCH7VN6DP', NEM_NETWORK_MAINNET)
+        self.assertTrue(validity)
+
+        validity = nem.validate_address('TAU5HO3DRQZNELFEMZZTUKQEZGQ7IUAHKPO7OOLK', NEM_NETWORK_TESTNET)
+        self.assertTrue(validity)
+
+        validity = nem.validate_address('nope', NEM_NETWORK_TESTNET)
+        self.assertFalse(validity)
+
+        # not valid on testnet
+        validity = nem.validate_address('NCUKWDY3J3THKQHAKOK5ALF6ANJQABZHCH7VN6DP', NEM_NETWORK_TESTNET)
+        self.assertFalse(validity)
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/tests/test_apps.nem.hdnode.py b/tests/test_apps.nem.hdnode.py
index ee35a93df1..ea109c6550 100644
--- a/tests/test_apps.nem.hdnode.py
+++ b/tests/test_apps.nem.hdnode.py
@@ -4,7 +4,7 @@ from trezor.crypto import bip32
 from apps.nem.helpers import NEM_NETWORK_MAINNET, NEM_CURVE
 
 
-class TestNemAddress(unittest.TestCase):
+class TestNemHDNode(unittest.TestCase):
 
     def test_addresses(self):
         # test vectors from https://raw.githubusercontent.com/NemProject/nem-test-vectors/master/1.test-keys.dat