Add descriptor class

bitcoin-core · Feb 2, 2019 · 1826049 · 1826049
1 parent d78fc1f
commit 1826049
Show file tree

Hide file tree

Showing 3 changed files with 158 additions and 0 deletions.
diff --git a/hwilib/descriptor.py b/hwilib/descriptor.py
@@ -0,0 +1,76 @@
+import re
+
+class Descriptor:
+    def __init__(self, origin_fingerprint, origin_path, base_key, path_suffix, testnet, sh_wpkh, wpkh):
+        self.origin_fingerprint = origin_fingerprint
+        self.origin_path = origin_path
+        self.path_suffix = path_suffix
+        self.base_key = base_key
+        self.testnet = testnet
+        self.sh_wpkh = sh_wpkh
+        self.wpkh = wpkh
+        self.m_path = None
+
+        if origin_path:
+            self.m_path_base = "m" + origin_path
+            self.m_path = "m" + origin_path + (path_suffix or "")
+
+    @classmethod
+    def parse(cls, desc, testnet = False):
+        sh_wpkh = None
+        wpkh = None
+        origin_fingerprint = None
+        origin_path = None
+        base_key_and_path_match = None
+        base_key = None
+
+        if desc.startswith("sh(wpkh("):
+            sh_wpkh = True
+        elif desc.startswith("wpkh("):
+            wpkh = True
+
+        origin_match = re.search(r"\[(.*)\]", desc)
+        if origin_match:
+            origin = origin_match.group(1)
+            match = re.search(r"^([0-9a-fA-F]{8})(\/.*)", origin)
+            if  match:
+              origin_fingerprint = match.group(1)
+              origin_path = match.group(2)
+              # Replace h with '
+              origin_path = origin_path.replace('h', '\'')
+
+            base_key_and_path_match = re.search(r"\[.*\](\w+)([\/\)][\d'\/\*]*)", desc)
+        else:
+            base_key_and_path_match = re.search(r"\((\w+)([\/\)][\d'\/\*]*)", desc)
+
+        if base_key_and_path_match:
+            base_key = base_key_and_path_match.group(1)
+            path_suffix = base_key_and_path_match.group(2)
+            if path_suffix == ")":
+                path_suffix = None
+        else:
+            if origin_match == None:
+                return None
+
+        return cls(origin_fingerprint, origin_path, base_key, path_suffix, testnet, sh_wpkh, wpkh)
+
+
+    def serialize(self):
+        descriptor_open = 'pkh('
+        descriptor_close = ')'
+        origin = ''
+        path_suffix = ''
+
+        if self.wpkh == True:
+            descriptor_open = 'wpkh('
+        elif self.sh_wpkh == True:
+            descriptor_open = 'sh(wpkh('
+            descriptor_close = '))'
+
+        if self.origin_fingerprint and self.origin_path:
+            origin = '[' + self.origin_fingerprint + self.origin_path + ']'
+
+        if self.path_suffix:
+            path_suffix = self.path_suffix
+
+        return descriptor_open + origin + self.base_key + path_suffix + descriptor_close
diff --git a/test/run_tests.py b/test/run_tests.py
@@ -7,6 +7,7 @@
 
 from test_bech32 import TestSegwitAddress
 from test_coldcard import coldcard_test_suite
+from test_descriptor import TestDescriptor
 from test_device import start_bitcoind
 from test_psbt import TestPSBT
 from test_trezor import trezor_test_suite
@@ -35,6 +36,7 @@
 
 # Run tests
 suite = unittest.TestSuite()
+suite.addTests(unittest.defaultTestLoader.loadTestsFromTestCase(TestDescriptor))
 suite.addTests(unittest.defaultTestLoader.loadTestsFromTestCase(TestSegwitAddress))
 suite.addTests(unittest.defaultTestLoader.loadTestsFromTestCase(TestPSBT))
 

