move transport code to its own file

electrum/lnbase.py

@@ -9,18 +9,14 @@ import json
 import asyncio
 import os
 import time
-import hashlib
-import hmac
 from functools import partial
 from typing import List
 
-import cryptography.hazmat.primitives.ciphers.aead as AEAD
 import aiorpcx
 
 from . import bitcoin
 from . import ecc
 from .ecc import sig_string_from_r_and_s, get_r_and_s_from_sig_string
-from .crypto import sha256
 from . import constants
 from .util import PrintError, bh2u, print_error, bfh, log_exceptions, list_enabled_bits, ignore_exceptions
 from .transaction import Transaction, TxOutput
@@ -29,10 +25,11 @@ from .lnaddr import lndecode
 from .lnchan import Channel, RevokeAndAck, htlcsum
 from .lnutil import (Outpoint, LocalConfig, ChannelConfig,
                      RemoteConfig, OnlyPubkeyKeypair, ChannelConstraints, RevocationStore,
-                     funding_output_script, get_ecdh, get_per_commitment_secret_from_seed,
+                     funding_output_script, get_per_commitment_secret_from_seed,
                      secret_to_pubkey, LNPeerAddr, PaymentFailure, LnLocalFeatures,
                      LOCAL, REMOTE, HTLCOwner, generate_keypair, LnKeyFamily,
                      get_ln_flag_pair_of_bit)
+from .lnutil import LightningPeerConnectionClosed, HandshakeFailed
 from .lnrouter import NotFoundChanAnnouncementForUpdate, RouteEdge
 
 
@@ -41,11 +38,6 @@ def channel_id_from_funding_tx(funding_txid, funding_index):
     i = int.from_bytes(funding_txid_bytes, 'big') ^ funding_index
     return i.to_bytes(32, 'big'), funding_txid_bytes
 
-class LightningError(Exception):
-    pass
-
-class LightningPeerConnectionClosed(LightningError):
-    pass
 
 message_types = {}
 
@@ -192,174 +184,6 @@ def gen_msg(msg_type, **kwargs):
     return data
 
 
