diff --git a/src/electrum_aionostr/key.py b/src/electrum_aionostr/key.py index 4d4e2dd..aa6f40f 100644 --- a/src/electrum_aionostr/key.py +++ b/src/electrum_aionostr/key.py @@ -7,11 +7,16 @@ import electrum_ecc as ecc -from .crypto_aes import aes_encrypt_with_iv, aes_decrypt_with_iv from .delegation import Delegation -from .event import Event from . import bech32 +AES_AVAILABLE = False +try: + from .crypto_aes import aes_encrypt_with_iv, aes_decrypt_with_iv + AES_AVAILABLE = True +except ImportError: + pass + class PublicKey: def __init__(self, raw_bytes: bytes) -> None: @@ -70,6 +75,8 @@ def compute_shared_secret(self, public_key_hex: str) -> bytes: return int.to_bytes(pt.x(), length=32, byteorder='big', signed=False) def encrypt_message(self, message: str, public_key_hex: str) -> str: + if not AES_AVAILABLE: + raise CryptoBackendUnavailableError iv = secrets.token_bytes(16) encrypted_message = aes_encrypt_with_iv( key=self.compute_shared_secret(public_key_hex), @@ -79,6 +86,8 @@ def encrypt_message(self, message: str, public_key_hex: str) -> str: return f"{base64.b64encode(encrypted_message).decode()}?iv={base64.b64encode(iv).decode()}" def decrypt_message(self, encoded_message: str, public_key_hex: str) -> str: + if not AES_AVAILABLE: + raise CryptoBackendUnavailableError encoded_data = encoded_message.split("?iv=") encoded_content, encoded_iv = encoded_data[0], encoded_data[1] @@ -121,3 +130,11 @@ def mine_vanity_key(prefix: str = None, suffix: str = None) -> PrivateKey: break return sk + + +class CryptoBackendUnavailableError(ImportError): + def __init__(self): + super().__init__( + "Error: at least one of ('pycryptodomex', 'cryptography') needs to be installed for NIP-04 functionality.\n" + "Install electrum-aionostr with with [crypto] feature: `pip install electrum-aionostr[crypto]`." + ) diff --git a/tests/test_key.py b/tests/test_key.py index d3062b7..78e7173 100644 --- a/tests/test_key.py +++ b/tests/test_key.py @@ -1,7 +1,8 @@ from hashlib import sha256 import unittest +from unittest.mock import patch -from electrum_aionostr.key import PrivateKey, PublicKey +from electrum_aionostr.key import PrivateKey, PublicKey, CryptoBackendUnavailableError bfh = bytes.fromhex @@ -50,3 +51,13 @@ def test_encrypt_message(self): ciphertext = privkey1.encrypt_message(msg1, privkey2.public_key.hex()) self.assertEqual(msg1, privkey2.decrypt_message(ciphertext, privkey1.public_key.hex())) self.assertEqual(msg1, privkey1.decrypt_message(ciphertext, privkey2.public_key.hex())) + + def test_crypto_backend_unavailable(self): + with patch('electrum_aionostr.key.AES_AVAILABLE', False): + key = PrivateKey() + random_pubkey = PrivateKey().public_key + key.compute_shared_secret(random_pubkey.hex()) # test some other functionality, shouldn't raise + with self.assertRaises(CryptoBackendUnavailableError): + key.encrypt_message(message="test", public_key_hex=random_pubkey.hex()) + with self.assertRaises(CryptoBackendUnavailableError): + key.decrypt_message(encoded_message="test", public_key_hex=random_pubkey.hex())