From 031ce14857ad994bc50ae261abc1d178ed8d2b68 Mon Sep 17 00:00:00 2001 From: Vlad Stan Date: Mon, 15 Jan 2024 11:51:15 +0200 Subject: [PATCH] refactor: extract `AESCipher` to `crypto.py` (#2202) --- lnbits/core/views/auth_api.py | 4 +- lnbits/utils/crypto.py | 75 +++++++++++++++++++++++++++++ lnbits/wallets/lndgrpc.py | 3 +- lnbits/wallets/lndrest.py | 3 +- lnbits/wallets/macaroon/__init__.py | 2 +- lnbits/wallets/macaroon/macaroon.py | 73 +--------------------------- 6 files changed, 82 insertions(+), 78 deletions(-) create mode 100644 lnbits/utils/crypto.py diff --git a/lnbits/core/views/auth_api.py b/lnbits/core/views/auth_api.py index 48e602b4..7e322076 100644 --- a/lnbits/core/views/auth_api.py +++ b/lnbits/core/views/auth_api.py @@ -20,9 +20,7 @@ from lnbits.helpers import ( is_valid_username, ) from lnbits.settings import AuthMethods, settings - -# todo: move this class to a `crypto.py` file -from lnbits.wallets.macaroon.macaroon import AESCipher +from lnbits.utils.crypto import AESCipher from ..crud import ( create_account, diff --git a/lnbits/utils/crypto.py b/lnbits/utils/crypto.py new file mode 100644 index 00000000..f60c4783 --- /dev/null +++ b/lnbits/utils/crypto.py @@ -0,0 +1,75 @@ +import base64 +import getpass +from hashlib import md5 + +from Cryptodome import Random +from Cryptodome.Cipher import AES + +BLOCK_SIZE = 16 + + +class AESCipher: + """This class is compatible with crypto-js/aes.js + + Encrypt and decrypt in Javascript using: + import AES from "crypto-js/aes.js"; + import Utf8 from "crypto-js/enc-utf8.js"; + AES.encrypt(decrypted, password).toString() + AES.decrypt(encrypted, password).toString(Utf8); + + """ + + def __init__(self, key=None, description=""): + self.key = key + self.description = description + " " + + def pad(self, data): + length = BLOCK_SIZE - (len(data) % BLOCK_SIZE) + return data + (chr(length) * length).encode() + + def unpad(self, data): + return data[: -(data[-1] if isinstance(data[-1], int) else ord(data[-1]))] + + @property + def passphrase(self): + passphrase = self.key if self.key is not None else None + if passphrase is None: + passphrase = getpass.getpass(f"Enter {self.description}password:") + return passphrase + + def bytes_to_key(self, data, salt, output=48): + # extended from https://gist.github.com/gsakkis/4546068 + assert len(salt) == 8, len(salt) + data += salt + key = md5(data).digest() + final_key = key + while len(final_key) < output: + key = md5(key + data).digest() + final_key += key + return final_key[:output] + + def decrypt(self, encrypted: str) -> str: # type: ignore + """Decrypts a string using AES-256-CBC.""" + passphrase = self.passphrase + encrypted = base64.b64decode(encrypted) # type: ignore + assert encrypted[0:8] == b"Salted__" + salt = encrypted[8:16] + key_iv = self.bytes_to_key(passphrase.encode(), salt, 32 + 16) + key = key_iv[:32] + iv = key_iv[32:] + aes = AES.new(key, AES.MODE_CBC, iv) + try: + return self.unpad(aes.decrypt(encrypted[16:])).decode() # type: ignore + except UnicodeDecodeError: + raise ValueError("Wrong passphrase") + + def encrypt(self, message: bytes) -> str: + passphrase = self.passphrase + salt = Random.new().read(8) + key_iv = self.bytes_to_key(passphrase.encode(), salt, 32 + 16) + key = key_iv[:32] + iv = key_iv[32:] + aes = AES.new(key, AES.MODE_CBC, iv) + return base64.b64encode( + b"Salted__" + salt + aes.encrypt(self.pad(message)) + ).decode() diff --git a/lnbits/wallets/lndgrpc.py b/lnbits/wallets/lndgrpc.py index abb554ad..df385117 100644 --- a/lnbits/wallets/lndgrpc.py +++ b/lnbits/wallets/lndgrpc.py @@ -12,6 +12,7 @@ import lnbits.wallets.lnd_grpc_files.lightning_pb2_grpc as lnrpc import lnbits.wallets.lnd_grpc_files.router_pb2 as router import lnbits.wallets.lnd_grpc_files.router_pb2_grpc as routerrpc from lnbits.settings import settings +from lnbits.utils.crypto import AESCipher from .base import ( InvoiceResponse, @@ -20,7 +21,7 @@ from .base import ( StatusResponse, Wallet, ) -from .macaroon import AESCipher, load_macaroon +from .macaroon import load_macaroon def b64_to_bytes(checking_id: str) -> bytes: diff --git a/lnbits/wallets/lndrest.py b/lnbits/wallets/lndrest.py index 7ae68be8..d50b548e 100644 --- a/lnbits/wallets/lndrest.py +++ b/lnbits/wallets/lndrest.py @@ -9,6 +9,7 @@ from loguru import logger from lnbits.nodes.lndrest import LndRestNode from lnbits.settings import settings +from lnbits.utils.crypto import AESCipher from .base import ( InvoiceResponse, @@ -17,7 +18,7 @@ from .base import ( StatusResponse, Wallet, ) -from .macaroon import AESCipher, load_macaroon +from .macaroon import load_macaroon class LndRestWallet(Wallet): diff --git a/lnbits/wallets/macaroon/__init__.py b/lnbits/wallets/macaroon/__init__.py index 16617aa6..44efcf4f 100644 --- a/lnbits/wallets/macaroon/__init__.py +++ b/lnbits/wallets/macaroon/__init__.py @@ -1 +1 @@ -from .macaroon import AESCipher, load_macaroon +from .macaroon import load_macaroon diff --git a/lnbits/wallets/macaroon/macaroon.py b/lnbits/wallets/macaroon/macaroon.py index c3712c20..8985c647 100644 --- a/lnbits/wallets/macaroon/macaroon.py +++ b/lnbits/wallets/macaroon/macaroon.py @@ -1,12 +1,8 @@ import base64 -import getpass -from hashlib import md5 -from Cryptodome import Random -from Cryptodome.Cipher import AES from loguru import logger -BLOCK_SIZE = 16 +from lnbits.utils.crypto import AESCipher def load_macaroon(macaroon: str) -> str: @@ -40,73 +36,6 @@ def load_macaroon(macaroon: str) -> str: # todo: move to its own (crypto.py) file -class AESCipher: - """This class is compatible with crypto-js/aes.js - - Encrypt and decrypt in Javascript using: - import AES from "crypto-js/aes.js"; - import Utf8 from "crypto-js/enc-utf8.js"; - AES.encrypt(decrypted, password).toString() - AES.decrypt(encrypted, password).toString(Utf8); - - """ - - def __init__(self, key=None, description=""): - self.key = key - self.description = description + " " - - def pad(self, data): - length = BLOCK_SIZE - (len(data) % BLOCK_SIZE) - return data + (chr(length) * length).encode() - - def unpad(self, data): - return data[: -(data[-1] if isinstance(data[-1], int) else ord(data[-1]))] - - @property - def passphrase(self): - passphrase = self.key if self.key is not None else None - if passphrase is None: - passphrase = getpass.getpass(f"Enter {self.description}password:") - return passphrase - - def bytes_to_key(self, data, salt, output=48): - # extended from https://gist.github.com/gsakkis/4546068 - assert len(salt) == 8, len(salt) - data += salt - key = md5(data).digest() - final_key = key - while len(final_key) < output: - key = md5(key + data).digest() - final_key += key - return final_key[:output] - - def decrypt(self, encrypted: str) -> str: # type: ignore - """Decrypts a string using AES-256-CBC.""" - passphrase = self.passphrase - encrypted = base64.b64decode(encrypted) # type: ignore - assert encrypted[0:8] == b"Salted__" - salt = encrypted[8:16] - key_iv = self.bytes_to_key(passphrase.encode(), salt, 32 + 16) - key = key_iv[:32] - iv = key_iv[32:] - aes = AES.new(key, AES.MODE_CBC, iv) - try: - return self.unpad(aes.decrypt(encrypted[16:])).decode() # type: ignore - except UnicodeDecodeError: - raise ValueError("Wrong passphrase") - - def encrypt(self, message: bytes) -> str: - passphrase = self.passphrase - salt = Random.new().read(8) - key_iv = self.bytes_to_key(passphrase.encode(), salt, 32 + 16) - key = key_iv[:32] - iv = key_iv[32:] - aes = AES.new(key, AES.MODE_CBC, iv) - return base64.b64encode( - b"Salted__" + salt + aes.encrypt(self.pad(message)) - ).decode() - - # if this file is executed directly, ask for a macaroon and encrypt it if __name__ == "__main__": macaroon = input("Enter macaroon: ")