from cryptography.fernet import Fernet
import pickle
import io

def writeKey(filename:str) -> Fernet:
  """
  Generates a key and save it into a file
  """
  key = Fernet.generate_key()
  with open(f'{filename}.key', 'wb') as key_file:
    key_file.write(key)

  return Fernet(key)

def generateKey() -> Fernet:
  """
  Generates a key and returns it
  """
  key = Fernet.generate_key()
  return Fernet(key)

def loadKey(filename:str) -> Fernet:
  """
  Loads the key from the current directory named `key.key`
  """
  with open(f'{filename}.key', 'rb') as f:
    data = f.read()
  return Fernet(data)

def dumps(obj, compress=None, encrypt:Fernet=None) -> bytes:
  """
  Dumps the object into a byte stream
  """
  data = pickle.dumps(obj)
  if compress:
    data = compress.compress(data)
  if encrypt:
    data = encrypt.encrypt(data)

  return data

def dump(obj, file: io.BytesIO, compress=None, encrypt:Fernet=None) -> None:
  """
  Dumps the object into a file
  """
  data = pickle.dumps(obj)
  data = dumps(data, compress, encrypt)

  file.write(data)
  file.flush()

def loads(data, compress=None, encrypt:Fernet=None) -> object:
  """
  Loads the object from a byte stream
  """
  if encrypt:
    data = encrypt.decrypt(data)
  if compress:
    data = compress.decompress(data)

  return pickle.loads(data)

def load(file: io.BytesIO, compress=None, encrypt:Fernet=None) -> object:
  """
  Loads the object from a file
  """
  data = file.read()
  data = loads(data, compress, encrypt)

  return pickle.loads(data)

if __name__ == "__main__":
  import lzma

  data = {'a': 'A', 'b': (1, 2, 3, 4, 5)}
  #writeKey('key')
  key = loadKey('key')
  dataCompressedAndEncrypted = dumps(data, compress=lzma, encrypt=key)
  print(dataCompressedAndEncrypted)

  del key

  key = loadKey('key')
  dataCompressedAndEncrypted = loads(dataCompressedAndEncrypted, compress=lzma, encrypt=key)
  print(dataCompressedAndEncrypted)