#!/usr/bin/env python3

import os, sys, tempfile,getopt
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(arg, '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()
    # We open the file, which contains the encrypted content
    crypt = open(args[0], 'rb').read()
    # And we decrypt it, and write it in the temp file
    clear = f.decrypt(crypt)
    tmp.write(clear)
    tmp.flush()
    # 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()
    os.unlink(tmp.name)
    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.json')

if __name__ == '__main__':
    main()