Add homemade key pair

CHANGELOG.md

@@ -1,5 +1,10 @@
 # Release Notes
 
+## 0.4.5
+
+- Add homemade key pair: `ecies.keys.PrivateKey` and `ecies.keys.PublicKey`
+- Bump dependencies
+
 ## 0.4.4
 
 - Make `eth-keys` optional

DETAILS.md

@@ -136,6 +136,6 @@ Now we have the shared key, and we can use the `nonce` and `tag` to decrypt. Thi
 b'helloworld'
 ```
 
-> Strictly speaking, `nonce` != `iv`, but this is a little bit off topic, if you are curious, you can check [the comment in `utils/symmetric.py`](./ecies/utils/symmetric.py#L86).
+> Strictly speaking, `nonce` != `iv`, but this is a little bit off topic, if you are curious, you can check [the comment in `utils/symmetric.py`](./ecies/utils/symmetric.py#L83).
 >
 > Warning: it's dangerous to reuse nonce, if you don't know what you are doing, just follow the default setting.

README.md

@@ -30,14 +30,18 @@ Or `pip install 'eciespy[eth]'` to install `eth-keys` as well.
 ## Quick Start
 
 ```python
->>> from ecies.utils import generate_key
+>>> from ecies.keys import PrivateKey
 >>> from ecies import encrypt, decrypt
 >>> data = 'hello world🌍'.encode()
->>> sk = generate_key()
+>>> sk = PrivateKey('secp256k1')
 >>> sk_bytes = sk.secret  # bytes
->>> pk_bytes = sk.public_key.format(True)  # bytes
+>>> pk_bytes = sk.public_key.to_bytes(True)  # bytes
 >>> decrypt(sk_bytes, encrypt(pk_bytes, data)).decode()
 'hello world🌍'
