Add public/private key encryption and decryption

.gitignore

@@ -8,6 +8,7 @@ __pycache__
 /.coverage
 /.pytest_cache
 /bitcash.egg-info
+/BitCash.egg-info
 /build
 /dist
 /docs/build

HISTORY.rst

@@ -16,6 +16,9 @@ Unreleased (see `master <https://github.com/ofek/bitcash>`_)
 - NetworkAPI.get_tx_amount() is now working and properly handles
   backends returning string or decimal values.
 
+- Add public/private key encryption using ECIES encryption decryption
+  methods
+
 0.5.2 (2018-05-16)
 ------------------
 

bitcash/crypto.py

@@ -1,5 +1,8 @@
-from hashlib import new, sha256 as _sha256
+from hashlib import new, sha256 as _sha256, sha512 as _sha512
+import base64
+import hmac
 
+import pyaes
 from coincurve import PrivateKey as ECPrivateKey, PublicKey as ECPublicKey
 
 
@@ -20,3 +23,110 @@ def ripemd160_sha256(bytestr):
 
 
 hash160 = ripemd160_sha256
+
+
+def sha512(bytestr):
+    return _sha512(bytestr).digest()
+
+
+# ECIES encryption/decryption methods; AES-128-CBC with PKCS7 is used
+# as the cipher; hmac-sha256 is used as the mac
+# Implementation follows the Electron-Cash implementation of the same
+def aes_encrypt_with_iv(key, iv, data):
+    """Provides AES-CBC encryption of data with key and iv
+
+    :param key: key for the encryption
+    :type key: ``bytes``
+    :param iv: Initialisation vector for the encryption
+    :type iv: ``bytes``
+    :param data: the data to be encrypted
+    :type data: ``bytes``
+    """
+    aes_cbc = pyaes.AESModeOfOperationCBC(key, iv=iv)
+    aes = pyaes.Encrypter(aes_cbc)
+    # empty aes.feed() flushes buffer
+    return aes.feed(data) + aes.feed()
+
+
+def aes_decrypt_with_iv(key, iv, data):
+    """Provides AES-CBC decryption of data with key and iv
+
+    :param key: key for the decryption
+    :type key: ``bytes``
+    :param iv: Initialisation vector for the decryption
+    :type iv: ``bytes``
+    :param data: the data to be decrypted
+    :type data: ``bytes``
+    :raises ValueError: if incorrect ``key`` or ``iv`` give a padding error
+                        during decryption
+    """
+    # assert_bytes(key, iv, data)
+    aes_cbc = pyaes.AESModeOfOperationCBC(key, iv=iv)
+    aes = pyaes.Decrypter(aes_cbc)
+    try:
+        # empty aes.feed() flushes buffer
+        return aes.feed(data) + aes.feed()
+    except ValueError:
+        raise ValueError('Invalid key or iv')
+
+
+def ecies_encrypt(message, pubkey):
+    """Encrypt message with the given pubkey
+
+    :param message: the message to be encrypted
+    :type message: ``bytes``
+    :param pubkey: the public key to be used
+    :type pubkey: ``bytes``
+    """
+    pk = ECPublicKey(pubkey)
+
+    # random key
+    ephemeral = ECPrivateKey()
+    ecdh_key = pk.multiply(ephemeral.secret).format()
+    key = sha512(ecdh_key)
+
+    # aes key and iv, and hmac key
+    iv, key_e, key_m = key[0:16], key[16:32], key[32:]
+    ciphertext = aes_encrypt_with_iv(key_e, iv, message)
+    encrypted = (
+        b'BIE1'
+        + ephemeral.public_key.format()
+        + ciphertext
+    )
+    mac = hmac.new(key_m, encrypted, _sha256).digest()
+
+    return base64.b64encode(encrypted + mac)
+
+
+def ecies_decrypt(encrypted, secret):
+    """Decrypt the encrypted message with the given private-key secret
+
+    :param encrypted: the message to be decrypted
+    :type encrypted: ``bytes``
+    :param secret: the private key secret to be used
+    :type secret: ``bytes``
+    :raises ValueError: if magic bytes or HMAC bytes are invalid
+    """
+    encrypted = base64.b64decode(encrypted)
+    if len(encrypted) < 85:
+        raise ValueError('Invalid cipher length')
+
+    # splitting data
+    magic = encrypted[:4]
+    ephemeral_pubkey = ECPublicKey(encrypted[4:37])
+    ciphertext = encrypted[37:-32]
+    mac = encrypted[-32:]
+    if magic != b'BIE1':
+        raise ValueError('Invalid magic bytes')
+
+    # retrieving keys
+    ecdh_key = ephemeral_pubkey.multiply(secret).format()
+    key = sha512(ecdh_key)
+    iv, key_e, key_m = key[0:16], key[16:32], key[32:]
+
+    # validating hmac
+    if mac != hmac.new(key_m, encrypted[:-32], _sha256).digest():
+        raise ValueError("Invalid HMAC bytes")
+
+    # decrypting
+    return aes_decrypt_with_iv(key_e, iv, ciphertext)

