KeyUtils.py

from ecdsa import SigningKey, VerifyingKey, SECP256k1
from Crypto.PublicKey import RSA
from Crypto.Cipher import PKCS1_OAEP
from base64 import b64encode, b64decode
from json import loads


# comments:
# BTC uses base58 encode to avoid special characters. It requires another library. So I use base64 here.


# -------------------------------------------ECC key functions-------------------------------------------------------


def generate_private_key_ecc():
    """
    generate a private key using SECP256k1 curve and sha1 encryption
    :return: a private key of type ecdsa.SigningKey
    """
    return SigningKey.generate(curve=SECP256k1)


def generate_public_key_ecc(private_key):
    """
    given a private key, generate its public key
    :param private_key: a private key of type ecdsa.SigningKey
    :return: a public key of type ecdsa.VerifyingKey
    """
    return private_key.get_verifying_key()


def key_to_address_ecc(key):
    """
    encode a key to address using base64
    :param key:
    :return:
    """
    return b64encode(key.to_string())


def address_to_key_ecc(address):
    """
    decode an address to key using base64 in hex format
    :param address:
    :return:
    """
    return b64decode(address, validate=True).hex()


def import_public_key_ecc(s):
    """
    import public ecc key from a hex string
    :param s: a hex string generated by export_keys()
    :return: a VerifyingKey object
    """
    return VerifyingKey.from_string(bytes.fromhex(s), curve=SECP256k1)


# -------------------------------------------RSA key functions-------------------------------------------------------


def generate_private_key_rsa(key_length=1024):
    """
    Generate a private rsa key with certain length
    :param key_length:
    :return:
    """
    return RSA.generate(key_length)


def generate_public_key_rsa(private_key):
    """
    Generate a public key from its private key
    :param private_key:
    :return:
    """
    return private_key.publickey()


def export_public_key_rsa(public_key):
    """
    Export the public rsa key to string format
    :param public_key: a public key object
    :return: a public key in string format
    """
    return public_key.exportKey().decode('utf-8')


def import_public_key_rsa(public_key):
    """
    Import a public key in string format, convert it to key object
    :param public_key: a string public key
    :return: a public key object
    """
    return RSA.importKey(public_key).publickey()

# ---------------------------------------------General functions----------------------------------------------


def string_to_keys(s):
    """
    Import keys from string.
    :param s: a json style string created by export_keys() method
    :return: ecc key and rsa key objects
    """
    j = loads(s)
    public_key_ecc = import_public_key_ecc(j["public_key_ecc"])
    public_key_rsa = import_public_key_rsa(j["public_key_rsa"])
    return public_key_ecc, public_key_rsa


def encrypt_message(public_key_recipient, message):
    """
    Encrypt a message using recipient's public key
    :param public_key_recipient: public key of recipient
    :param message: message in string format
    :return: an rsa encrypted message
    """
    return PKCS1_OAEP.new(public_key_recipient).encrypt(message.encode('utf-8')).hex()


def test():
    sk = generate_private_key_ecc()  # generate private key

    vk = generate_public_key_ecc(sk)  # generate public key

    sadd = key_to_address_ecc(sk)  # convert private key to address

    vadd = key_to_address_ecc(vk)  # convert public key to address

    print('\nprivate key: ' + sk.to_string().hex())

    print('\nconvert to address: ')
    print(sadd)

    print('\nconvert back to key: ' + address_to_key_ecc(sadd))

    print('\npublic key: ' + vk.to_string().hex())

    print(type(vk.to_string().hex()))

    print('\nconvert to address: ')
    print(vadd)

    print('\nconvert back to key: ' + address_to_key_ecc(vadd))

    message = "hello world."  # get a message
    print('\nmessage: ' + message)

    signature = sk.sign(message.encode('utf-8'))  # sign a signature using private key
    print('\nsign a signature: ' + signature.hex())

    # verify the message using the public key
    print('\nverify the message and signature using the public key: ')
    print(vk.verify(signature, message.encode('utf-8')))

    print(vk)


if __name__ == '__main__':
    test()