From fc5bda43d04fcccd557b056ad2e2b6bd79d10e1a Mon Sep 17 00:00:00 2001
From: Kristaps Kaupe <kristaps@blogiem.lv>
Date: Wed, 11 May 2022 16:37:29 +0300
Subject: [PATCH] Basic output descriptor functions (generate only)

---
 jmbitcoin/jmbitcoin/__init__.py           |  2 +
 jmbitcoin/jmbitcoin/output_descriptors.py | 85 +++++++++++++++++++++++
 jmbitcoin/test/test_output_descriptors.py | 23 ++++++
 3 files changed, 110 insertions(+)
 create mode 100644 jmbitcoin/jmbitcoin/output_descriptors.py
 create mode 100644 jmbitcoin/test/test_output_descriptors.py

diff --git a/jmbitcoin/jmbitcoin/__init__.py b/jmbitcoin/jmbitcoin/__init__.py
index 6efa373..e4fed6e 100644
--- a/jmbitcoin/jmbitcoin/__init__.py
+++ b/jmbitcoin/jmbitcoin/__init__.py
@@ -21,6 +21,8 @@ from jmbitcoin.secp256k1_deterministic import *
 from jmbitcoin.snicker import *
 from jmbitcoin.amount import *
 from jmbitcoin.bip21 import *