-class HandshakeState(object):
-    prologue = b"lightning"
-    protocol_name = b"Noise_XK_secp256k1_ChaChaPoly_SHA256"
-    handshake_version = b"\x00"
-
-    def __init__(self, responder_pub):
-        self.responder_pub = responder_pub
-        self.h = sha256(self.protocol_name)
-        self.ck = self.h
-        self.update(self.prologue)
-        self.update(self.responder_pub)
-
-    def update(self, data):
-        self.h = sha256(self.h + data)
-        return self.h
-
-
-class HandshakeFailed(Exception): pass
-
-
-def get_nonce_bytes(n):
-    """BOLT 8 requires the nonce to be 12 bytes, 4 bytes leading
-    zeroes and 8 bytes little endian encoded 64 bit integer.
-    """
-    return b"\x00"*4 + n.to_bytes(8, 'little')
-
-def aead_encrypt(k, nonce, associated_data, data):
-    nonce_bytes = get_nonce_bytes(nonce)
-    a = AEAD.ChaCha20Poly1305(k)
-    return a.encrypt(nonce_bytes, data, associated_data)
-
-def aead_decrypt(k, nonce, associated_data, data):
-    nonce_bytes = get_nonce_bytes(nonce)
-    a = AEAD.ChaCha20Poly1305(k)
-    #raises InvalidTag exception if it's not valid
-    return a.decrypt(nonce_bytes, data, associated_data)
-
-def get_bolt8_hkdf(salt, ikm):
-    """RFC5869 HKDF instantiated in the specific form
-    used in Lightning BOLT 8:
-    Extract and expand to 64 bytes using HMAC-SHA256,
-    with info field set to a zero length string as per BOLT8
-    Return as two 32 byte fields.
-    """
-    #Extract
-    prk = hmac.new(salt, msg=ikm, digestmod=hashlib.sha256).digest()
-    assert len(prk) == 32
-    #Expand
-    info = b""
-    T0 = b""
-    T1 = hmac.new(prk, T0 + info + b"\x01", digestmod=hashlib.sha256).digest()
-    T2 = hmac.new(prk, T1 + info + b"\x02", digestmod=hashlib.sha256).digest()
-    assert len(T1 + T2) == 64
-    return T1, T2
-
-def act1_initiator_message(hs, epriv, epub):
-    hs.update(epub)
-    ss = get_ecdh(epriv, hs.responder_pub)
-    ck2, temp_k1 = get_bolt8_hkdf(hs.ck, ss)
-    hs.ck = ck2
-    c = aead_encrypt(temp_k1, 0, hs.h, b"")
-    #for next step if we do it
-    hs.update(c)
-    msg = hs.handshake_version + epub + c
-    assert len(msg) == 50
-    return msg
-
-def privkey_to_pubkey(priv: bytes) -> bytes:
-    return ecc.ECPrivkey(priv[:32]).get_public_key_bytes()
-
-def create_ephemeral_key() -> (bytes, bytes):
-    privkey = ecc.ECPrivkey.generate_random_key()
-    return privkey.get_secret_bytes(), privkey.get_public_key_bytes()
-
-class InitiatorSession:
-    def __init__(self, privkey, remote_pubkey, reader, writer):
-        self.privkey = privkey
-        self.remote_pubkey = remote_pubkey
-        self.reader = reader
-        self.writer = writer
-
-    def send_bytes(self, msg):
-        l = len(msg).to_bytes(2, 'big')
-        lc = aead_encrypt(self.sk, self.sn(), b'', l)
-        c = aead_encrypt(self.sk, self.sn(), b'', msg)
-        assert len(lc) == 18
-        assert len(c) == len(msg) + 16
-        self.writer.write(lc+c)
-
-    async def read_messages(self):
-        read_buffer = b''
-        while True:
-            rn_l, rk_l = self.rn()
-            rn_m, rk_m = self.rn()
-            while True:
-                if len(read_buffer) >= 18:
-                    lc = read_buffer[:18]
-                    l = aead_decrypt(rk_l, rn_l, b'', lc)
-                    length = int.from_bytes(l, 'big')
-                    offset = 18 + length + 16
-                    if len(read_buffer) >= offset:
-                        c = read_buffer[18:offset]
-                        read_buffer = read_buffer[offset:]
-                        msg = aead_decrypt(rk_m, rn_m, b'', c)
-                        yield msg
-                        break
-                try:
-                    s = await self.reader.read(2**10)
-                except:
-                    s = None
-                if not s:
-                    raise LightningPeerConnectionClosed()
-                read_buffer += s
-
-    async def handshake(self):
-        hs = HandshakeState(self.remote_pubkey)
-        # Get a new ephemeral key
-        epriv, epub = create_ephemeral_key()
-
-        msg = act1_initiator_message(hs, epriv, epub)
-        # act 1
-        self.writer.write(msg)
-        rspns = await self.reader.read(2**10)
-        if len(rspns) != 50:
-            raise HandshakeFailed("Lightning handshake act 1 response has bad length, are you sure this is the right pubkey? " + str(bh2u(self.pubkey)))
-        hver, alice_epub, tag = rspns[0], rspns[1:34], rspns[34:]
-        if bytes([hver]) != hs.handshake_version:
-            raise HandshakeFailed("unexpected handshake version: {}".format(hver))
-        # act 2
-        hs.update(alice_epub)
-        ss = get_ecdh(epriv, alice_epub)
-        ck, temp_k2 = get_bolt8_hkdf(hs.ck, ss)
-        hs.ck = ck
-        p = aead_decrypt(temp_k2, 0, hs.h, tag)
-        hs.update(tag)
-        # act 3
-        my_pubkey = privkey_to_pubkey(self.privkey)
-        c = aead_encrypt(temp_k2, 1, hs.h, my_pubkey)
-        hs.update(c)
-        ss = get_ecdh(self.privkey[:32], alice_epub)
-        ck, temp_k3 = get_bolt8_hkdf(hs.ck, ss)
-        hs.ck = ck
-        t = aead_encrypt(temp_k3, 0, hs.h, b'')
-        self.sk, self.rk = get_bolt8_hkdf(hs.ck, b'')
-        msg = hs.handshake_version + c + t
-        self.writer.write(msg)
-        # init counters
-        self._sn = 0
-        self._rn = 0
-        self.r_ck = ck
-        self.s_ck = ck
-
-    def rn(self):
-        o = self._rn, self.rk
-        self._rn += 1
-        if self._rn == 1000:
-            self.r_ck, self.rk = get_bolt8_hkdf(self.r_ck, self.rk)
-            self._rn = 0
-        return o
-
-    def sn(self):
-        o = self._sn
-        self._sn += 1
-        if self._sn == 1000:
-            self.s_ck, self.sk = get_bolt8_hkdf(self.s_ck, self.sk)
-            self._sn = 0
-        return o
-
 
 class Peer(PrintError):
 

electrum/lntransport.py

