ThomasV
7 years ago
4 changed files with 189 additions and 182 deletions
@ -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 |
||||
Loading…
Reference in new issue