EJSON decryption in Python with pynacl

ejson.py

import json
import os
import pathlib
from base64 import b64decode
from binascii import unhexlify

from nacl.public import Box, PrivateKey, PublicKey


PRIVATE_KEYS_DIR = '/opt/ejson/keys'


def decrypt_message(msg, privkey_as_hex_str):
    privkey = unhexlify(privkey_as_hex_str)

    header, b64_encpub, b64_nonce, b64_box = msg.split(':')
    encpub = b64decode(b64_encpub)
    nonce = b64decode(b64_nonce)
    box = b64decode(b64_box)

    b = Box(PrivateKey(privkey), PublicKey(encpub))

    decrypted = b.decrypt(box, nonce)
    return decrypted.decode('utf-8')


def decrypt_dict(d, privkey_as_hex_str):
    return {name: decrypt_message(secret, privkey_as_hex_str) for name, secret in d.items()}


def fetch_private_key(public_key):
    path = pathlib.Path(PRIVATE_KEYS_DIR).joinpath(public_key)
    with path.open():
        return path.read_text().strip()


def load_from_filename(filename):
    with open(filename) as fh:
        data = json.load(fh)

    public_key = data.pop('_public_key')
    enc_environment = data.pop('environment', {})
    enc_secrets = data.copy()

    private_key = fetch_private_key(public_key)

    return {
        'environment': decrypt_dict(enc_environment, private_key),
        'secrets': decrypt_dict(enc_secrets, private_key),
    }


def load_into_environ_from_filename(filename):
    os.environ.update(load_from_filename(filename)['environment'])