pyln-proto: Added a couple of utilities to manage onions and zbase32

contrib/pyln-proto/pyln/proto/onion.py

@@ -0,0 +1,236 @@
+from .primitives import varint_decode, varint_encode
+from io import BytesIO, SEEK_CUR
+from binascii import hexlify, unhexlify
+import struct
+
+
+class OnionPayload(object):
+
+    @classmethod
+    def from_bytes(cls, b):
+        if isinstance(b, bytes):
+            b = BytesIO(b)
+
+        realm = b.read(1)
+        b.seek(-1, SEEK_CUR)
+
+        if realm == b'\x00':
+            return LegacyOnionPayload.from_bytes(b)
+        elif realm != b'\x01':
+            return TlvPayload.from_bytes(b, skip_length=False)
+        else:
+            raise ValueError("Onion payloads with realm 0x01 are unsupported")
+
+    @classmethod
+    def from_hex(cls, s):
+        if isinstance(s, str):
+            s = s.encode('ASCII')
+        return cls.from_bytes(bytes(unhexlify(s)))
+
+    def to_bytes(self):
+        raise ValueError("OnionPayload is an abstract class, use "
+                         "LegacyOnionPayload or TlvPayload instead")
+
+    def to_hex(self):
+        return hexlify(self.to_bytes()).decode('ASCII')
+
+
+class LegacyOnionPayload(OnionPayload):
+
+    def __init__(self, amt_to_forward, outgoing_cltv_value,
+                 short_channel_id=None, padding=None):
+        assert(padding is None or len(padding) == 12)
+        self.padding = b'\x00' * 12 if padding is None else padding
+
+        if isinstance(amt_to_forward, str):
+            self.amt_to_forward = int(amt_to_forward)
+        else:
+            self.amt_to_forward = amt_to_forward
+
+        self.outgoing_cltv_value = outgoing_cltv_value
+
+        if isinstance(short_channel_id, str) and 'x' in short_channel_id:
+            # Convert the short_channel_id from its string representation to its numeric representation
+            block, tx, out = short_channel_id.split('x')
+            num_scid = int(block) << 40 | int(tx) << 16 | int(out)
+            self.short_channel_id = num_scid
+        elif isinstance(short_channel_id, int):
+            self.short_channel_id = short_channel_id
+        else:
+            raise ValueError("short_channel_id format cannot be recognized: {}".format(short_channel_id))
+
+    @classmethod
+    def from_bytes(cls, b):
+        if isinstance(b, bytes):
+            b = BytesIO(b)
+
+        assert(b.read(1) == b'\x00')
+
+        s, a, o = struct.unpack("!QQL", b.read(20))
+        padding = b.read(12)
+        return LegacyOnionPayload(a, o, s, padding)
+
+    def to_bytes(self, include_realm=True):
+        b = b''
+        if include_realm:
+            b += b'\x00'
+
+        b += struct.pack("!Q", self.short_channel_id)
+        b += struct.pack("!Q", self.amt_to_forward)
+        b += struct.pack("!L", self.outgoing_cltv_value)
+        b += self.padding
+        assert(len(b) == 32 + include_realm)
+        return b
+
+    def to_hex(self, include_realm=True):
+        return hexlify(self.to_bytes(include_realm)).decode('ASCII')
+
+    def __str__(self):
+        return ("LegacyOnionPayload[scid={self.short_channel_id}, "
+                "amt_to_forward={self.amt_to_forward}, "
+                "outgoing_cltv={self.outgoing_cltv_value}]").format(self=self)
+
+
+class TlvPayload(OnionPayload):
+
+    def __init__(self, fields=None):
+        self.fields = [] if fields is None else fields
+
+    @classmethod
+    def from_bytes(cls, b, skip_length=False):
+        if isinstance(b, str):
+            b = b.encode('ASCII')
+        if isinstance(b, bytes):
+            b = BytesIO(b)
+
+        if skip_length:
+            # Consume the entire remainder of the buffer.
+            payload_length = len(b.getvalue()) - b.tell()
+        else:
+            payload_length = varint_decode(b)
+
+        instance = TlvPayload()
+
+        start = b.tell()
+        while b.tell() < start + payload_length:
+            typenum = varint_decode(b)
+            if typenum is None:
+                break
+            length = varint_decode(b)
+            if length is None:
+                raise ValueError(
+                    "Unable to read length at position {}".format(b.tell())
+                )
+            val = b.read(length)
+
+            # Get the subclass that is the correct interpretation of this
+            # field. Default to the binary field type.
+            c = tlv_types.get(typenum, (TlvField, "unknown"))
+            cls = c[0]
+            field = cls.from_bytes(typenum=typenum, b=val, description=c[1])
+            instance.fields.append(field)
+
+        return instance
+
+    @classmethod
+    def from_hex(cls, h):
+        return cls.from_bytes(unhexlify(h))
+
+    def add_field(self, typenum, value):
+        self.fields.append(TlvField(typenum=typenum, value=value))
+
+    def get(self, key, default=None):
+        for f in self.fields:
+            if f.typenum == key:
+                return f
+        return default
+
+    def to_bytes(self):
+        ser = [f.to_bytes() for f in self.fields]
+        b = BytesIO()
+        varint_encode(sum([len(b) for b in ser]), b)
+        for f in ser:
+            b.write(f)
+        return b.getvalue()
+
+    def __str__(self):
+        return "TlvPayload[" + ', '.join([str(f) for f in self.fields]) + "]"
+
+
+class TlvField(object):
+
+    def __init__(self, typenum, value=None, description=None):
+        self.typenum = typenum
+        self.value = value
+        self.description = description
+
+    @classmethod
+    def from_bytes(cls, typenum, b, description=None):
+        return TlvField(typenum=typenum, value=b, description=description)
+
+    def __str__(self):
+        return "TlvField[{description},{num}={hex}]".format(
+            description=self.description,
+            num=self.typenum,
+            hex=hexlify(self.value).decode('ASCII')
+        )
+
+    def to_bytes(self):
+        b = BytesIO()
+        varint_encode(self.typenum, b)
+        varint_encode(len(self.value), b)
+        b.write(self.value)
+        return b.getvalue()
+
+
+class Tu32Field(TlvField):
+    pass
+
+
+class Tu64Field(TlvField):
+    pass
+
+
+class ShortChannelIdField(TlvField):
+    pass
+
+
+class TextField(TlvField):
+
+    @classmethod
+    def from_bytes(cls, typenum, b, description=None):
+        val = b.decode('UTF-8')
+        return TextField(typenum, value=val, description=description)
+
+    def to_bytes(self):
+        b = BytesIO()
+        val = self.value.encode('UTF-8')
+        varint_encode(self.typenum, b)
+        varint_encode(len(val), b)
+        b.write(val)
+        return b.getvalue()
+
+    def __str__(self):
+        return "TextField[{description},{num}=\"{val}\"]".format(
+            description=self.description,
+            num=self.typenum,
+            val=self.value,
+        )
+
+
+class HashField(TlvField):
+    pass
+
+
+class SignatureField(TlvField):
+    pass
+
+
+# A mapping of known TLV types
+tlv_types = {
+    2: (Tu64Field, 'amt_to_forward'),
+    4: (Tu32Field, 'outgoing_cltv_value'),
+    6: (ShortChannelIdField, 'short_channel_id'),
+    34349334: (TextField, 'noise_message_body'),
+    34349336: (SignatureField, 'noise_message_signature'),
+}