custom config added as an optional argument (#385)

Author: pseusys

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)

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,17 +95,17 @@ 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
     )
     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)

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]