From 18f3a37032627bee9247df397204e068ec62c7ea Mon Sep 17 00:00:00 2001
From: SomberNight <somber.night@protonmail.com>
Date: Wed, 4 Mar 2020 16:15:22 +0100
Subject: [PATCH] crypto: move LN-related chacha20/poly1305 code into crypto.py

---
 electrum/crypto.py      | 21 +++++++++++++++++++++
 electrum/lnonion.py     |  9 ++++-----
 electrum/lntransport.py | 20 +++++++++-----------
 3 files changed, 34 insertions(+), 16 deletions(-)

diff --git a/electrum/crypto.py b/electrum/crypto.py
index 1d345e6c1..044d35b94 100644
--- a/electrum/crypto.py
+++ b/electrum/crypto.py
@@ -40,6 +40,8 @@ try:
 except:
     AES = None
 
+from Cryptodome.Cipher import ChaCha20_Poly1305, ChaCha20
+
 
 class InvalidPadding(Exception):
     pass
@@ -216,3 +218,22 @@ def hmac_oneshot(key: bytes, msg: bytes, digest) -> bytes:
         return hmac.digest(key, msg, digest)
     else:
         return hmac.new(key, msg, digest).digest()
+
+
+def chacha20_poly1305_encrypt(*, key: bytes, nonce: bytes, associated_data: bytes, data: bytes) -> bytes:
+    cipher = ChaCha20_Poly1305.new(key=key, nonce=nonce)
+    cipher.update(associated_data)
+    ciphertext, mac = cipher.encrypt_and_digest(plaintext=data)
+    return ciphertext + mac
+
+
+def chacha20_poly1305_decrypt(*, key: bytes, nonce: bytes, associated_data: bytes, data: bytes) -> bytes:
+    cipher = ChaCha20_Poly1305.new(key=key, nonce=nonce)
+    cipher.update(associated_data)
+    # raises ValueError if not valid (e.g. incorrect MAC)
+    return cipher.decrypt_and_verify(ciphertext=data[:-16], received_mac_tag=data[-16:])
+
+
+def chacha20_encrypt(*, key: bytes, nonce: bytes, data: bytes) -> bytes:
+    cipher = ChaCha20.new(key=key, nonce=nonce)
+    return cipher.encrypt(data)
diff --git a/electrum/lnonion.py b/electrum/lnonion.py
index a198fe364..3c1753f55 100644
--- a/electrum/lnonion.py
+++ b/electrum/lnonion.py
@@ -27,10 +27,8 @@ import hashlib
 from typing import Sequence, List, Tuple, NamedTuple, TYPE_CHECKING
 from enum import IntEnum, IntFlag
 
-from Cryptodome.Cipher import ChaCha20
-
 from . import ecc
-from .crypto import sha256, hmac_oneshot
+from .crypto import sha256, hmac_oneshot, chacha20_encrypt
 from .util import bh2u, profiler, xor_bytes, bfh
 from .lnutil import (get_ecdh, PaymentFailure, NUM_MAX_HOPS_IN_PAYMENT_PATH,
                      NUM_MAX_EDGES_IN_PAYMENT_PATH, ShortChannelID)
@@ -227,8 +225,9 @@ def generate_filler(key_type: bytes, num_hops: int, hop_size: int,
 
 
 def generate_cipher_stream(stream_key: bytes, num_bytes: int) -> bytes:
-    cipher = ChaCha20.new(key=stream_key, nonce=bytes(8))
-    return cipher.encrypt(bytes(num_bytes))
+    return chacha20_encrypt(key=stream_key,
+                            nonce=bytes(8),
+                            data=bytes(num_bytes))
 
 
 class ProcessedOnionPacket(NamedTuple):
diff --git a/electrum/lntransport.py b/electrum/lntransport.py
index 8bf830c72..11ba5b155 100644
--- a/electrum/lntransport.py
+++ b/electrum/lntransport.py
@@ -9,9 +9,7 @@ import hashlib
 import asyncio
 from asyncio import StreamReader, StreamWriter
 
-from Cryptodome.Cipher import ChaCha20_Poly1305
-
-from .crypto import sha256, hmac_oneshot
+from .crypto import sha256, hmac_oneshot, chacha20_poly1305_encrypt, chacha20_poly1305_decrypt
 from .lnutil import (get_ecdh, privkey_to_pubkey, LightningPeerConnectionClosed,
                      HandshakeFailed, LNPeerAddr)
 from . import ecc
@@ -41,17 +39,17 @@ def get_nonce_bytes(n):
 
 def aead_encrypt(key: bytes, nonce: int, associated_data: bytes, data: bytes) -> bytes:
     nonce_bytes = get_nonce_bytes(nonce)
-    cipher = ChaCha20_Poly1305.new(key=key, nonce=nonce_bytes)
-    cipher.update(associated_data)
-    ciphertext, mac = cipher.encrypt_and_digest(plaintext=data)
-    return ciphertext + mac
+    return chacha20_poly1305_encrypt(key=key,
+                                     nonce=nonce_bytes,
+                                     associated_data=associated_data,
+                                     data=data)
 
 def aead_decrypt(key: bytes, nonce: int, associated_data: bytes, data: bytes) -> bytes:
     nonce_bytes = get_nonce_bytes(nonce)
-    cipher = ChaCha20_Poly1305.new(key=key, nonce=nonce_bytes)
-    cipher.update(associated_data)
-    # raises ValueError if not valid (e.g. incorrect MAC)
-    return cipher.decrypt_and_verify(ciphertext=data[:-16], received_mac_tag=data[-16:])
+    return chacha20_poly1305_decrypt(key=key,
+                                     nonce=nonce_bytes,
+                                     associated_data=associated_data,
+                                     data=data)
 
 def get_bolt8_hkdf(salt, ikm):
     """RFC5869 HKDF instantiated in the specific form