Add consts (#394)

Author: kigawas

README.md

@@ -48,7 +48,7 @@ Or just use a builtin command `eciespy` in your favorite [command line](#command
 
 Parameters:
 
-- **receiver_pk** - Receiver's public key (hex str or bytes)
+- **receiver_pk** - Receiver's public key (hex `str` or `bytes`)
 - **data** - Data to encrypt
 - **config** - Optional configuration object
 
@@ -58,7 +58,7 @@ Returns: **bytes**
 
 Parameters:
 
-- **receiver_sk** - Receiver's private key (hex str or bytes)
+- **receiver_sk** - Receiver's private key (hex `str` or `bytes`)
 - **data** - Data to decrypt
 - **config** - Optional configuration object
 
@@ -116,12 +116,11 @@ $ rm sk pk data enc_data
 Ephemeral key format in the payload and shared key in the key derivation can be configured as compressed or uncompressed format.
 
 ```py
+from .consts import COMPRESSED_PUBLIC_KEY_SIZE, UNCOMPRESSED_PUBLIC_KEY_SIZE
+
 SymmetricAlgorithm = Literal["aes-256-gcm", "xchacha20"]
 NonceLength = Literal[12, 16]  # only for aes-256-gcm, xchacha20 will always be 24
 
-COMPRESSED_PUBLIC_KEY_SIZE = 33
-UNCOMPRESSED_PUBLIC_KEY_SIZE = 65
-
 
 @dataclass()
 class Config:

ecies/__init__.py

@@ -26,7 +26,7 @@ def encrypt(
     Parameters
     ----------
     receiver_pk: Union[str, bytes]
-        Receiver's public key (hex str or bytes)
+        Receiver's public key (hex `str` or `bytes`)
     data: bytes
         Data to encrypt
     config: Config
@@ -63,7 +63,7 @@ def decrypt(
     Parameters
     ----------
     receiver_sk: Union[str, bytes]
-        Receiver's private key (hex str or bytes)
+        Receiver's private key (hex `str` or `bytes`)
     data: bytes
         Data to decrypt
     config: Config

ecies/config.py

@@ -1,12 +1,11 @@
 from dataclasses import dataclass
 from typing import Literal
 
+from .consts import COMPRESSED_PUBLIC_KEY_SIZE, UNCOMPRESSED_PUBLIC_KEY_SIZE
+
 SymmetricAlgorithm = Literal["aes-256-gcm", "xchacha20"]
 NonceLength = Literal[12, 16]  # only for aes-256-gcm, xchacha20 will always be 24
 
-COMPRESSED_PUBLIC_KEY_SIZE = 33
-UNCOMPRESSED_PUBLIC_KEY_SIZE = 65
-
 
 @dataclass()
 class Config:

ecies/consts.py

@@ -0,0 +1,8 @@
+# elliptic
+COMPRESSED_PUBLIC_KEY_SIZE = 33
+UNCOMPRESSED_PUBLIC_KEY_SIZE = 65
+ETH_PUBLIC_KEY_LENGTH = 64
+
+# symmetric
+XCHACHA20_NONCE_LENGTH = 24
+AEAD_TAG_LENGTH = 16

ecies/utils/eth.py

@@ -1,9 +1,8 @@
 from coincurve import PublicKey
 from coincurve.utils import get_valid_secret
 
-from .hash import keccak_256
-
-ETH_PUBLIC_KEY_LENGTH = 64
+from ..consts import ETH_PUBLIC_KEY_LENGTH
+from .hash import keccak256
 
 
 def generate_eth_key():
@@ -23,13 +22,14 @@ def generate_eth_key():
     return keys.PrivateKey(get_valid_secret())
 
 
+# for cli only
 def to_eth_public_key(pk: PublicKey) -> bytes:
     return pk.format(False)[1:]
 
 
 def to_eth_address(pk: PublicKey) -> str:
     pk_bytes = to_eth_public_key(pk)
-    return encode_checksum(keccak_256(pk_bytes)[-20:].hex())
+    return encode_checksum(keccak256(pk_bytes)[-20:].hex())
 
 
 # private below
@@ -42,7 +42,7 @@ def convert_eth_public_key(data: bytes):
 def encode_checksum(raw_address: str) -> str:
     # https://github.com/ethereum/ercs/blob/master/ERCS/erc-55.md
     address = raw_address.lower().replace("0x", "")
-    address_hash = keccak_256(address.encode()).hex()
+    address_hash = keccak256(address.encode()).hex()
 
     res = []
     for a, h in zip(address, address_hash):

ecies/utils/hash.py

@@ -25,11 +25,12 @@ def sha256(data: bytes) -> bytes:
 
 
 def derive_key(master: bytes, salt: bytes = b"") -> bytes:
+    # for aes256 and xchacha20
     derived = HKDF(master, 32, salt, SHA256, num_keys=1)
     return derived  # type: ignore
 
 
 # private below
-def keccak_256(b: bytes) -> bytes:
+def keccak256(b: bytes) -> bytes:
     h = keccak.new(data=b, digest_bits=256)
     return h.digest()

ecies/utils/symmetric.py

@@ -3,17 +3,14 @@
 from Crypto.Cipher import AES, ChaCha20_Poly1305
 
 from ..config import NonceLength, SymmetricAlgorithm
-
-AES_CIPHER_MODE = AES.MODE_GCM
-AEAD_TAG_LENGTH = 16
-XCHACHA20_NONCE_LENGTH = 24
+from ..consts import AEAD_TAG_LENGTH, XCHACHA20_NONCE_LENGTH
 
 
 def sym_encrypt(
     key: bytes,
     plain_text: bytes,
     algorithm: SymmetricAlgorithm = "aes-256-gcm",
-    aes_nonce_length: NonceLength = 16,
+    nonce_length: NonceLength = 16,
 ) -> bytes:
     """
     Symmetric encryption. AES-256-GCM or XChaCha20-Poly1305.
@@ -33,8 +30,8 @@ def sym_encrypt(
         nonce + tag(16 bytes) + encrypted data
     """
     if algorithm == "aes-256-gcm":
-        nonce = os.urandom(aes_nonce_length)
-        aes_cipher = AES.new(key, AES_CIPHER_MODE, nonce)
+        nonce = os.urandom(nonce_length)
+        aes_cipher = AES.new(key, AES.MODE_GCM, nonce)
         encrypted, tag = aes_cipher.encrypt_and_digest(plain_text)
     elif algorithm == "xchacha20":
         nonce = os.urandom(XCHACHA20_NONCE_LENGTH)
@@ -92,11 +89,11 @@ def sym_decrypt(
     # If it's 16 bytes, the nonce will be used to hash, so it's meaningless to increment
 
     if algorithm == "aes-256-gcm":
-        nonce, tag, ciphered_data = _split_cipher_text(cipher_text, nonce_length)
-        aes_cipher = AES.new(key, AES_CIPHER_MODE, nonce)
+        nonce, tag, ciphered_data = __split_cipher_text(cipher_text, nonce_length)
+        aes_cipher = AES.new(key, AES.MODE_GCM, nonce)
         return aes_cipher.decrypt_and_verify(ciphered_data, tag)
     elif algorithm == "xchacha20":
-        nonce, tag, ciphered_data = _split_cipher_text(
+        nonce, tag, ciphered_data = __split_cipher_text(
             cipher_text, XCHACHA20_NONCE_LENGTH
         )
         xchacha_cipher = ChaCha20_Poly1305.new(key=key, nonce=nonce)
@@ -105,7 +102,7 @@ def sym_decrypt(
         raise NotImplementedError
 
 
-def _split_cipher_text(cipher_text: bytes, nonce_length: int):
+def __split_cipher_text(cipher_text: bytes, nonce_length: int):
     nonce_tag_length = nonce_length + AEAD_TAG_LENGTH
     nonce = cipher_text[:nonce_length]
     tag = cipher_text[nonce_length:nonce_tag_length]

tests/utils/test_symmetric.py

@@ -23,19 +23,30 @@ def test_symmetric_big(algorithm, big_data):
     __check_symmetric_random(big_data, algorithm)
 
 
-def test_aes_known():
-    key = decode_hex("0000000000000000000000000000000000000000000000000000000000000000")
-    nonce = decode_hex("0xf3e1ba810d2c8900b11312b7c725565f")
-    tag = decode_hex("0Xec3b71e17c11dbe31484da9450edcf6c")
-    encrypted = decode_hex("02d2ffed93b856f148b9")
-    data = b"".join([nonce, tag, encrypted])
-    assert b"helloworld" == sym_decrypt(key, data)
-
+@pytest.mark.parametrize(
+    "key,nonce,tag,encrypted,algorithm",
+    [
+        (
+            "0000000000000000000000000000000000000000000000000000000000000000",
+            "0xf3e1ba810d2c8900b11312b7c725565f",
+            "0Xec3b71e17c11dbe31484da9450edcf6c",
+            "02d2ffed93b856f148b9",
+            "aes-256-gcm",
+        ),
+        (
+            "27bd6ec46292a3b421cdaf8a3f0ca759cbc67bcbe7c5855aa0d1e0700fd0e828",
+            "0xfbd5dd10431af533c403d6f4fa629931e5f31872d2f7e7b6",
+            "0X5b5ccc27324af03b7ca92dd067ad6eb5",
+            "aa0664f3c00a09d098bf",
+            "xchacha20",
+        ),
+    ],
+)
+def test_known(key, nonce, tag, encrypted, algorithm):
+    key = decode_hex(key)
+    nonce = decode_hex(nonce)
+    tag = decode_hex(tag)
+    encrypted = decode_hex(encrypted)
 
-def test_xchacha20_known():
-    key = decode_hex("27bd6ec46292a3b421cdaf8a3f0ca759cbc67bcbe7c5855aa0d1e0700fd0e828")
-    nonce = decode_hex("0xfbd5dd10431af533c403d6f4fa629931e5f31872d2f7e7b6")
-    tag = decode_hex("0X5b5ccc27324af03b7ca92dd067ad6eb5")
-    encrypted = decode_hex("aa0664f3c00a09d098bf")
     data = b"".join([nonce, tag, encrypted])
-    assert b"helloworld" == sym_decrypt(key, data, "xchacha20")
+    assert b"helloworld" == sym_decrypt(key, data, algorithm)