+from jmbitcoin.output_descriptors import (get_address_descriptor,
+                                          get_xpub_descriptor)
 from bitcointx import select_chain_params
 from bitcointx.core import (x, b2x, b2lx, lx, COutPoint, CTxOut, CTxIn,
                             CTxInWitness, CTxWitness, CTransaction,
diff --git a/jmbitcoin/jmbitcoin/output_descriptors.py b/jmbitcoin/jmbitcoin/output_descriptors.py
new file mode 100644
index 0000000..d590234
--- /dev/null
+++ b/jmbitcoin/jmbitcoin/output_descriptors.py
@@ -0,0 +1,85 @@
+
+# Descriptor checksum calculation code taken from
+# https://github.com/bitcoin-core/HWI/blob/master/hwilib/descriptor.py
+
+def poly_mod(c: int, val: int) -> int:
+    """
+    :meta private:
+    Function to compute modulo over the polynomial used for descriptor checksums
+    From: https://github.com/bitcoin/bitcoin/blob/master/src/script/descriptor.cpp
+    """
+    c0 = c >> 35
+    c = ((c & 0x7ffffffff) << 5) ^ val
+    if (c0 & 1):
+        c ^= 0xf5dee51989
+    if (c0 & 2):
+        c ^= 0xa9fdca3312
+    if (c0 & 4):
+        c ^= 0x1bab10e32d
+    if (c0 & 8):
+        c ^= 0x3706b1677a
+    if (c0 & 16):
+        c ^= 0x644d626ffd
+    return c
+
+
+def descriptor_checksum(desc: str) -> str:
+    """
+    Compute the checksum for a descriptor
+    :param desc: The descriptor string to compute a checksum for
+    :return: A checksum
+    """
+    INPUT_CHARSET = "0123456789()[],'/*abcdefgh@:$%{}IJKLMNOPQRSTUVWXYZ&+-.;<=>?!^_|~ijklmnopqrstuvwxyzABCDEFGH`#\"\\ "
+    CHECKSUM_CHARSET = "qpzry9x8gf2tvdw0s3jn54khce6mua7l"
+
+    c = 1
+    cls = 0
+    clscount = 0
+    for ch in desc:
+        pos = INPUT_CHARSET.find(ch)
+        if pos == -1:
+            return ""
+        c = poly_mod(c, pos & 31)
+        cls = cls * 3 + (pos >> 5)
+        clscount += 1
+        if clscount == 3:
+            c = poly_mod(c, cls)
+            cls = 0
+            clscount = 0
+    if clscount > 0:
+        c = poly_mod(c, cls)
+    for j in range(0, 8):
+        c = poly_mod(c, 0)
+    c ^= 1
+
+    ret = [''] * 8
+    for j in range(0, 8):
+        ret[j] = CHECKSUM_CHARSET[(c >> (5 * (7 - j))) & 31]
+    return ''.join(ret)
+
+
+def add_checksum(desc: str) -> str:
+    """
+    Compute and attach the checksum for a descriptor
+    :param desc: The descriptor string to add a checksum to
+    :return: Descriptor with checksum
+    """
+    return desc + "#" + descriptor_checksum(desc)
+
+
+def get_address_descriptor(address: str) -> str:
+    return add_checksum("addr({})".format(address))
+
+
+def get_xpub_descriptor(xpub_key: str, address_type: str) -> str:
+    if address_type == "p2pkh":
+        function = "pkh"
+    elif address_type == "p2sh-p2wpkh" or address_type == "p2wpkh":
+        function = "wpkh"
+    else:
+        raise Exception("Unsupported address type {}".format(address_type))
+    descriptor = "{}({}/*)".format(function, xpub_key)
+    if address_type == "p2sh-p2wpkh":
+        descriptor = "sh({})".format(descriptor)
+    return add_checksum(descriptor)
+
diff --git a/jmbitcoin/test/test_output_descriptors.py b/jmbitcoin/test/test_output_descriptors.py
new file mode 100644
index 0000000..a35deba
--- /dev/null
+++ b/jmbitcoin/test/test_output_descriptors.py
@@ -0,0 +1,23 @@
+import jmbitcoin as btc
+
+
+def test_address_descriptors():
+    assert(btc.get_address_descriptor("1AGNa15ZQXAZUgFiqJ2i7Z2DPU2J6hW62i") ==
+        "addr(1AGNa15ZQXAZUgFiqJ2i7Z2DPU2J6hW62i)#ns3f5w84")
+    assert(btc.get_address_descriptor("3CMNFxN1oHBc4R1EpboAL5yzHGgE611Xou") ==
+        "addr(3CMNFxN1oHBc4R1EpboAL5yzHGgE611Xou)#swk5gt6w")
+    assert(btc.get_address_descriptor("bc1qt493axn3wl4gzjxvfg03vkacre0m6f2gzfhv5t") ==
+        "addr(bc1qt493axn3wl4gzjxvfg03vkacre0m6f2gzfhv5t)#q8mdrmlw")
+
+
+def test_xpub_descriptors():
+    assert(btc.get_xpub_descriptor(
+        "xpub6CMAJ67vZWVXuzjzYXUoJgWrmuvFRiqiUG4dwoXNFmJtpTH3WgviANNxGyZYo27zxbMuqhDDym6fnBxmGaYoxr6LHgNDo1eEghkXHTX4Jnx", "p2pkh") ==
+        "pkh(xpub6CMAJ67vZWVXuzjzYXUoJgWrmuvFRiqiUG4dwoXNFmJtpTH3WgviANNxGyZYo27zxbMuqhDDym6fnBxmGaYoxr6LHgNDo1eEghkXHTX4Jnx/*)#flej8438")
+    assert(btc.get_xpub_descriptor(
+        "xpub6CMAJ67vZWVXuzjzYXUoJgWrmuvFRiqiUG4dwoXNFmJtpTH3WgviANNxGyZYo27zxbMuqhDDym6fnBxmGaYoxr6LHgNDo1eEghkXHTX4Jnx", "p2sh-p2wpkh") ==
+        "sh(wpkh(xpub6CMAJ67vZWVXuzjzYXUoJgWrmuvFRiqiUG4dwoXNFmJtpTH3WgviANNxGyZYo27zxbMuqhDDym6fnBxmGaYoxr6LHgNDo1eEghkXHTX4Jnx/*))#f2940w8j")
+    assert(btc.get_xpub_descriptor(
+        "xpub6CMAJ67vZWVXuzjzYXUoJgWrmuvFRiqiUG4dwoXNFmJtpTH3WgviANNxGyZYo27zxbMuqhDDym6fnBxmGaYoxr6LHgNDo1eEghkXHTX4Jnx", "p2wpkh") ==
+        "wpkh(xpub6CMAJ67vZWVXuzjzYXUoJgWrmuvFRiqiUG4dwoXNFmJtpTH3WgviANNxGyZYo27zxbMuqhDDym6fnBxmGaYoxr6LHgNDo1eEghkXHTX4Jnx/*)#keekqdjc")
+