import os
import json
import base64
from cryptography.fernet import Fernet, InvalidToken
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC

# https://cryptography.io/en/latest/fernet/#using-passwords-with-fernet

rescue_msg = {
    "readme": "The salt (bytes) and the encrypted object (bytes) are both stored in base64 format.\
               The key, 32 bytes long, is derived using PBKDF2 using the salt and your password.\
               Decrypt the object using the AES-CBC algorithm with the key.",
    "code_to_decrypt": """
                import base64
                from cryptography.fernet import Fernet
                from cryptography.hazmat.primitives import hashes
                from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
                def s2b(s:str)->bytes:
                    return base64.decodebytes(s.encode('utf-8'))
                salt = s2b(save_obj["salt"])
                obj_encrypted = s2b(save_obj["obj"])
                password = input("password:").encode('utf-8')
                kdf = PBKDF2HMAC(
                        algorithm=hashes.SHA256(),
                        length=32,
                        salt=salt,
                        iterations=480000,
                )
                key = base64.urlsafe_b64encode(kdf.derive(password))
                obj_decrypted = Fernet(key).decrypt(obj_encrypted).decode('utf-8')
                print(obj_decrypted)
            """
}

def get_fernet(password:bytes, salt:bytes):
    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=32,
        salt=salt,
        iterations=480000,
    )
    key = base64.urlsafe_b64encode(kdf.derive(password))
    return Fernet(key)

def b2s(b:bytes)->str:
    return base64.encodebytes(b).decode('utf-8')

def s2b(s:str)->bytes:
    # s must be constructed from b2s
    return base64.decodebytes(s.encode('utf-8'))

def save_obj(obj, password:bytes, salt:bytes) -> str:
    f = get_fernet(password, salt)
    save_obj = dict()
    save_obj["salt"] = b2s(salt)
    obj_encrypted = f.encrypt(json.dumps(obj).encode('utf-8'))
    save_obj["obj"] = b2s(obj_encrypted)
    save_obj["rescue"] = rescue_msg
    return json.dumps(save_obj,indent=4)

def load_obj(s, password):
    save_obj = json.loads(s)
    salt = s2b(save_obj["salt"])
    obj_encrypted = s2b(save_obj["obj"])
    f = get_fernet(password, salt)
    try:
        obj_decrypted = f.decrypt(obj_encrypted)
    except InvalidToken:
        print("Invalid Password")
        exit(-1)
    obj = json.loads(obj_decrypted.decode('utf-8'))
    return obj

def perform_encrypt(infile, outfile):
    password = input("password:").encode('utf-8')
    salt = os.urandom(16)
    obj = json.load(open(infile,'r'))
    d = save_obj(obj, password, salt)
    f = open(outfile,'w'); f.write(d); f.close()

def perform_decrypt(infile, outfile):
    f = open(infile,'r'); d=f.read(); f.close()
    password = input("password:").encode('utf-8')
    obj = load_obj(d, password)
    json.dump(obj, open(outfile,'w'), indent=4)

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('mode', choices=['e', 'd'], help="e for encrypt, d for decrypt")
    parser.add_argument('infile')
    parser.add_argument('outfile')
    args = parser.parse_args()
    if args.mode == 'e':
        perform_encrypt(args.infile, args.outfile)
    elif args.mode == 'd':
        perform_decrypt(args.infile, args.outfile)