Browse Source

pyln: Add type-annotations to pyln-proto

travis-experimental
Christian Decker 4 years ago
committed by Rusty Russell
parent
commit
cafaad741b
  1. 33
      contrib/pyln-proto/pyln/proto/bech32.py
  2. 6
      contrib/pyln-proto/pyln/proto/invoice.py
  3. 2
      contrib/pyln-proto/pyln/proto/zbase32.py

33
contrib/pyln-proto/pyln/proto/bech32.py

@ -19,12 +19,13 @@
# THE SOFTWARE. # THE SOFTWARE.
"""Reference implementation for Bech32 and segwit addresses.""" """Reference implementation for Bech32 and segwit addresses."""
from typing import Tuple
CHARSET = "qpzry9x8gf2tvdw0s3jn54khce6mua7l" CHARSET = "qpzry9x8gf2tvdw0s3jn54khce6mua7l"
def bech32_polymod(values): def bech32_polymod(values: bytes) -> int:
"""Internal function that computes the Bech32 checksum.""" """Internal function that computes the Bech32 checksum."""
generator = [0x3b6a57b2, 0x26508e6d, 0x1ea119fa, 0x3d4233dd, 0x2a1462b3] generator = [0x3b6a57b2, 0x26508e6d, 0x1ea119fa, 0x3d4233dd, 0x2a1462b3]
chk = 1 chk = 1
@ -36,43 +37,47 @@ def bech32_polymod(values):
return chk return chk
def bech32_hrp_expand(hrp): def bech32_hrp_expand(hrp: str) -> bytes:
"""Expand the HRP into values for checksum computation.""" """Expand the HRP into values for checksum computation."""
return [ord(x) >> 5 for x in hrp] + [0] + [ord(x) & 31 for x in hrp] return bytes([ord(x) >> 5 for x in hrp] + [0] + [ord(x) & 31 for x in hrp])
def bech32_verify_checksum(hrp, data): def bech32_verify_checksum(hrp: str, data: bytes) -> bool:
"""Verify a checksum given HRP and converted data characters.""" """Verify a checksum given HRP and converted data characters."""
return bech32_polymod(bech32_hrp_expand(hrp) + data) == 1 return bech32_polymod(bech32_hrp_expand(hrp) + data) == 1
def bech32_create_checksum(hrp, data): def bech32_create_checksum(hrp: str, data: bytes) -> bytes:
"""Compute the checksum values given HRP and data.""" """Compute the checksum values given HRP and data."""
values = bech32_hrp_expand(hrp) + data values = bech32_hrp_expand(hrp) + data
polymod = bech32_polymod(values + [0, 0, 0, 0, 0, 0]) ^ 1 polymod = bech32_polymod(values + bytes([0, 0, 0, 0, 0, 0])) ^ 1
return [(polymod >> 5 * (5 - i)) & 31 for i in range(6)] return bytes([(polymod >> 5 * (5 - i)) & 31 for i in range(6)])
def bech32_encode(hrp, data): def bech32_encode(hrp: str, data: bytes) -> str:
"""Compute a Bech32 string given HRP and data values.""" """Compute a Bech32 string given HRP and data values."""
combined = data + bech32_create_checksum(hrp, data) combined = data + bech32_create_checksum(hrp, data)
return hrp + '1' + ''.join([CHARSET[d] for d in combined]) return hrp + '1' + ''.join([CHARSET[d] for d in combined])
def bech32_decode(bech): def bech32_decode(bech: str) -> Tuple[str, bytes]:
"""Validate a Bech32 string, and determine HRP and data.""" """Validate a Bech32 string, and determine HRP and data."""
if ((any(ord(x) < 33 or ord(x) > 126 for x in bech)) or (bech.lower() != bech and bech.upper() != bech)): if ((any(ord(x) < 33 or ord(x) > 126 for x in bech)) or (bech.lower() != bech and bech.upper() != bech)):
return (None, None) raise ValueError("Not a bech32-encoded string: {}".format(bech))
bech = bech.lower() bech = bech.lower()
pos = bech.rfind('1') pos = bech.rfind('1')
if pos < 1 or pos + 7 > len(bech): if pos < 1 or pos + 7 > len(bech):
return (None, None) raise ValueError("Could not locate hrp separator '1' in {}".format(bech))
if not all(x in CHARSET for x in bech[pos + 1:]): if not all(x in CHARSET for x in bech[pos + 1:]):
return (None, None) raise ValueError("Non-bech32 character found in {}".format(bech))
hrp = bech[:pos] hrp = bech[:pos]
data = [CHARSET.find(x) for x in bech[pos + 1:]] data = bytes([CHARSET.find(x) for x in bech[pos + 1:]])
if not bech32_verify_checksum(hrp, data): if not bech32_verify_checksum(hrp, data):
return (None, None) raise ValueError("Checksum verification failed for {}".format(bech))
return (hrp, data[:-6]) return (hrp, data[:-6])

6
contrib/pyln-proto/pyln/proto/invoice.py

@ -59,20 +59,20 @@ def unshorten_amount(amount):
# Bech32 spits out array of 5-bit values. Shim here. # Bech32 spits out array of 5-bit values. Shim here.
def u5_to_bitarray(arr): def u5_to_bitarray(arr: bytes):
ret = bitstring.BitArray() ret = bitstring.BitArray()
for a in arr: for a in arr:
ret += bitstring.pack("uint:5", a) ret += bitstring.pack("uint:5", a)
return ret return ret
def bitarray_to_u5(barr): def bitarray_to_u5(barr) -> bytes:
assert barr.len % 5 == 0 assert barr.len % 5 == 0
ret = [] ret = []
s = bitstring.ConstBitStream(barr) s = bitstring.ConstBitStream(barr)
while s.pos != s.len: while s.pos != s.len:
ret.append(s.read(5).uint) ret.append(s.read(5).uint)
return ret return bytes(ret)
def encode_fallback(fallback, currency): def encode_fallback(fallback, currency):

2
contrib/pyln-proto/pyln/proto/zbase32.py

@ -1,4 +1,4 @@
import bitstring import bitstring # type: ignore
zbase32_chars = b'ybndrfg8ejkmcpqxot1uwisza345h769' zbase32_chars = b'ybndrfg8ejkmcpqxot1uwisza345h769'

Loading…
Cancel
Save