#!/usr/bin/env python3

import os, sys, tempfile, getopt, json, yaml
from cryptography.fernet import Fernet

def main():
    try:
        opts, args = getopt.getopt(sys.argv[1:],'k:K:n')
    except:
        usage()
        sys.exit(2)
    
    key = ''
    for opt, val in opts:
        if opt == '-n':
            # Print a new key and exit
            print(Fernet.generate_key().decode())
            exit(0)
        if opt == '-k':
            key = val
        elif opt == '-K':
            key = open(val, 'rb').read()
    
    if key == '':
        usage()
        sys.exit(2)

    # Create a fernet object using our key
    f = Fernet(key)
    # This temp file will hold the decrypted content while we edit it
    tmp = tempfile.NamedTemporaryFile(suffix='.yml')
    # We open the file which contains the encrypted content, if it exists
    if os.path.exists(args[0]):
        crypt = open(args[0], 'rb').read()
        # And we decrypt it, and write it in the temp file
        clear = f.decrypt(crypt)
        try:
            convert = yaml.dump(json.loads(clear)).encode('utf-8')
            clear = convert
        except:
            pass
        tmp.write(clear)
        tmp.flush()
    loop = 1
    while loop == 1:
      # Now, lets open our favorite editor to edit the file
      os.system(os.getenv('EDITOR', 'vim') + ' ' + tmp.name)
      # We closed the editor, we just have to open the cleartext file, encrypt its content
      # and save it
      clear = open(tmp.name, 'rb').read()
      loop = 0
      # Validate YAML data. Edit agin if not valid
      try:
          yaml.safe_load(clear)
      except:
          loop = 1
    wcrypt = open(args[0], 'wb')
    wcrypt.write(f.encrypt(clear))
    wcrypt.flush()
    wcrypt.close()

def usage():
    print(argv[0] + ' [-k secret] [-K ./secret.txt] variables.txt')

if __name__ == '__main__':
    main()