from tinyec import registry
from Crypto.Cipher import AES
import hashlib, secrets, binascii

def encrypt_AES_GCM(msg, secretKey):
    aesCipher = AES.new(secretKey, AES.MODE_GCM)
    ciphertext, authTag = aesCipher.encrypt_and_digest(msg)
    return (ciphertext, aesCipher.nonce, authTag)

def decrypt_AES_GCM(ciphertext, nonce, authTag, secretKey):
    aesCipher = AES.new(secretKey, AES.MODE_GCM, nonce)
    plaintext = aesCipher.decrypt_and_verify(ciphertext, authTag)
    return plaintext

def ecc_point_to_256_bit_key(point):
    sha = hashlib.sha256(int.to_bytes(point.x, 32, 'big'))
    sha.update(int.to_bytes(point.y, 32, 'big'))
    return sha.digest()


def encrypt_ECC(curve, msg, pubKey):
    ciphertextPrivKey = secrets.randbelow(curve.field.n)
    sharedECCKey = ciphertextPrivKey * pubKey
    secretKey = ecc_point_to_256_bit_key(sharedECCKey)
    ciphertext, nonce, authTag = encrypt_AES_GCM(msg, secretKey)
    ciphertextPubKey = ciphertextPrivKey * curve.g
    return (ciphertext, nonce, authTag, ciphertextPubKey)

def decrypt_ECC(encryptedMsg, privKey):
    (ciphertext, nonce, authTag, ciphertextPubKey) = encryptedMsg
    sharedECCKey = privKey * ciphertextPubKey
    secretKey = ecc_point_to_256_bit_key(sharedECCKey)
    plaintext = decrypt_AES_GCM(ciphertext, nonce, authTag, secretKey)
    return plaintext

def main():
	curve = registry.get_curve('brainpoolP256r1')
	msg = b'Text to be encrypted by ECC public key and ' \
		  b'decrypted by its corresponding ECC private key'
	print("original msg:", msg)
	privKey = secrets.randbelow(curve.field.n)
	pubKey = privKey * curve.g

	encryptedMsg = encrypt_ECC(curve, msg, pubKey)
	encryptedMsgObj = {
		'ciphertext': binascii.hexlify(encryptedMsg[0]),
		'nonce': binascii.hexlify(encryptedMsg[1]),
		'authTag': binascii.hexlify(encryptedMsg[2]),
		'ciphertextPubKey': hex(encryptedMsg[3].x) + hex(encryptedMsg[3].y % 2)[2:]
	}
	print("encrypted msg:", encryptedMsgObj)

	decryptedMsg = decrypt_ECC(encryptedMsg, privKey)
	print("decrypted msg:", decryptedMsg)

if __name__ == "__main__":
	main()