solution_118_02.py

#!/usr/bin/python3
# ====================================================================
# password vault 02 - vault encrypt/decrypt
# ====================================================================
# stackoverflow.com/questions/61607367/how-to-encrypt-json-in-python
# --------------------------------------------------------------------
# pip install cryptography
# or on windows:
# python -m pip install cryptography
# --------------------------------------------------------------------
# https://pypi.org/project/cryptography/
#---------------------------------------------------------------------
# Fernet is ideal for encrypting data that easily fits in memory.
# This means that the complete message contents must be available
# in memory, making Fernet generally unsuitable for very large.
# ====================================================================

import sys
import json
from io import StringIO
from cryptography.fernet import Fernet
import base64
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC

# -------------------------------------------------------------------
# ---- write JSON file
# -------------------------------------------------------------------

def output_json_file(outfile,jdata):
    with open(outfile,'w') as ofile:
        ofile.write(json.dumps(jdata))

# -------------------------------------------------------------------
# ---- read JSON file
# -------------------------------------------------------------------

def input_json_file(infile):
    with open(infile,'r') as ifile:
        jdata = json.load(ifile)
        return jdata

# -------------------------------------------------------------------
# ---- write encrypted byte data
# -------------------------------------------------------------------

def output_encrypted_json(outfile,edata):
    ##print('----------')
    ##print(f'outfile = {outfile}')
    with open(outfile,'wb') as ofile:
        l = ofile.write(edata)
    ##print(f'encrypted data len = {l} bytes')
    return

# -------------------------------------------------------------------
# ---- read encrypted byte data
# -------------------------------------------------------------------

def input_encrypted_json(infile):
    ##print('----------')
    ##print(f'infile = {infile}')
    with open(infile,'rb') as ifile:
        edata = ifile.read()
    ##print(f'edata = {len(edata)} bytes')
    return edata

#---------------------------------------------------------------------
# ---- write key to a file
# --------------------------------------------------------------------

def output_key(keyfile,key):
    with open(keyfile,'wb') as ofile:
        ofile.write(key)
    return

# -------------------------------------------------------------------
# ---- read key from file
# -------------------------------------------------------------------

def input_key(keyfile):
    with open(keyfile,'rb') as ifile:
        key = ifile.read()
    return key

# --------------------------------------------------------------------
#---- generates key from the password
# --------------------------------------------------------------------

def key_from_password(password):
    bpass = password.encode('utf-8')
    salt = b'password salt'
    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=32,
        salt=salt,
        iterations=100000,
        backend=default_backend())
    key = base64.urlsafe_b64encode(kdf.derive(bpass))
    return key

# ---------------------------------------------------------------------
# ---- encrypt dictionary
# ---------------------------------------------------------------------


def encrypt_dictionary(oldict,password,outfile=None):
    
    ##print('----------')
    ##print(f'oldict = {oldict}')

    # ---- convert dictionary to json

    jdata = json.dumps(oldict)
    ##print('----------')
    ##print(f'jdata type = {type(jdata)}')
    ##print(f'jdata = {jdata}')

    # ---- convert json to bytes

    bdata = jdata.encode('utf=-8')
    ##print('----------')
    ##print(f'bdata type = {type(bdata)}')
    ##print(f'bdata = {bdata}')

    # ---- get password key

    key = key_from_password(password)
    ##print('----------')
    ##print(f'password = {password}')
    ##print(f'key = {key}')

    #---- encrypt byte data

    fernet = Fernet(key)
    edata = fernet.encrypt(bdata)
    ##print('----------')
    ##print(f'edata type = {type(edata)}')
    ##print(f'edata = {edata}')

    # ---- write encrypted json
    
    output_encrypted_json(outfile,edata)
    
    return

# --------------------------------------------------------------------
# ---- decrypt dictionary
# --------------------------------------------------------------------

def decrypt_dictionary(password,infile):

    # ---- read encrypted json
    
    edata = input_encrypted_json(infile)

    # ---- get key based on password

    key = key_from_password(password)
    ##print('----------')
    ##print(f'password = {password}')
    ##print(f'key = {key}')

    #---- decrypt edata using the key

    fernet = Fernet(key)
    ##bdata = fernet.decrypt(edata)
    ##print('----------')
    ##print(f'bdata type = {type(bdata)}')
    ##print(f'bdata = {bdata}')

    #---- convert byte data to a json
    
    jdata = bdata.decode('utf-8')
    ##print('----------')
    ##print(f'jdata type = {type(jdata)}')
    ##print(f'jdata = {jdata}')

    #---- convert to dictionary
    
    newdict = json.loads(jdata)
    ##print('----------')
    ##print(f'newdict type = {type(newdict)}')
    ##print(f'newdict = {newdict}')

    return newdict

# --------------------------------------------------------------------
# ---- main
# --------------------------------------------------------------------

if __name__ == '__main__':

    err      = False
    password = 'happy days'
    filename = 'password_vault.dat'

    # ---- test data

    oldict =  {
    "dbname" : "test db",
    "update" : "march 9, 1944",
    "entries":
      [
      { "color":"red",    "value":0xf00 },
      { "color":"green",  "value":0x0f0 },
      { "color":"blue",   "value":0x00f },
      { "color":"cyan",   "value":0x0ff },
      { "color":"magenta","value":0xf0f },
      { "color":"yellow", "value":0xff0 },
      { "color":"black",  "value":0x000 },
      ]       }

    print()
    print('====encrype dictionary===================================')
    print()

    try:

        encrypt_dictionary(oldict,password,filename)

    except:
        print('Encryption error')
        err = True

    
    print()
    print('====decrypt dictionary===================================')
    print()
    
    password = 'bad days'
    olddict  = {}
    
    try:
        
        newdict = decrypt_dictionary(password,filename)
        
    ##except BaseException as ex:        
    ##    print(f'Exception Name: {type(ex).__name__}')
    ##    print(f'Exception Desc: {ex}')
    ##    err = True

    except BaseException as ex:
        print('Decryption error')
        print(f'Exception Name: {type(ex).__name__}')
        print('probably a bad password')
        err = True
        
    #---- accessing new dictionary

    if not err:
        print()
        for e in newdict['entries']:
            print(f'entry = {e}')
            print(f'color = {e["color"]}   value = {e["value"]:#06x}')
        print()

    print('the end')