From abd20741f62d53480d9df53e5745b8d61f29d3f3 Mon Sep 17 00:00:00 2001
From: Jan Pochyla <jpochyla@gmail.com>
Date: Thu, 6 Oct 2016 15:04:38 +0200
Subject: [PATCH] add seed module, GetPublicKey now works

---
 src/apps/common/seed.py                  | 69 ++++++++++++++++++++++++
 src/apps/wallet/layout_get_public_key.py | 27 +++++-----
 src/trezor/crypto/__init__.py            |  2 +
 3 files changed, 83 insertions(+), 15 deletions(-)
 create mode 100644 src/apps/common/seed.py

diff --git a/src/apps/common/seed.py b/src/apps/common/seed.py
new file mode 100644
index 00000000..1aa08e4c
--- /dev/null
+++ b/src/apps/common/seed.py
@@ -0,0 +1,69 @@
+from trezor import wire
+
+# FIXME: this is a stub
+
+# TODO: decomplect the MVC layers
+# TODO: most likely storage sensitive data in c
+# TODO: check pin in constant time
+# TODO: pin failure counter
+
+_cached_seed = None
+_cached_root_node = None
+
+
+async def get_node(session_id: int, path: list):
+    from trezor import ui
+    ui.display.clear()
+    ui.display.text_center(120, 120, 'Deriving key...',
+                           ui.NORMAL, ui.GREY, ui.BLACK)
+    ui.display.refresh()
+    node = await get_root_node(session_id)
+    node.derive_path(path)
+    return node
+
+
+async def get_root_node(session_id: int):
+    global _cached_root_node
+    if _cached_root_node is None:
+        _cached_root_node = await compute_root_node(session_id)
+    return _cached_root_node
+
+
+async def compute_root_node(session_id: int):
+    from trezor.crypto import bip32
+    seed = await get_seed(session_id)
+    return bip32.from_seed(seed, 'secp256k1')
+
+
+async def get_seed(session_id: int) -> bytes:
+    global _cached_seed
+    if _cached_seed is None:
+        _cached_seed = await compute_seed(session_id)
+    return _cached_seed
+
+
+async def compute_seed(session_id):
+    from trezor.crypto import bip39
+    from trezor.messages.FailureType import PinInvalid, Other
+    from .request_pin import request_pin
+    from . import storage
+
+    try:
+        st = storage.get(session_id)
+    except KeyError:
+        raise wire.FailureError(Other, 'Device is not initialized')
+
+    st_pin = getattr(st, 'pin', '')
+    if st_pin and st_pin != await request_pin(session_id):
+        raise wire.FailureError(PinInvalid, 'PIN is incorrect')
+
+    st_passphrase_protection = getattr(st, 'passphrase_protection', False)
+    if st_passphrase_protection:
+        from trezor.messages.PassphraseRequest import PassphraseRequest
+        from trezor.messages.wire_types import PassphraseAck
+        ack = await wire.reply_message(session_id, PassphraseRequest(), PassphraseAck)
+        passphrase = ack.passphrase
+    else:
+        passphrase = ''
+
+    return bip39.seed(st.mnemonic, passphrase)
diff --git a/src/apps/wallet/layout_get_public_key.py b/src/apps/wallet/layout_get_public_key.py
index 6f3e3db2..6b77484f 100644
--- a/src/apps/wallet/layout_get_public_key.py
+++ b/src/apps/wallet/layout_get_public_key.py
@@ -1,22 +1,19 @@
-from trezor import wire, ui
 from trezor.utils import unimport
 
 
 @unimport
-async def layout_get_public_key(session_id, message):
-    from trezor.messages.PublicKey import PublicKey
+async def layout_get_public_key(message, session_id):
     from trezor.messages.HDNodeType import HDNodeType
+    from trezor.messages.PublicKey import PublicKey
+    from ..common.seed import get_node
 
-    # TODO: protect with pin
-    # TODO: fail if not initialized
-    # TODO: derive correct node
+    node = await get_node(session_id, message.address_n)
 
-    pubkey = PublicKey()
-    pubkey.node = HDNodeType()
-    pubkey.node.depth = 0
-    pubkey.node.child_num = 0
-    pubkey.node.fingerprint = 0
-    pubkey.node.chain_code = 'deadbeef'
-    pubkey.node.public_key = 'deadbeef'
-
-    return pubkey
+    node_xpub = node.serialize_public(0x0488B21E)
+    node_type = HDNodeType(
+        depth=node.depth(),
+        child_num=node.child_num(),
+        fingerprint=node.fingerprint(),
+        chain_code=node.chain_code(),
+        public_key=node.public_key())
+    return PublicKey(node=node_type, xpub=node_xpub)
diff --git a/src/trezor/crypto/__init__.py b/src/trezor/crypto/__init__.py
index 334b024e..85622438 100644
--- a/src/trezor/crypto/__init__.py
+++ b/src/trezor/crypto/__init__.py
@@ -1,8 +1,10 @@
+from TrezorCrypto import Bip32
 from TrezorCrypto import Bip39
 from TrezorCrypto import Pbkdf2 as pbkdf2
 from TrezorCrypto import Random
 from TrezorCrypto import SSSS
 
+bip32 = Bip32()
 bip39 = Bip39()
 random = Random()
 ssss = SSSS()