@@ -0,0 +1,173 @@
+import hmac
+import hashlib
+import cryptography.hazmat.primitives.ciphers.aead as AEAD
+
+from .crypto import sha256
+from .lnutil import get_ecdh
+from .lnutil import LightningPeerConnectionClosed, HandshakeFailed
+from . import ecc
+
+class HandshakeState(object):
+    prologue = b"lightning"
+    protocol_name = b"Noise_XK_secp256k1_ChaChaPoly_SHA256"
+    handshake_version = b"\x00"
+
+    def __init__(self, responder_pub):
+        self.responder_pub = responder_pub
+        self.h = sha256(self.protocol_name)
+        self.ck = self.h
+        self.update(self.prologue)
+        self.update(self.responder_pub)
+
+    def update(self, data):
+        self.h = sha256(self.h + data)
+        return self.h
+
+
+def get_nonce_bytes(n):
+    """BOLT 8 requires the nonce to be 12 bytes, 4 bytes leading
+    zeroes and 8 bytes little endian encoded 64 bit integer.
+    """
+    return b"\x00"*4 + n.to_bytes(8, 'little')
+
+def aead_encrypt(k, nonce, associated_data, data):
+    nonce_bytes = get_nonce_bytes(nonce)
+    a = AEAD.ChaCha20Poly1305(k)
+    return a.encrypt(nonce_bytes, data, associated_data)
+
+def aead_decrypt(k, nonce, associated_data, data):
+    nonce_bytes = get_nonce_bytes(nonce)
+    a = AEAD.ChaCha20Poly1305(k)
+    #raises InvalidTag exception if it's not valid
+    return a.decrypt(nonce_bytes, data, associated_data)
+
+def get_bolt8_hkdf(salt, ikm):
+    """RFC5869 HKDF instantiated in the specific form
+    used in Lightning BOLT 8:
+    Extract and expand to 64 bytes using HMAC-SHA256,
+    with info field set to a zero length string as per BOLT8
+    Return as two 32 byte fields.
+    """
+    #Extract
+    prk = hmac.new(salt, msg=ikm, digestmod=hashlib.sha256).digest()
+    assert len(prk) == 32
+    #Expand
+    info = b""
+    T0 = b""
+    T1 = hmac.new(prk, T0 + info + b"\x01", digestmod=hashlib.sha256).digest()
+    T2 = hmac.new(prk, T1 + info + b"\x02", digestmod=hashlib.sha256).digest()
+    assert len(T1 + T2) == 64
+    return T1, T2
+
+def act1_initiator_message(hs, epriv, epub):
+    hs.update(epub)
+    ss = get_ecdh(epriv, hs.responder_pub)
+    ck2, temp_k1 = get_bolt8_hkdf(hs.ck, ss)
+    hs.ck = ck2
+    c = aead_encrypt(temp_k1, 0, hs.h, b"")
+    #for next step if we do it
+    hs.update(c)
+    msg = hs.handshake_version + epub + c
+    assert len(msg) == 50
+    return msg
+
+def privkey_to_pubkey(priv: bytes) -> bytes:
+    return ecc.ECPrivkey(priv[:32]).get_public_key_bytes()
+
+def create_ephemeral_key() -> (bytes, bytes):
+    privkey = ecc.ECPrivkey.generate_random_key()
+    return privkey.get_secret_bytes(), privkey.get_public_key_bytes()
+
+class LNTransport:
+    def __init__(self, privkey, remote_pubkey, reader, writer):
+        self.privkey = privkey
+        self.remote_pubkey = remote_pubkey
+        self.reader = reader
+        self.writer = writer
+
+    def send_bytes(self, msg):
+        l = len(msg).to_bytes(2, 'big')
+        lc = aead_encrypt(self.sk, self.sn(), b'', l)
+        c = aead_encrypt(self.sk, self.sn(), b'', msg)
+        assert len(lc) == 18
+        assert len(c) == len(msg) + 16
+        self.writer.write(lc+c)
+
+    async def read_messages(self):
+        read_buffer = b''
+        while True:
+            rn_l, rk_l = self.rn()
+            rn_m, rk_m = self.rn()
+            while True:
+                if len(read_buffer) >= 18:
+                    lc = read_buffer[:18]
+                    l = aead_decrypt(rk_l, rn_l, b'', lc)
+                    length = int.from_bytes(l, 'big')
+                    offset = 18 + length + 16
+                    if len(read_buffer) >= offset:
+                        c = read_buffer[18:offset]
+                        read_buffer = read_buffer[offset:]
+                        msg = aead_decrypt(rk_m, rn_m, b'', c)
+                        yield msg
+                        break
+                try:
+                    s = await self.reader.read(2**10)
+                except:
+                    s = None
+                if not s:
+                    raise LightningPeerConnectionClosed()
+                read_buffer += s
+
+    async def handshake(self):
+        hs = HandshakeState(self.remote_pubkey)
+        # Get a new ephemeral key
+        epriv, epub = create_ephemeral_key()
+
+        msg = act1_initiator_message(hs, epriv, epub)
+        # act 1
+        self.writer.write(msg)
+        rspns = await self.reader.read(2**10)
+        if len(rspns) != 50:
+            raise HandshakeFailed("Lightning handshake act 1 response has bad length, are you sure this is the right pubkey? " + str(bh2u(self.pubkey)))
+        hver, alice_epub, tag = rspns[0], rspns[1:34], rspns[34:]
+        if bytes([hver]) != hs.handshake_version:
+            raise HandshakeFailed("unexpected handshake version: {}".format(hver))
+        # act 2
+        hs.update(alice_epub)
+        ss = get_ecdh(epriv, alice_epub)
+        ck, temp_k2 = get_bolt8_hkdf(hs.ck, ss)
+        hs.ck = ck
+        p = aead_decrypt(temp_k2, 0, hs.h, tag)
+        hs.update(tag)
+        # act 3
+        my_pubkey = privkey_to_pubkey(self.privkey)
+        c = aead_encrypt(temp_k2, 1, hs.h, my_pubkey)
+        hs.update(c)
+        ss = get_ecdh(self.privkey[:32], alice_epub)
+        ck, temp_k3 = get_bolt8_hkdf(hs.ck, ss)
+        hs.ck = ck
+        t = aead_encrypt(temp_k3, 0, hs.h, b'')
+        self.sk, self.rk = get_bolt8_hkdf(hs.ck, b'')
+        msg = hs.handshake_version + c + t
+        self.writer.write(msg)
+        # init counters
+        self._sn = 0
+        self._rn = 0
+        self.r_ck = ck
+        self.s_ck = ck
+
+    def rn(self):
+        o = self._rn, self.rk
+        self._rn += 1
+        if self._rn == 1000:
+            self.r_ck, self.rk = get_bolt8_hkdf(self.r_ck, self.rk)
+            self._rn = 0
+        return o
+
+    def sn(self):
+        o = self._sn
+        self._sn += 1
+        if self._sn == 1000:
+            self.s_ck, self.sk = get_bolt8_hkdf(self.s_ck, self.sk)
+            self._sn = 0
+        return o

