1  #!/usr/bin/env python3
  2  # Copyright (c) 2017 Pieter Wuille
  3  # Distributed under the MIT software license, see the accompanying
  4  # file COPYING or http://www.opensource.org/licenses/mit-license.php.
  5  """Reference implementation for Bech32/Bech32m and segwit addresses."""
  6  import unittest
  7  from enum import Enum
  8  
  9  CHARSET = "qpzry9x8gf2tvdw0s3jn54khce6mua7l"
 10  BECH32_CONST = 1
 11  BECH32M_CONST = 0x2bc830a3
 12  
 13  class Encoding(Enum):
 14      """Enumeration type to list the various supported encodings."""
 15      BECH32 = 1
 16      BECH32M = 2
 17  
 18  
 19  def bech32_polymod(values):
 20      """Internal function that computes the Bech32 checksum."""
 21      generator = [0x3b6a57b2, 0x26508e6d, 0x1ea119fa, 0x3d4233dd, 0x2a1462b3]
 22      chk = 1
 23      for value in values:
 24          top = chk >> 25
 25          chk = (chk & 0x1ffffff) << 5 ^ value
 26          for i in range(5):
 27              chk ^= generator[i] if ((top >> i) & 1) else 0
 28      return chk
 29  
 30  
 31  def bech32_hrp_expand(hrp):
 32      """Expand the HRP into values for checksum computation."""
 33      return [ord(x) >> 5 for x in hrp] + [0] + [ord(x) & 31 for x in hrp]
 34  
 35  
 36  def bech32_verify_checksum(hrp, data):
 37      """Verify a checksum given HRP and converted data characters."""
 38      check = bech32_polymod(bech32_hrp_expand(hrp) + data)
 39      if check == BECH32_CONST:
 40          return Encoding.BECH32
 41      elif check == BECH32M_CONST:
 42          return Encoding.BECH32M
 43      else:
 44          return None
 45  
 46  def bech32_create_checksum(encoding, hrp, data):
 47      """Compute the checksum values given HRP and data."""
 48      values = bech32_hrp_expand(hrp) + data
 49      const = BECH32M_CONST if encoding == Encoding.BECH32M else BECH32_CONST
 50      polymod = bech32_polymod(values + [0, 0, 0, 0, 0, 0]) ^ const
 51      return [(polymod >> 5 * (5 - i)) & 31 for i in range(6)]
 52  
 53  
 54  def bech32_encode(encoding, hrp, data):
 55      """Compute a Bech32 or Bech32m string given HRP and data values."""
 56      combined = data + bech32_create_checksum(encoding, hrp, data)
 57      return hrp + '1' + ''.join([CHARSET[d] for d in combined])
 58  
 59  
 60  def bech32_decode(bech):
 61      """Validate a Bech32/Bech32m string, and determine HRP and data."""
 62      if ((any(ord(x) < 33 or ord(x) > 126 for x in bech)) or
 63              (bech.lower() != bech and bech.upper() != bech)):
 64          return (None, None, None)
 65      bech = bech.lower()
 66      pos = bech.rfind('1')
 67      if pos < 1 or pos + 7 > len(bech) or len(bech) > 90:
 68          return (None, None, None)
 69      if not all(x in CHARSET for x in bech[pos+1:]):
 70          return (None, None, None)
 71      hrp = bech[:pos]
 72      data = [CHARSET.find(x) for x in bech[pos+1:]]
 73      encoding = bech32_verify_checksum(hrp, data)
 74      if encoding is None:
 75          return (None, None, None)
 76      return (encoding, hrp, data[:-6])
 77  
 78  
 79  def convertbits(data, frombits, tobits, pad=True):
 80      """General power-of-2 base conversion."""
 81      acc = 0
 82      bits = 0
 83      ret = []
 84      maxv = (1 << tobits) - 1
 85      max_acc = (1 << (frombits + tobits - 1)) - 1
 86      for value in data:
 87          if value < 0 or (value >> frombits):
 88              return None
 89          acc = ((acc << frombits) | value) & max_acc
 90          bits += frombits
 91          while bits >= tobits:
 92              bits -= tobits
 93              ret.append((acc >> bits) & maxv)
 94      if pad:
 95          if bits:
 96              ret.append((acc << (tobits - bits)) & maxv)
 97      elif bits >= frombits or ((acc << (tobits - bits)) & maxv):
 98          return None
 99      return ret
100  
101  
102  def decode_segwit_address(hrp, addr):
103      """Decode a segwit address."""
104      encoding, hrpgot, data = bech32_decode(addr)
105      if hrpgot != hrp:
106          return (None, None)
107      decoded = convertbits(data[1:], 5, 8, False)
108      if decoded is None or len(decoded) < 2 or len(decoded) > 40:
109          return (None, None)
110      if data[0] > 16:
111          return (None, None)
112      if data[0] == 0 and len(decoded) != 20 and len(decoded) != 32:
113          return (None, None)
114      if (data[0] == 0 and encoding != Encoding.BECH32) or (data[0] != 0 and encoding != Encoding.BECH32M):
115          return (None, None)
116      return (data[0], decoded)
117  
118  
119  def encode_segwit_address(hrp, witver, witprog):
120      """Encode a segwit address."""
121      encoding = Encoding.BECH32 if witver == 0 else Encoding.BECH32M
122      ret = bech32_encode(encoding, hrp, [witver] + convertbits(witprog, 8, 5))
123      if decode_segwit_address(hrp, ret) == (None, None):
124          return None
125      return ret
126  
127  class TestFrameworkScript(unittest.TestCase):
128      def test_segwit_encode_decode(self):
129          def test_python_bech32(addr):
130              hrp = addr[:4]
131              self.assertEqual(hrp, "bcrt")
132              (witver, witprog) = decode_segwit_address(hrp, addr)
133              self.assertEqual(encode_segwit_address(hrp, witver, witprog), addr)
134  
135          # P2WPKH
136          test_python_bech32('bcrt1qthmht0k2qnh3wy7336z05lu2km7emzfpm3wg46')
137          # P2WSH
138          test_python_bech32('bcrt1qqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqqq3xueyj')
139          test_python_bech32('bcrt1qft5p2uhsdcdc3l2ua4ap5qqfg4pjaqlp250x7us7a8qqhrxrxfsqseac85')
140          # P2TR
141          test_python_bech32('bcrt1p0xlxvlhemja6c4dqv22uapctqupfhlxm9h8z3k2e72q4k9hcz7vqc8gma6')