From 729fac20773d027dbf63e012fc094beefeb1025f Mon Sep 17 00:00:00 2001
From: Pavol Rusnak <stick@gk2.sk>
Date: Fri, 21 Feb 2014 18:56:51 +0100
Subject: [PATCH] add coin param to get_public_node call; new test for
 getpublickey message

---
 cmd.py                         |  3 ++-
 tests/test_msg_getaddress.py   |  1 -
 tests/test_msg_getpublickey.py | 28 ++++++++++++++++++++++++++++
 trezorlib/ckd_public.py        | 14 ++++++++++++++
 trezorlib/client.py            |  5 +++--
 5 files changed, 47 insertions(+), 4 deletions(-)
 create mode 100644 tests/test_msg_getpublickey.py

diff --git a/cmd.py b/cmd.py
index 62fdb2c00b..6c3161e977 100755
--- a/cmd.py
+++ b/cmd.py
@@ -102,7 +102,7 @@ class Commands(object):
 
     def get_public_node(self, args):
         address_n = self.client.expand_path(args.n)
-        return self.client.get_public_node(address_n)
+        return self.client.get_public_node(args.coin, address_n)
     
     def set_label(self, args):
         return self.client.apply_settings(label=args.label)
@@ -238,6 +238,7 @@ class Commands(object):
     )
 
     get_public_node.arguments = (
+        (('-c', '--coin'), {'type': str, 'default': 'Bitcoin'}),
         (('-n', '-address'), {'type': str}),
     )
 
diff --git a/tests/test_msg_getaddress.py b/tests/test_msg_getaddress.py
index b6419e02d4..b13a4b4223 100644
--- a/tests/test_msg_getaddress.py
+++ b/tests/test_msg_getaddress.py
@@ -1,7 +1,6 @@
 import unittest
 import common
 import trezorlib.ckd_public as bip32
-from trezorlib import tools
 
 class TestMsgGetaddress(common.TrezorTest):
 
