From d0bb7edbfa282746d9c481ede035e8612ef7a9ff Mon Sep 17 00:00:00 2001
From: Jan Pochyla <jpochyla@gmail.com>
Date: Thu, 6 Oct 2016 15:03:54 +0200
Subject: [PATCH] storage: change api to protobuf

---
 src/apps/common/storage.py                 | 23 +++++++++++------
 src/apps/management/layout_load_device.py  | 29 +++++++++++-----------
 src/apps/management/layout_reset_device.py |  9 +++----
 src/apps/management/layout_wipe_device.py  |  4 +--
 4 files changed, 36 insertions(+), 29 deletions(-)

diff --git a/src/apps/common/storage.py b/src/apps/common/storage.py
index 5d8ca17f..4609b4f5 100644
--- a/src/apps/common/storage.py
+++ b/src/apps/common/storage.py
@@ -1,18 +1,27 @@
 from micropython import const
 from trezor import config
+from trezor.messages.Storage import Storage
 
 
-APP_MANAGEMENT = const(1)
+APP_COMMON = const(1)
 CFG_STORAGE = const(1)
 
 
-def get_storage(session_id):
-    return config.get(session_id, APP_MANAGEMENT, CFG_STORAGE)
+def has(session_id):
+    buf = config.get(session_id, APP_COMMON, CFG_STORAGE)
+    return bool(buf)
 
 
-def set_storage(session_id, buf):
-    config.set(session_id, APP_MANAGEMENT, CFG_STORAGE, buf)
+def get(session_id):
+    buf = config.get(session_id, APP_COMMON, CFG_STORAGE)
+    if not buf:
+        raise KeyError('Storage is not initialized')
+    return Storage.loads(buf)
 
 
-def clear_storage(session_id):
-    set_storage(session_id, b'')
+def set(session_id, st):
+    config.set(session_id, APP_COMMON, CFG_STORAGE, st.dumps())
+
+
+def clear(session_id):
+    config.set(session_id, APP_COMMON, CFG_STORAGE, b'')
diff --git a/src/apps/management/layout_load_device.py b/src/apps/management/layout_load_device.py
index 202588dd..46324aca 100644
--- a/src/apps/management/layout_load_device.py
+++ b/src/apps/management/layout_load_device.py
@@ -9,34 +9,33 @@ async def layout_load_device(message, session_id):
     from trezor.messages.Success import Success
     from trezor.messages.FailureType import UnexpectedMessage, Other
     from trezor.ui.text import Text
-    from ..common.storage import get_storage, set_storage
     from ..common.confirm import require_confirm
+    from ..common import storage
 
-    if get_storage(session_id):
+    if storage.has(session_id):
         raise wire.FailureError(UnexpectedMessage, 'Already initialized')
-    storage = Storage()
-    storage.imported = True
+
+    st = Storage()
+    st.imported = True
+    st.version = 1
+    st.pin = message.pin
+    st.passphrase_protection = message.passphrase_protection,
+    st.language = message.language
+    st.label = message.label
 
     if hasattr(message, 'node'):
-        storage.node = message.node
+        st.node = message.node
+
     elif hasattr(message, 'mnemonic'):
+        st.mnemonic = message.mnemonic
         if not message.skip_checksum and not bip39.check(message.mnemonic):
             raise wire.FailureError(Other, 'Mnemonic is not valid')
-        storage.mnemonic = message.mnemonic
-
-    ui.display.clear()
 
     await require_confirm(session_id, Text(
         'Loading seed',
         ui.BOLD, 'Loading private seed', 'is not recommended.',
         ui.NORMAL, 'Continue only if you', 'know what you are doing!'))
 
-    storage.version = 1
-    storage.pin = message.pin
-    storage.passphrase_protection = message.passphrase_protection,
-    storage.language = message.language
-    storage.label = message.label
-
-    set_storage(session_id, await storage.dumps())
+    storage.set(session_id, st)
 
     return Success(message='Device loaded')
diff --git a/src/apps/management/layout_reset_device.py b/src/apps/management/layout_reset_device.py
index 72ddd85d..997c7026 100644
--- a/src/apps/management/layout_reset_device.py
+++ b/src/apps/management/layout_reset_device.py
@@ -8,10 +8,10 @@ async def layout_reset_device(message, session_id):
     from trezor.messages.Success import Success
     from trezor.messages.Storage import Storage
     from trezor.messages.FailureType import UnexpectedMessage
-    from ..common.storage import get_storage, set_storage
     from ..common.request_pin import request_pin_twice
+    from ..common import storage
 
-    if get_storage(session_id):
+    if storage.has(session_id):
         raise wire.FailureError(UnexpectedMessage, 'Already initialized')
 
     mnemonic = await generate_mnemonic(
@@ -23,11 +23,10 @@ async def layout_reset_device(message, session_id):
     else:
         pin = ''
 
-    storage = Storage(
+    storage.set(session_id, Storage(
         version=1, pin=pin, mnemonic=mnemonic,
         passphrase_protection=message.passphrase_protection,
-        language=message.language, label=message.label)
-    set_storage(session_id, await storage.dumps())
+        language=message.language, label=message.label))
 
     return Success(message='Initialized')
 
diff --git a/src/apps/management/layout_wipe_device.py b/src/apps/management/layout_wipe_device.py
index 48aaa88b..a78f41c4 100644
--- a/src/apps/management/layout_wipe_device.py
+++ b/src/apps/management/layout_wipe_device.py
@@ -7,7 +7,7 @@ async def layout_wipe_device(message, session_id):
     from trezor.messages.Success import Success
     from trezor.ui.text import Text
     from ..common.confirm import hold_to_confirm
-    from ..common.storage import clear_storage
+    from ..common import storage
 
     ui.display.clear()
 
@@ -17,6 +17,6 @@ async def layout_wipe_device(message, session_id):
         ui.NORMAL, '', 'All data will be lost.')
     await hold_to_confirm(session_id, content)
 
-    clear_storage(session_id)
+    storage.clear(session_id)
 
     return Success(message='Device wiped')