encrypted-tables/peek.py

"""
Peek:
Operate on a list of dicts in the cmdline and store them (optionally encrypted) on the filesystem
I AM NOT A SECURITY EXPERT. NO GUARANTEE. USE AT YOUR OWN RISK
"""

import os
import json
import base64
from encrypt_json import load_obj, save_obj

def base64_encode(objs) -> str:
    # Storing in the illegible base64 format instead of clear text creates
    # A FALSE SENSE OF DATA SECURITY, SHOULD BE DISABlED

    # In Python, bytes objects are immutable sequences of single bytes
    # string <-e/d-> bytes <--base64 e/d--> bytes <-e/d-> string
    s = json.dumps(objs)
    return base64.encodebytes(s.encode('utf-8')).decode('utf-8')

def base64_decode(s:str):
    s = base64.decodebytes(s.encode('utf-8')).decode('utf-8')
    return json.loads(s)

def printraw(mem):
    for i,obj in enumerate(mem):
        print(i)
        for k,v in obj.items():
            print(k+":"+v,end='\t')
        print('')

def printtable(table):
    # Prints a table in markdown style
    table.insert(1,["---"]*len(table[0]))
    for row in table:
        print("|" + "|".join(row) + "|")

def printfmt(mem, keys, shortcuts):
    # Parse which keys are required
    key_order = list()
    list_others = False
    keys = keys.split()
    for c in keys[0]:
        if c == '+':
            list_others = True
        elif c in shortcuts.keys():
            key_order.append(shortcuts[c])
        else:
            print("unknown character shortcut '%s', cannot print formatted"%c)
            return
    key_order.extend(keys[1:])
    # List dicts that has the required keys
    table = list()
    table.append(["#"] + key_order + (["+"] if list_others else []))
    for i,obj in enumerate(mem):
        if all([key in obj for key in key_order]):
            row = [str(i)]
            for key in key_order:
                row.append(obj[key])
            if list_others:
                s = ""
                for k,v in obj.items():
                    if k not in key_order:
                        s += k+":"+v+";"
                row.append(s)
            table.append(row)
    printtable(table)

def inputkv(obj, shortcuts):
    while True:
        k = input("k:")
        if k=='':
            break # end of input
        if len(k)==1:
            try:
                k = shortcuts[k]
            except KeyError:
                print("unknown character shortcut %s, one-character key is reserved for shortcuts, \
                        choose another key"%k)
                continue # invalid input
        if k in obj:
            print("rewriting key %s, original value is "%k + obj[k])
        v = input("v:")
        obj[k] = v

def inputshortcuts():
    sc = {}
    while True:
        k = input("k:")
        if k=='':
            break
        if len(k)!=1:
            print("shortcut must be one-character")
            continue
        v = input("v:")
        sc[k] = v
    return sc

def loaddb(db_name, encode_method, password):
    if encode_method=="base64":
        try:
            f = open(db_name+".store",'r')
            obj = base64_decode(f.read())
            f.close()
        except Exception as e:
            print(type(e).__name__ + ":" + str(e))
            print("Corrupted save file")
            exit(-1)
    elif encode_method=="fernet":
        f = open(db_name+".json",'r')
        d=f.read(); f.close()
        obj = load_obj(d, password)
    else:
        print("unsupported encode method, exiting")
        exit(-1)
    return obj

def savedb(db_name, encode_method, obj, password):
    if encode_method=="base64":
        f = open(db_name+".store",'w')
        f.write(base64_encode(obj))
    elif encode_method=="fernet":
        f = open(db_name+".json",'w')
        salt = os.urandom(16)
        d = save_obj(obj, password, salt)
        f.write(d)
    else:
        print("unsupported encode method, exiting")
    f.close()

def interact(db_name, decode_method, encode_method, password, create):
    db_obj = {"mem":[], "shortcuts":{}} if create else loaddb(db_name, decode_method, password)
    if type(db_obj)==list:
        db_obj = {"mem":db_obj, "shortcuts":{}}
    mem, shortcuts = db_obj["mem"], db_obj["shortcuts"]
    # Main loop
    help_msg = "p: Print Raw, f: Print Formatted, a: Add, m: Modify, d: Delete,\
                wq: Save & Quit, q!: Don't Save, sc: define shortcuts"
    print(help_msg)
    modified = False
    cmd = ""
    while True:
        cmd = input('>')
        if cmd=='p':
            printraw(mem)
        elif cmd=='f':
            keys = input("keys:")
            printfmt(mem, keys, shortcuts)
        elif cmd=='a':
            modified = True
            obj = dict()
            inputkv(obj, shortcuts)
            mem.append(obj)
        elif cmd=='m':
            modified = True
            pos = int(input("index:"))
            inputkv(mem[pos], shortcuts)
        elif cmd=='d':
            modified = True
            pos = int(input("index:"))
            print("deleting %d: "%pos + '\t'.join([k+":"+v for k,v in mem[pos].items()]))
            del mem[pos]
        elif cmd=='q':
            if modified:
                print("Memory has been modified, to save, use 'wq', to exit without save, use 'q!'")
            else:
                break
        elif cmd=='wq':
            savedb(db_name, encode_method, db_obj, password)
            break
        elif cmd=='q!':
            break
        elif cmd=="sc":
            print(shortcuts)
            db_obj["shortcuts"] = inputshortcuts()
        else:
            print("Invalid Command, "+help_msg)

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description='Operate on a list of dicts in the cmdline and\
                                    store them encrypted on the filesystem')
    parser.add_argument('db_name')
    parser.add_argument('decode_method')
    parser.add_argument('encode_method', nargs="?")
    parser.add_argument('--create', action='store_true')
    # db_pass should not be passed in as cmdline args as they are logged and appear in .bash_history
    args = parser.parse_args()
    if args.encode_method==None:
        args.encode_method = args.decode_method
    if args.create and (os.path.isfile(args.db_name+'.store') or os.path.isfile(args.db_name+'.json')):
        print("Cannot overwrite existing db")
        exit(-1)
    password=None
    if args.decode_method=="fernet" or args.encode_method=="fernet":
        password = input("password:").encode('utf-8')
    interact(args.db_name, args.decode_method, args.encode_method, password, args.create)