'''
Implement the fernet algorithm from pycryptography
'''
# Import python libs
import hashlib
import base64
# Import third party libs
import cryptography.fernet


def gen(hub):
    '''
    Generate a raw key to use for fernet. This is used by seals that require
    a key to be generated by the cipher, like shamir and not by seals that
    derive the key weakly from use input, like passwd
    '''
    return cryptography.fernet.Fernet.generate_key()


async def setup(hub, unit, seal_raw):
    '''
    Using the seal data derive a key
    '''
    seal_raw = seal_raw.encode()
    seal = hub.takara.UNITS[unit]['seal']
    if hasattr(hub, f'takara.seal.{seal}.ENCODE_RAW'):
        if getattr(hub, f'takara.seal.{seal}.ENCODE_RAW'):
            key = base64.urlsafe_b64encode(hashlib.blake2s(seal_raw).digest())
    elif hasattr(hub, f'takara.seal.{seal}.derive'):
        key = getattr(hub, f'takara.seal.{seal}.derive')(seal_raw)
    else:
        key = seal_raw
    hub.takara.CRYPT[unit] = {}
    hub.takara.CRYPT[unit]['crypter'] = cryptography.fernet.Fernet(key)


async def encrypt(hub, unit, data):
    '''
    Using the key from the given unit encrypt the raw bytes found in data
    and return the encryption string
    '''
    crypter = hub.takara.CRYPT[unit]['crypter']
    return crypter.encrypt(data)


async def decrypt(hub, unit, data):
    '''
    Using the key from the given unit decrypt the raw bytes found in data
    and return the clear string
    '''
    crypter = hub.takara.CRYPT[unit]['crypter']
    return crypter.decrypt(data)

