From dd209c71603271bebd264b67d4552673683eac43 Mon Sep 17 00:00:00 2001 From: Alexander Sergeev <22302418+pseusys@users.noreply.github.com> Date: Mon, 27 Jan 2025 07:47:21 +0100 Subject: [PATCH] custom config added as an optional argument (#385) --- ecies/__init__.py | 18 +++++++++--------- ecies/utils/elliptic.py | 10 +++++----- ecies/utils/symmetric.py | 14 +++++++------- 3 files changed, 21 insertions(+), 21 deletions(-) diff --git a/ecies/__init__.py b/ecies/__init__.py index 2a41ebf..46afaf2 100644 --- a/ecies/__init__.py +++ b/ecies/__init__.py @@ -2,7 +2,7 @@ from coincurve import PrivateKey, PublicKey -from .config import ECIES_CONFIG +from .config import ECIES_CONFIG, Config from .utils import ( decapsulate, encapsulate, @@ -16,7 +16,7 @@ __all__ = ["encrypt", "decrypt", "ECIES_CONFIG"] -def encrypt(receiver_pk: Union[str, bytes], msg: bytes) -> bytes: +def encrypt(receiver_pk: Union[str, bytes], msg: bytes, config: Config = ECIES_CONFIG) -> bytes: """ Encrypt with receiver's secp256k1 public key @@ -41,15 +41,15 @@ def encrypt(receiver_pk: Union[str, bytes], msg: bytes) -> bytes: ephemeral_sk = generate_key() ephemeral_pk = ephemeral_sk.public_key.format( - ECIES_CONFIG.is_ephemeral_key_compressed + config.is_ephemeral_key_compressed ) - sym_key = encapsulate(ephemeral_sk, pk) - encrypted = sym_encrypt(sym_key, msg) + sym_key = encapsulate(ephemeral_sk, pk, config) + encrypted = sym_encrypt(sym_key, msg, config) return ephemeral_pk + encrypted -def decrypt(receiver_sk: Union[str, bytes], msg: bytes) -> bytes: +def decrypt(receiver_sk: Union[str, bytes], msg: bytes, config: Config = ECIES_CONFIG) -> bytes: """ Decrypt with receiver's secp256k1 private key @@ -72,8 +72,8 @@ def decrypt(receiver_sk: Union[str, bytes], msg: bytes) -> bytes: else: raise TypeError("Invalid secret key type") - key_size = ECIES_CONFIG.ephemeral_key_size + key_size = config.ephemeral_key_size ephemeral_pk, encrypted = PublicKey(msg[0:key_size]), msg[key_size:] - sym_key = decapsulate(ephemeral_pk, sk) - return sym_decrypt(sym_key, encrypted) + sym_key = decapsulate(ephemeral_pk, sk, config) + return sym_decrypt(sym_key, encrypted, config) diff --git a/ecies/utils/elliptic.py b/ecies/utils/elliptic.py index d7dd403..eadb897 100644 --- a/ecies/utils/elliptic.py +++ b/ecies/utils/elliptic.py @@ -2,7 +2,7 @@ from coincurve.utils import get_valid_secret from eth_keys import keys -from ..config import ECIES_CONFIG +from ..config import ECIES_CONFIG, Config from .hex import decode_hex from .symmetric import derive_key @@ -95,8 +95,8 @@ def hex2sk(sk_hex: str) -> PrivateKey: # private below -def encapsulate(private_key: PrivateKey, peer_public_key: PublicKey) -> bytes: - is_compressed = ECIES_CONFIG.is_hkdf_key_compressed +def encapsulate(private_key: PrivateKey, peer_public_key: PublicKey, config: Config = ECIES_CONFIG) -> bytes: + is_compressed = config.is_hkdf_key_compressed shared_point = peer_public_key.multiply(private_key.secret) master = private_key.public_key.format(is_compressed) + shared_point.format( is_compressed @@ -104,8 +104,8 @@ def encapsulate(private_key: PrivateKey, peer_public_key: PublicKey) -> bytes: return derive_key(master) -def decapsulate(public_key: PublicKey, peer_private_key: PrivateKey) -> bytes: - is_compressed = ECIES_CONFIG.is_hkdf_key_compressed +def decapsulate(public_key: PublicKey, peer_private_key: PrivateKey, config: Config = ECIES_CONFIG) -> bytes: + is_compressed = config.is_hkdf_key_compressed shared_point = public_key.multiply(peer_private_key.secret) master = public_key.format(is_compressed) + shared_point.format(is_compressed) return derive_key(master) diff --git a/ecies/utils/symmetric.py b/ecies/utils/symmetric.py index 8fe3a23..04837c0 100644 --- a/ecies/utils/symmetric.py +++ b/ecies/utils/symmetric.py @@ -4,14 +4,14 @@ from Crypto.Hash import SHA256 from Crypto.Protocol.KDF import HKDF -from ..config import ECIES_CONFIG +from ..config import ECIES_CONFIG, Config AES_CIPHER_MODE = AES.MODE_GCM AEAD_TAG_LENGTH = 16 XCHACHA20_NONCE_LENGTH = 24 -def sym_encrypt(key: bytes, plain_text: bytes) -> bytes: +def sym_encrypt(key: bytes, plain_text: bytes, config: Config = ECIES_CONFIG) -> bytes: """ Symmetric encryption. AES-256-GCM or XChaCha20-Poly1305. @@ -29,9 +29,9 @@ def sym_encrypt(key: bytes, plain_text: bytes) -> bytes: bytes nonce + tag(16 bytes) + encrypted data """ - algorithm = ECIES_CONFIG.symmetric_algorithm + algorithm = config.symmetric_algorithm if algorithm == "aes-256-gcm": - nonce_length = ECIES_CONFIG.symmetric_nonce_length + nonce_length = config.symmetric_nonce_length nonce = os.urandom(nonce_length) cipher = AES.new(key, AES_CIPHER_MODE, nonce) elif algorithm == "xchacha20": @@ -48,7 +48,7 @@ def sym_encrypt(key: bytes, plain_text: bytes) -> bytes: return bytes(cipher_text) -def sym_decrypt(key: bytes, cipher_text: bytes) -> bytes: +def sym_decrypt(key: bytes, cipher_text: bytes, config: Config = ECIES_CONFIG) -> bytes: """ AES-GCM decryption. AES-256-GCM or XChaCha20-Poly1305. @@ -84,9 +84,9 @@ def sym_decrypt(key: bytes, cipher_text: bytes) -> bytes: # If it's 12 bytes, the nonce can be incremented by 1 for each encryption # If it's 16 bytes, the nonce will be used to hash, so it's meaningless to increment - algorithm = ECIES_CONFIG.symmetric_algorithm + algorithm = config.symmetric_algorithm if algorithm == "aes-256-gcm": - nonce_length = ECIES_CONFIG.symmetric_nonce_length + nonce_length = config.symmetric_nonce_length nonce_tag_length = nonce_length + AEAD_TAG_LENGTH nonce = cipher_text[:nonce_length] tag = cipher_text[nonce_length:nonce_tag_length]