Skip to content

Commit

Permalink
custom config added as an optional argument (#385)
Browse files Browse the repository at this point in the history
  • Loading branch information
pseusys authored Jan 27, 2025
1 parent 6e8de54 commit dd209c7
Show file tree
Hide file tree
Showing 3 changed files with 21 additions and 21 deletions.
18 changes: 9 additions & 9 deletions ecies/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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)
10 changes: 5 additions & 5 deletions ecies/utils/elliptic.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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)
14 changes: 7 additions & 7 deletions ecies/utils/symmetric.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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":
Expand All @@ -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.
Expand Down Expand Up @@ -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]
Expand Down

0 comments on commit dd209c7

Please sign in to comment.