From 601d3b49c337cea0fcc5dddb89a55b9aea4252e9 Mon Sep 17 00:00:00 2001
From: matejcik <ja@matejcik.cz>
Date: Fri, 2 Nov 2018 17:06:04 +0100
Subject: [PATCH] trezorlib: add some utility features

---
 trezorlib/protobuf.py | 13 ++++++++++++-
 trezorlib/tools.py    | 18 +++++++++++++++++-
 2 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/trezorlib/protobuf.py b/trezorlib/protobuf.py
index 19214fe99..c082e41af 100644
--- a/trezorlib/protobuf.py
+++ b/trezorlib/protobuf.py
@@ -148,7 +148,7 @@ class MessageType:
         return "<%s: %s>" % (self.__class__.__name__, d)
 
     def __iter__(self):
-        return self.__dict__.__iter__()
+        return iter(self.keys())
 
     def keys(self):
         return (name for name, _, _ in self.get_fields().values())
@@ -412,3 +412,14 @@ def dict_to_proto(message_type, d):
 
         params[fname] = newvalue
     return message_type(**params)
+
+
+def to_dict(msg):
+    res = {}
+    for key, value in msg.__dict__.items():
+        if value is None or value == []:
+            continue
+        if isinstance(value, MessageType):
+            value = to_dict(value)
+        res[key] = value
+    return res
diff --git a/trezorlib/tools.py b/trezorlib/tools.py
index 72f8a01ec..7851845aa 100644
--- a/trezorlib/tools.py
+++ b/trezorlib/tools.py
@@ -102,8 +102,11 @@ def b58encode(v):
     return (__b58chars[0] * nPad) + result
 
 
-def b58decode(v, length):
+def b58decode(v, length=None):
     """ decode v into a string of len bytes."""
+    if isinstance(v, bytes):
+        v = v.decode()
+
     long_value = 0
     for (i, c) in enumerate(v[::-1]):
         long_value += __b58chars.find(c) * (__b58base ** i)
@@ -129,6 +132,19 @@ def b58decode(v, length):
     return result
 
 
+def b58check_encode(v):
+    checksum = btc_hash(v)[:4]
+    return b58encode(v + checksum)
+
+
+def b58check_decode(v, length=None):
+    dec = b58decode(v, length)
+    data, checksum = dec[:-4], dec[-4:]
+    if btc_hash(data)[:4] != checksum:
+        raise ValueError("invalid checksum")
+    return data
+
+
 def parse_path(nstr: str) -> Address:
     """
     Convert BIP32 path string to list of uint32 integers with hardened flags.