diff --git a/tests/test_msg_getpublickey.py b/tests/test_msg_getpublickey.py
new file mode 100644
index 0000000000..40547c20d2
--- /dev/null
+++ b/tests/test_msg_getpublickey.py
@@ -0,0 +1,28 @@
+import unittest
+import common
+import trezorlib.ckd_public as bip32
+
+class TestMsgGetpublic_key(common.TrezorTest):
+
+    def test_btc(self):
+        self.setup_mnemonic_nopin_nopassphrase()
+        self.assertEqual(bip32.serialize( self.client.get_public_node('Bitcoin', []) ), 'xpub661MyMwAqRbcF1zGijBb2K6x9YiJPh58xpcCeLvTxMX6spkY3PcpJ4ABcCyWfskq5DDxM3e6Ez5ePCqG5bnPUXR4wL8TZWyoDaUdiWW7bKy')
+        self.assertEqual(bip32.serialize( self.client.get_public_node('Bitcoin', [1]) ), 'xpub68zNxjsTrV8y9AadThLW7dTAqEpZ7xBLFSyJ3X9pjTv6Njg6kxgjXJkzxq8u3ttnjBw1jupQHMP3gpGZzZqd1eh5S4GjkaMhPR18vMyUi8N')
+        self.assertEqual(bip32.serialize( self.client.get_public_node('Bitcoin', [0, -1]) ), 'xpub6A3FoZqYXj1AbW4thRwBh26YwZWbmoyjTaZwwxJjY1oKUpefLepL3RFS9DHKQrjAfxDrzDepYMDZPqXN6upQm3bHQ9xaXD5a3mqni3goF4v')
+        self.assertEqual(bip32.serialize( self.client.get_public_node('Bitcoin', [-9, 0]) ), 'xpub6A2h5mzLDfYginoD7q7wCWbq18wTbN9gducRr2w5NRTwdLeoT3cJSwefFqW7uXTpVFGtpUyDMBNYs3DNvvXx6NPjF9YEbUQrtxFSWnPtVrv')
+        self.assertEqual(bip32.serialize( self.client.get_public_node('Bitcoin', [0, 9999999]) ), 'xpub6A3FoZqQEK6iwLZ4HFkqSo5fb35BH4bpjC4SPZ63prfLdGYPwYxEuC6o91bUvFFdMzKWe5rs3axHRUjxJaSvBnKKFtnfLwDACRxPxabsv2r')
+
+    def test_ltc(self):
+        self.setup_mnemonic_nopin_nopassphrase()
+        self.assertEqual(bip32.serialize( self.client.get_public_node('Litecoin', []) ), 'Ltub2SSUS19CirucVPGDKDBatBDBEM2s9UbH66pBURfaKrMocCPLhQ7Z7hecy5VYLHA5fRdXwB2e61j2VJCNzVsqKTCVEU1vECjqi5EyczFX9xp')
+        self.assertEqual(bip32.serialize( self.client.get_public_node('Litecoin', [1]) ), 'Ltub2VRVRP5VjvSyPXra4BLVyVZPv397sjhUNjBGsbtw6xko77JuQyBULxFSKheviJJ3KQLbL3Cx8P2RnudguTw4raUVjCACRG7jsumUptYx55C')
+        self.assertEqual(bip32.serialize( self.client.get_public_node('Litecoin', [0, -1]) ), 'Ltub2WUNGD3aRAKAqsLqHuwBYtCn2MqAXbVsarmvn33quWe2DCHTzfK4s4jsW5oM5G8RGAdSaM3NPNrwVvtV1ourbyNhhHr3BtqcYGc8caf5GoT')
+        self.assertEqual(bip32.serialize( self.client.get_public_node('Litecoin', [-9, 0]) ), 'Ltub2WToYRCN76rgyA59iK7w4Ni45wG2M9fpmBpQg7gBjvJeMiHc7473Gb96ci29Zvs55TgUQcMmCD1vy8aVqpdPwJB9YHRhGAAuPT1nRLLXmFu')
+        self.assertEqual(bip32.serialize( self.client.get_public_node('Litecoin', [0, 9999999]) ), 'Ltub2WUNGD3S7kQjBhpzsjkqJfBtfqPk2r7xrUGRDdqACMW3MeBCbZSyiqbEVt7WaeesxCj6EDFQtcbfXa75DUYN2i6jZ2g81cyCgvijs9J2u2n')
+
+    def test_tbtc(self):
+        self.setup_mnemonic_nopin_nopassphrase()
+        self.assertEqual(bip32.serialize( self.client.get_public_node('Testnet', [111, 42]) ), 'tpubDAgixSyai5PWbc8N1mBkHDR5nLgAnHFtY7r4y5EzxqAxrt9YUDpZL3kaRoHVvCfrcwNo31c2isBP2uTHcZxEosuKbyJhCAbrvGoPuLUZ7Mz')
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/trezorlib/ckd_public.py b/trezorlib/ckd_public.py
index 83b7b7d344..bec0bad1ad 100644
--- a/trezorlib/ckd_public.py
+++ b/trezorlib/ckd_public.py
@@ -90,3 +90,17 @@ def get_subnode(node, i):
     node_out.public_key = point_to_pubkey(point)
 
     return node_out
+
+def serialize(node):
+    s = ''
+    s += struct.pack('>I', node.version)
+    s += struct.pack('>B', node.depth)
+    s += struct.pack('>I', node.fingerprint)
+    s += struct.pack('>I', node.child_num)
+    s += node.chain_code
+    if node.private_key:
+        s += '\x00' + node.private_key
+    else :
+        s += node.public_key
+    s += tools.Hash(s)[:4]
+    return tools.b58encode(s)
diff --git a/trezorlib/client.py b/trezorlib/client.py
index a5dfff445d..aa20348ed8 100644
--- a/trezorlib/client.py
+++ b/trezorlib/client.py
@@ -306,8 +306,9 @@ class ProtocolMixin(object):
 
     @field('node')
     @expect(proto.PublicKey)
-    def get_public_node(self, n):
-        return self.call(proto.GetPublicKey(address_n=n))
+    def get_public_node(self, coin_name, n):
+        n = self._convert_prime(n)
+        return self.call(proto.GetPublicKey(address_n=n, coin_name=coin_name))
 
     @field('address')
     @expect(proto.Address)