lockbox.py

import json
import base64
import os
from argparse import ArgumentParser
import getpass
import pprint

from Crypto.Cipher import AES
from Crypto import Random

random = Random.new()


def generate_key():
    key = random.read(AES.key_size[0])
    return base64.b64encode(key)


def get_key():
    return getpass.getpass("Key? ")


def get_input(msg):
    valid = False
    s = None
    while not valid:
        s = raw_input(msg)
        s.strip()
        valid = bool(s)
    return s

def open_lockbox(path, key):
    with open(path) as f:
        iv = f.read(AES.block_size)
        cipher = AES.AESCipher(key, AES.MODE_CFB, iv)
        return json.loads(cipher.decrypt(f.read()))


def commit_lockbox(path, key, payload):
    iv = random.read(AES.block_size)
    cipher = AES.AESCipher(key, AES.MODE_CFB, iv)
    tmp_path = path + '.tmp'
    f = open(tmp_path, "w")
    f.write(iv)
    f.write(cipher.encrypt(json.dumps(payload)))
    f.flush()
    os.fsync(f.fileno())
    f.close()
    os.rename(tmp_path, path)


def init(args):
    key = generate_key()
    commit_lockbox(args.lockbox, key, {})
    print u"Key: {}".format(key)


def dump(args):
    key = get_key()
    payload = open_lockbox(args.lockbox, key)
    pprint.pprint(payload)


def add(args):
    key = get_key()
    payload = open_lockbox(args.lockbox, key)

    k = get_input("k? ")
    v = get_input("v? ")

    if not (k and v):
        raise ValueError("abort")

    payload[k] = v

    commit_lockbox(args.lockbox, key, payload)
    print u"set \"{}\" -> \"{}\"".format(k, v)


def remove(args):
    key = get_key()
    payload = open_lockbox(args.lockbox, key)

    k = get_input("k? ")
    if k not in payload:
        raise ValueError("abort")
    del payload[k]

    commit_lockbox(args.lockbox, key, payload)
    print u"removed \"{}\"".format(k)


if __name__ == "__main__":

    parser = ArgumentParser()
    subparsers = parser.add_subparsers()

    parser_init = subparsers.add_parser("init")
    parser_init.add_argument("lockbox")
    parser_init.set_defaults(func=init)

    parser_init = subparsers.add_parser("dump")
    parser_init.add_argument("lockbox")
    parser_init.set_defaults(func=dump)

    parser_init = subparsers.add_parser("add")
    parser_init.add_argument("lockbox")
    parser_init.set_defaults(func=add)

    parser_init = subparsers.add_parser("remove")
    parser_init.add_argument("lockbox")
    parser_init.set_defaults(func=remove)

    args = parser.parse_args()
    args.func(args)