electrum/lnutil.py

@@ -57,7 +57,17 @@ class Outpoint(NamedTuple("Outpoint", [('txid', str), ('output_index', int)])):
         return "{}:{}".format(self.txid, self.output_index)
 
 
-class UnableToDeriveSecret(Exception): pass
+class LightningError(Exception):
+    pass
+
+class LightningPeerConnectionClosed(LightningError):
+    pass
+
+class UnableToDeriveSecret(LightningError):
+    pass
+
+class HandshakeFailed(LightningError):
+    pass
 
 
 class RevocationStore:

electrum/lnworker.py

@@ -16,7 +16,8 @@ from . import bitcoin
 from .keystore import BIP32_KeyStore
 from .bitcoin import sha256, COIN
 from .util import bh2u, bfh, PrintError, InvoiceError, resolve_dns_srv, is_ip_address, log_exceptions
-from .lnbase import Peer, InitiatorSession
+from .lntransport import LNTransport
+from .lnbase import Peer
 from .lnaddr import lnencode, LnAddr, lndecode
 from .ecc import der_sig_from_sig_string
 from .lnchan import Channel
@@ -29,7 +30,6 @@ from .lnaddr import lndecode
 from .i18n import _
 from .lnrouter import RouteEdge
 
-
 NUM_PEERS_TARGET = 4
 PEER_RETRY_INTERVAL = 600  # seconds
 PEER_RETRY_INTERVAL_FOR_CHANNELS = 30  # seconds
@@ -114,7 +114,7 @@ class LNWorker(PrintError):
         self.print_error("adding peer", peer_addr)
         async def _init_peer():
             reader, writer = await asyncio.open_connection(peer_addr.host, peer_addr.port)
-            transport = InitiatorSession(self.node_keypair.privkey, node_id, reader, writer)
+            transport = LNTransport(self.node_keypair.privkey, node_id, reader, writer)
             peer.transport = transport
             await self.network.main_taskgroup.spawn(peer.main_loop())
         asyncio.ensure_future(_init_peer())