bitcash/wallet.py

@@ -1,6 +1,6 @@
 import json
 
-from bitcash.crypto import ECPrivateKey
+from bitcash.crypto import ECPrivateKey, ecies_encrypt, ecies_decrypt
 from bitcash.curve import Point
 from bitcash.exceptions import InvalidNetwork
 from bitcash.format import (
@@ -143,6 +143,22 @@ def is_compressed(self):
     def __eq__(self, other):
         return self.to_int() == other.to_int()
 
+    def encrypt_message(self, message):
+        """Encrypt message with the instance's public key
+
+        :param message: the message to be encrypted
+        :type message: ``bytes``
+        """
+        return ecies_encrypt(message, self.public_key)
+
+    def decrypt_message(self, encrypted):
+        """Decrypt the encrypted message using the instance's private key
+
+        :param encrypted: the message to be decrypted
+        :type encrypted: ``bytes``
+        """
+        return ecies_decrypt(encrypted, self._pk.secret)
+
 
 class PrivateKey(BaseKey):
     """This class represents a BitcoinCash private key. ``Key`` is an alias.

setup.py

@@ -48,7 +48,7 @@
         'Programming Language :: Python :: Implementation :: PyPy'
     ],
 
-    install_requires=['coincurve>=4.3.0', 'requests'],
+    install_requires=['coincurve>=4.3.0', 'requests', 'pyaes'],
     extras_require={
         'cli': ('appdirs', 'click', 'privy', 'tinydb'),
         'cache': ('lmdb', ),

tests/test_crypto.py

@@ -0,0 +1,60 @@
+import pytest
+
+from bitcash.crypto import (
+    aes_encrypt_with_iv,
+    aes_decrypt_with_iv,
+    ecies_encrypt,
+    ecies_decrypt
+)
+
+
+KEY_AES = b'$\x99\xd7-\x10\x1aY\xa4"\xd6\x9c\x7f\x0f\xd7\x0aT'
+
+KEY_AES2 = b'$\x99\xd7-\x10\x1aY\xa4"\xd6\x9c\x7f\x0f\xd7\x0bT'
+
+IV = b'\\\xdf\x8c\xdd\xebA\xa6\x7f\xfa\xbfq\x0cn\xccr\xc8'
+
+PUBKEY = (
+    b"\x03=\\(u\xc9\xbd\x11hu\xa7\x1a]\xb6L\xff\xcb\x139k\x16=\x03\x9b"
+    + b"\x1d\x93'\x82H\x91\x80C4"
+)
+
+SECRET = (
+    b'\xc2\x8a\x9f\x80s\x8fw\rRx\x03\xa5f\xcfo\xc3\xed\xf6\xce\xa5\x86'
+    + b'\xc4\xfcJR#\xa5\xady~\x1a\xc3'
+)
+
+SECRET2 = (
+    b'\xc2\x8a\x9f\x80s\x8fw\rRx\x03\xa5f\xcfo\xc3\xed\xf6\xce\xa5\x86'
+    + b'\xc4\xfcJR#\xa5\xady~\x1a\xc4'
+)
+
+
+class TestAes:
+    def test_aes_success(self):
+        message = b'test'
+        encrypted_message = aes_encrypt_with_iv(KEY_AES, IV, message)
+        decrypted_message = aes_decrypt_with_iv(KEY_AES, IV, encrypted_message)
+        assert message == decrypted_message
+
+    def test_aes_fail(self):
+        message = b'test'
+        encrypted_message = aes_encrypt_with_iv(KEY_AES, IV, message)
+        with pytest.raises(ValueError):
+            decrypted_message = aes_decrypt_with_iv(KEY_AES2,
+                                                    IV,
+                                                    encrypted_message)
+
+
+class TestEcies:
+    def test_ecies_success(self):
+        message = b'test'
+        encrypted_message = ecies_encrypt(message, PUBKEY)
+        decrypted_message = ecies_decrypt(encrypted_message, SECRET)
+        assert message == decrypted_message
+
+    def test_ecies_fail(self):
+        message = b'test'
+        encrypted_message = ecies_encrypt(message, PUBKEY)
+        with pytest.raises(ValueError):
+            decrypted_message = ecies_decrypt(encrypted_message, SECRET2)

tests/test_wallet.py

@@ -148,6 +148,21 @@ def test_equal(self):
             WALLET_FORMAT_COMPRESSED_MAIN
         )
 
+    def test_encrypt_decrypt(self):
+        # successful test
+        key = BaseKey()
+        message = b'test'
+        encrypted_message = key.encrypt_message(message)
+        decrypted_message = key.decrypt_message(encrypted_message)
+        assert message == decrypted_message
+
+        # failed test
+        key = BaseKey()
+        message = b'test'
+        encrypted_message = key.encrypt_message(message)
+        with pytest.raises(ValueError):
+            decrypted_message = BaseKey().decrypt_message(encrypted_message)
+
 
 class TestPrivateKey:
     def test_alias(self):