maintainsensitive.py

"""                                                                                                                                                                                                               
    Description: Rotate encryption key to keep sensitive data more secure
                 1. Read 'secret key' and decrypt the 'secret file'.
                 2. Generate new 'secret key' and encrypt the 'secret file' with new key and store it back to the 'secret file'.
                 3. Update new 'secret key' into the hidden file.
"""

import os
import sys 
import yaml

from cryptography.fernet import Fernet

def generate_encrypt_newkey(filename, secretfile):
    """ 
        Description: Create new key and update encryption key.

        @filename   : File that contains data that needs to be protected/encrypted
        @secretfile : File that holds current key that needs to be rotated
                      periodically.
    """

    keycfg = {}

    ##--------------------------------------------##
    ## Check if secretfile exists else create new ## 
    ##--------------------------------------------##
    if os.path.exists(secretfile):
        with open(secretfile) as fd: 
            keycfg = yaml.safe_load(fd)

    ##---------------------------------------------------##
    ## Prepare a new key for encryption & store current  ## 
    ## key for decryption                                ##
    ##---------------------------------------------------##
    keycfg['prevkey'] = keycfg.get('currentkey', None)
    keycfg['currentkey'] = Fernet.generate_key()

    fernetnew = Fernet(keycfg['currentkey'])

    if keycfg['prevkey']:
        fernetold = Fernet(keycfg['prevkey'])
        with open(filename, 'rb') as fd: 
            encryptedata = fernetnew.encrypt(fernetold.decrypt(fd.read()))

    else:
        with open(filename, 'rb') as fd: 
            encryptedata = fernetnew.encrypt(fd.read())

    ## --------------------- ##
    ## Update encrypted file ##
    ## --------------------- ##
    with open(filename, 'wb') as fd:
        fd.write(encryptedata)

    ## ------------------ ##
    ## Update secret keys ## 
    ## ------------------ ##
    with open(secretfile, 'w') as fdwr:
        fdwr.write(yaml.safe_dump(keycfg, explicit_start=True))

def decrypt_and_print(filename, secretfile):
    """
        Description: Test if current key could decrypt your file

        @filename   : File that contains data that needs to be protected/encrypted
        @secretfile : File that holds current key that needs to be rotated
                      periodically.
    """
    keycfg = {}

    if os.path.exists(secretfile):
        with open(secretfile) as fd:
            keycfg = yaml.safe_load(fd)

    currentkey = keycfg.get('currentkey', None)

    if currentkey:
        fernetobj = Fernet(currentkey)

        with open(filename, 'rb') as fd:
            decrypted = fernetobj.decrypt(fd.read())
            print(decrypted)

if __name__ == '__main__':

    #------------------------------------------#
    # How to use following options             # 
    # 1. python3.9 maintainsensitive.py rotate # 
    # 2. python3.9 maintainsensitive.py test   #
    #------------------------------------------#
    if len(sys.argv) == 2 and sys.argv[1] == 'rotate':
        generate_encrypt_newkey("test.txt", ".sensitive.yml")
        print("INFO: Key file rotated successfully")

    elif len(sys.argv) == 2 and sys.argv[1] == 'test':
        decrypt_and_print("test.txt", ".sensitive.yml")

    else:
        filename = __file__.split("/")[-1]                                                                                                                                                                         

        print("\n" + "-"*80)
        print("Periodically rotate key on an encrypted file".center(80))
        print("-"*80 + "\n")
        print("1. To rotate keys on a sensitive file\n")                                                                                                                                                           
        print(f"\tpython {filename} rotate\n\n")
        print("2. Test if key is effective in reading sensitive data\n")
        print(f"\tpython {filename} test\n\n")
        print("."*80)