diff --git a/test/test_descriptor.py b/test/test_descriptor.py
@@ -0,0 +1,80 @@
+#! /usr/bin/env python3
+
+from hwilib.descriptor import Descriptor
+import unittest
+
+class TestDescriptor(unittest.TestCase):
+    def test_parse_descriptor_with_origin(self):
+        desc = Descriptor.parse("wpkh([00000001/84'/1'/0']tpubD6NzVbkrYhZ4WaWSyoBvQwbpLkojyoTZPRsgXELWz3Popb3qkjcJyJUGLnL4qHHoQvao8ESaAstxYSnhyswJ76uZPStJRJCTKvosUCJZL5B/0/0)", True)
+        self.assertIsNotNone(desc)
+        self.assertEqual(desc.wpkh, True)
+        self.assertEqual(desc.sh_wpkh, None)
+        self.assertEqual(desc.origin_fingerprint, "00000001")
+        self.assertEqual(desc.origin_path, "/84'/1'/0'")
+        self.assertEqual(desc.base_key, "tpubD6NzVbkrYhZ4WaWSyoBvQwbpLkojyoTZPRsgXELWz3Popb3qkjcJyJUGLnL4qHHoQvao8ESaAstxYSnhyswJ76uZPStJRJCTKvosUCJZL5B")
+        self.assertEqual(desc.path_suffix, "/0/0")
+        self.assertEqual(desc.testnet, True)
+        self.assertEqual(desc.m_path, "m/84'/1'/0'/0/0")
+
+    def test_parse_descriptor_without_origin(self):
+        desc = Descriptor.parse("wpkh(tpubD6NzVbkrYhZ4WaWSyoBvQwbpLkojyoTZPRsgXELWz3Popb3qkjcJyJUGLnL4qHHoQvao8ESaAstxYSnhyswJ76uZPStJRJCTKvosUCJZL5B/0/0)", True)
+        self.assertIsNotNone(desc)
+        self.assertEqual(desc.wpkh, True)
+        self.assertEqual(desc.sh_wpkh, None)
+        self.assertEqual(desc.origin_fingerprint, None)
+        self.assertEqual(desc.origin_path, None)
+        self.assertEqual(desc.base_key, "tpubD6NzVbkrYhZ4WaWSyoBvQwbpLkojyoTZPRsgXELWz3Popb3qkjcJyJUGLnL4qHHoQvao8ESaAstxYSnhyswJ76uZPStJRJCTKvosUCJZL5B")
+        self.assertEqual(desc.path_suffix, "/0/0")
+        self.assertEqual(desc.testnet, True)
+        self.assertEqual(desc.m_path, None)
+
+    def test_parse_descriptor_with_key_at_end_with_origin(self):
+        desc = Descriptor.parse("wpkh([00000001/84'/1'/0'/0/0]0297dc3f4420402e01a113984311bf4a1b8de376cac0bdcfaf1b3ac81f13433c7)", True)
+        self.assertIsNotNone(desc)
+        self.assertEqual(desc.wpkh, True)
+        self.assertEqual(desc.sh_wpkh, None)
+        self.assertEqual(desc.origin_fingerprint, "00000001")
+        self.assertEqual(desc.origin_path, "/84'/1'/0'/0/0")
+        self.assertEqual(desc.base_key, "0297dc3f4420402e01a113984311bf4a1b8de376cac0bdcfaf1b3ac81f13433c7")
+        self.assertEqual(desc.path_suffix, None)
+        self.assertEqual(desc.testnet, True)
+        self.assertEqual(desc.m_path, "m/84'/1'/0'/0/0")
+
+    def test_parse_descriptor_with_key_at_end_without_origin(self):
+        desc = Descriptor.parse("wpkh(0297dc3f4420402e01a113984311bf4a1b8de376cac0bdcfaf1b3ac81f13433c7)", True)
+        self.assertIsNotNone(desc)
+        self.assertEqual(desc.wpkh, True)
+        self.assertEqual(desc.sh_wpkh, None)
+        self.assertEqual(desc.origin_fingerprint, None)
+        self.assertEqual(desc.origin_path, None)
+        self.assertEqual(desc.base_key, "0297dc3f4420402e01a113984311bf4a1b8de376cac0bdcfaf1b3ac81f13433c7")
+        self.assertEqual(desc.path_suffix, None)
+        self.assertEqual(desc.testnet, True)
+        self.assertEqual(desc.m_path, None)
+
+    def test_parse_empty_descriptor(self):
+        desc = Descriptor.parse("", True)
+        self.assertIsNone(desc)
+
+    def test_parse_descriptor_replace_h(self):
+        desc = Descriptor.parse("wpkh([00000001/84h/1h/0']tpubD6NzVbkrYhZ4WaWSyoBvQwbpLkojyoTZPRsgXELWz3Popb3qkjcJyJUGLnL4qHHoQvao8ESaAstxYSnhyswJ76uZPStJRJCTKvosUCJZL5B/0/0)", True)
+        self.assertIsNotNone(desc)
+        self.assertEqual(desc.origin_path, "/84'/1'/0'")
+
+    def test_serialize_descriptor_with_origin(self):
+        descriptor = "wpkh([00000001/84'/1'/0']tpubD6NzVbkrYhZ4WaWSyoBvQwbpLkojyoTZPRsgXELWz3Popb3qkjcJyJUGLnL4qHHoQvao8ESaAstxYSnhyswJ76uZPStJRJCTKvosUCJZL5B/0/0)"
+        desc = Descriptor.parse(descriptor, True)
+        self.assertEqual(desc.serialize(), descriptor)
+
+    def test_serialize_descriptor_without_origin(self):
+        descriptor = "wpkh(tpubD6NzVbkrYhZ4WaWSyoBvQwbpLkojyoTZPRsgXELWz3Popb3qkjcJyJUGLnL4qHHoQvao8ESaAstxYSnhyswJ76uZPStJRJCTKvosUCJZL5B/0/0)"
+        desc = Descriptor.parse(descriptor, True)
+        self.assertEqual(desc.serialize(), descriptor)
+
+    def test_serialize_descriptor_with_key_at_end_with_origin(self):
+        descriptor = "wpkh([00000001/84'/1'/0'/0/0]0297dc3f4420402e01a113984311bf4a1b8de376cac0bdcfaf1b3ac81f13433c7)"
+        desc = Descriptor.parse(descriptor, True)
+        self.assertEqual(desc.serialize(), descriptor)
+
+if __name__ == "__main__":
+    unittest.main()