+>>> sk_hex = sk.to_hex() # hex str
+>>> pk_hex = sk.public_key.to_hex() # hex str
+>>> decrypt(sk_hex, encrypt(pk_hex, data)).decode()
+'hello world🌍'
 ```
 
 Or just use a builtin command `eciespy` in your favorite [command line](#command-line-interface).

ecies/__init__.py

@@ -1,18 +1,8 @@
 from typing import Union
 
-from coincurve import PrivateKey, PublicKey
-
 from .config import ECIES_CONFIG, Config
-from .utils import (
-    bytes2pk,
-    decapsulate,
-    encapsulate,
-    generate_key,
-    hex2pk,
-    hex2sk,
-    sym_decrypt,
-    sym_encrypt,
-)
+from .keys import PrivateKey, PublicKey
+from .utils import sym_decrypt, sym_encrypt
 
 __all__ = ["encrypt", "decrypt", "ECIES_CONFIG"]
 
@@ -37,17 +27,18 @@ def encrypt(
     bytes
         Encrypted data
     """
+    curve = config.elliptic_curve
     if isinstance(receiver_pk, str):
-        pk = hex2pk(receiver_pk)
+        _receiver_pk = PublicKey.from_hex(curve, receiver_pk)
     elif isinstance(receiver_pk, bytes):
-        pk = bytes2pk(receiver_pk)
+        _receiver_pk = PublicKey(curve, receiver_pk)
     else:
         raise TypeError("Invalid public key type")
 
-    ephemeral_sk = generate_key()
-    ephemeral_pk = ephemeral_sk.public_key.format(config.is_ephemeral_key_compressed)
+    ephemeral_sk = PrivateKey(curve)
+    ephemeral_pk = ephemeral_sk.public_key.to_bytes(config.is_ephemeral_key_compressed)
 
-    sym_key = encapsulate(ephemeral_sk, pk, config.is_hkdf_key_compressed)
+    sym_key = ephemeral_sk.encapsulate(_receiver_pk, config.is_hkdf_key_compressed)
     encrypted = sym_encrypt(
         sym_key, data, config.symmetric_algorithm, config.symmetric_nonce_length
     )
@@ -74,17 +65,18 @@ def decrypt(
     bytes
         Plain text
     """
+    curve = config.elliptic_curve
     if isinstance(receiver_sk, str):
-        sk = hex2sk(receiver_sk)
+        _receiver_sk = PrivateKey.from_hex(curve, receiver_sk)
     elif isinstance(receiver_sk, bytes):
-        sk = PrivateKey(receiver_sk)
+        _receiver_sk = PrivateKey(curve, receiver_sk)
     else:
         raise TypeError("Invalid secret key type")
 
     key_size = config.ephemeral_key_size
-    ephemeral_pk, encrypted = PublicKey(data[0:key_size]), data[key_size:]
+    ephemeral_pk, encrypted = PublicKey(curve, data[0:key_size]), data[key_size:]
 
-    sym_key = decapsulate(ephemeral_pk, sk, config.is_hkdf_key_compressed)
+    sym_key = ephemeral_pk.decapsulate(_receiver_sk, config.is_hkdf_key_compressed)
     return sym_decrypt(
         sym_key, encrypted, config.symmetric_algorithm, config.symmetric_nonce_length
     )

ecies/__main__.py

@@ -14,7 +14,8 @@
 import sys
 
 from ecies import decrypt, encrypt
-from ecies.utils import generate_key, to_eth_address, to_eth_public_key
+from ecies.keys import PrivateKey
+from ecies.utils import to_eth_address
 
 __description__ = "Elliptic Curve Integrated Encryption Scheme for secp256k1 in Python"
 
@@ -68,11 +69,12 @@ def main():
 
     args = parser.parse_args()
     if args.generate:
-        k = generate_key()
+        k = PrivateKey("secp256k1")
+        eth_pk_bytes = k.public_key.to_bytes()[1:]
         sk, pk, addr = (
             k.to_hex(),
-            f"0x{to_eth_public_key(k.public_key).hex()}",
-            to_eth_address(k.public_key),
+            f"0x{eth_pk_bytes.hex()}",
+            to_eth_address(eth_pk_bytes),
         )
         print("Private: {}\nPublic: {}\nAddress: {}".format(sk, pk, addr))
         return

ecies/config.py

@@ -3,12 +3,14 @@
 
 from .consts import COMPRESSED_PUBLIC_KEY_SIZE, UNCOMPRESSED_PUBLIC_KEY_SIZE
 
+EllipticCurve = Literal["secp256k1"]
 SymmetricAlgorithm = Literal["aes-256-gcm", "xchacha20"]
 NonceLength = Literal[12, 16]  # only for aes-256-gcm, xchacha20 will always be 24
 
 
 @dataclass()
 class Config:
+    elliptic_curve: EllipticCurve = "secp256k1"
     is_ephemeral_key_compressed: bool = False
     is_hkdf_key_compressed: bool = False
     symmetric_algorithm: SymmetricAlgorithm = "aes-256-gcm"

ecies/keys/__init__.py

@@ -0,0 +1,4 @@
+from .private import PrivateKey
+from .public import PublicKey
+
+__all__ = ["PrivateKey", "PublicKey"]

ecies/keys/helper.py

@@ -0,0 +1,55 @@
+from coincurve import PublicKey
+from coincurve.utils import GROUP_ORDER_INT
+from Crypto.Random import get_random_bytes
+
+from ..config import EllipticCurve
+from ..consts import ETH_PUBLIC_KEY_LENGTH
+
+
+def is_valid_secret(curve: EllipticCurve, secret: bytes) -> bool:
+    if curve == "secp256k1":
+        return 0 < bytes_to_int(secret) < GROUP_ORDER_INT
+    raise NotImplementedError
+
+
+def get_valid_secret(curve: EllipticCurve) -> bytes:
+    while True:
+        key = get_random_bytes(32)
+        if is_valid_secret(curve, key):
+            return key
+
+
+def get_public_key(
+    curve: EllipticCurve, secret: bytes, compressed: bool = False
+) -> bytes:
+    if curve == "secp256k1":
+        return PublicKey.from_secret(secret).format(compressed)
+    raise NotImplementedError
+
+
+def get_shared_point(
+    curve: EllipticCurve, sk: bytes, pk: bytes, compressed: bool = False
+) -> bytes:
+    if curve == "secp256k1":
+        return PublicKey(pk).multiply(sk).format(compressed)
+    raise NotImplementedError
+
+
+def convert_public_key(
+    curve: EllipticCurve, data: bytes, compressed: bool = False
+) -> bytes:
+    if curve == "secp256k1":
+        # handle 33/65/64 bytes
+        return PublicKey(pad_eth_public_key(data)).format(compressed)
+    raise NotImplementedError
+
+
+# private below
+def bytes_to_int(data: bytes) -> int:
+    return int.from_bytes(data, "big")
+
+
+def pad_eth_public_key(data: bytes):
+    if len(data) == ETH_PUBLIC_KEY_LENGTH:
+        data = b"\x04" + data
+    return data

ecies/keys/private.py

@@ -0,0 +1,58 @@
+from __future__ import annotations
+
+from typing import Optional
+
+from ..config import EllipticCurve
+from ..utils import decode_hex, derive_key
+from .helper import (
+    bytes_to_int,
+    get_public_key,
+    get_shared_point,
+    get_valid_secret,
+    is_valid_secret,
+)
+from .public import PublicKey
+
+
+class PrivateKey:
+    def __init__(self, curve: EllipticCurve, secret: Optional[bytes] = None):
+        self._curve = curve
+        if not secret:
+            self._secret = get_valid_secret(curve)
+        elif is_valid_secret(curve, secret):
+            self._secret = secret
+        else:
+            raise ValueError(f"Invalid {curve} secret key")
+        self._public_key = PublicKey(curve, get_public_key(curve, self._secret))
+
+    def __eq__(self, value):
+        return self._secret == value._secret if isinstance(value, PrivateKey) else False
+
+    @classmethod
+    def from_hex(cls, curve: EllipticCurve, sk_hex: str) -> PrivateKey:
+        """
+        For secp256k1, `sk_hex` can only be 32 bytes. `0x` prefix is optional.
+        """
+        return cls(curve, decode_hex(sk_hex))
+
+    def to_hex(self) -> str:
+        return self._secret.hex()
+
+    @property
+    def secret(self) -> bytes:
+        return self._secret
+
+    @property
+    def public_key(self) -> PublicKey:
+        return self._public_key
+
+    def to_int(self) -> int:
+        return bytes_to_int(self._secret)
+
+    def multiply(self, pk: PublicKey, compressed: bool = False) -> bytes:
+        return get_shared_point(self._curve, self._secret, pk.to_bytes(), compressed)
+
+    def encapsulate(self, pk: PublicKey, compressed: bool = False) -> bytes:
+        sender_point = self.public_key.to_bytes(compressed)
+        shared_point = self.multiply(pk, compressed)
+        return derive_key(sender_point + shared_point)

ecies/keys/public.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from ..config import EllipticCurve
+from ..utils import decode_hex, derive_key
+from .helper import convert_public_key
+
+if TYPE_CHECKING:
+    from .private import PrivateKey
+
+
+class PublicKey:
+    def __init__(self, curve: EllipticCurve, data: bytes):
+        self._curve = curve
+        compressed = convert_public_key(curve, data, True)
+        uncompressed = convert_public_key(curve, data, False)
+        self._data = compressed
+        self._data_uncompressed = (
+            uncompressed if len(compressed) != len(uncompressed) else b""
+        )
+
+    def __eq__(self, value):
+        return self._data == value._data if isinstance(value, PublicKey) else False
+
+    @classmethod
+    def from_hex(cls, curve: EllipticCurve, pk_hex: str) -> PublicKey:
+        """
+        For secp256k1, `pk_hex` can be 33(compressed)/65(uncompressed)/64(ethereum) bytes
+        """
+        return cls(curve, decode_hex(pk_hex))
+
+    def to_hex(self, compressed: bool = False) -> str:
+        """
+        For secp256k1, `pk_hex` can be 33(compressed)/65(uncompressed) bytes
+        """
+        return self.to_bytes(compressed).hex()
+
+    def to_bytes(self, compressed: bool = False) -> bytes:
+        """
+        For secp256k1, return uncompressed public key (65 bytes) by default
+        """
+        return self._data if compressed else self._data_uncompressed
+
+    def decapsulate(self, sk: PrivateKey, compressed: bool = False) -> bytes:
+        sender_point = self.to_bytes(compressed)
+        shared_point = sk.multiply(self, compressed)
+        return derive_key(sender_point + shared_point)

ecies/utils/__init__.py

@@ -1,5 +1,5 @@
-from .elliptic import bytes2pk, decapsulate, encapsulate, generate_key, hex2pk, hex2sk
-from .eth import generate_eth_key, to_eth_address, to_eth_public_key
+from .elliptic import decapsulate, encapsulate, generate_key, hex2pk, hex2sk
+from .eth import generate_eth_key, to_eth_address
 from .hash import derive_key, sha256
 from .hex import decode_hex
 from .symmetric import sym_decrypt, sym_encrypt
@@ -10,13 +10,11 @@
     "generate_key",
     "hex2sk",
     "hex2pk",
-    "bytes2pk",
-    "decapsulate",
     "encapsulate",
+    "decapsulate",
     # eth
     "generate_eth_key",
     "to_eth_address",
-    "to_eth_public_key",
     # hex
     "decode_hex",
     # hash