from cryptography.exceptions import *
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import padding
from cryptography.hazmat.backends import default_backend

def generateKeys():
    private_key = rsa.generate_private_key(public_exponent=65537,key_size=2048)
    public_key = private_key.public_key()

    return private_key, public_key

def keysToBytes(private_key, public_key):
    prv_ser = privateKeyToBytes(private_key)
    pbc_ser = publicKeyToBytes(public_key)
    
    return prv_ser, pbc_ser

def publicKeyToBytes(public_key):
    pbc_ser = public_key.public_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PublicFormat.SubjectPublicKeyInfo)
    
    return pbc_ser

def privateKeyToBytes(private_key):
    prv_ser = private_key.private_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PrivateFormat.TraditionalOpenSSL,
        encryption_algorithm=serialization.NoEncryption())
    
    return prv_ser

def bytesToKeys(private_ser, public_ser):
    private_key = serialization.load_pem_private_key(
        private_ser,
        password=None
    )
    public_key = serialization.load_pem_public_key(
        public_ser
    )
    return private_key, public_key

def publicBytesToKey(public_ser):
    public_key = serialization.load_pem_public_key(
        public_ser,
        backend=default_backend()
    )
    return public_key

def privateBytesToKey(private_ser):
    private_key = serialization.load_pem_private_key(
        private_ser,
        password=None,
        backend=default_backend()
    )
    return private_key

def sign(message, private_key):
    message = bytes(str(message), 'utf-8')
    signature = private_key.sign(
        message,
        padding.PSS(mgf=padding.MGF1(hashes.SHA256()), salt_length=padding.PSS.MAX_LENGTH),
        hashes.SHA256()
        )
    return signature

def verify(message, signature, public_ser):
    message = bytes(str(message), 'utf-8')
    public_key = publicBytesToKey(public_ser)
    try:
        public_key.verify(
            signature,
            message,
            padding.PSS(mgf=padding.MGF1(hashes.SHA256()),
            salt_length=padding.PSS.MAX_LENGTH),
            hashes.SHA256()
            )
        return True
    except InvalidSignature:
         return False
    except:
        print("Error executing 'public_key.verify'")
        return False