demo_general_database.py

import json, sys, cmd
import codecs
try:
  import cPickle as pickle
except ImportError:
  import pickle

from verifiable_base import VerifiableBase
from verifiable_log import VerifiableLog
from verifiable_map import VerifiableMap, recalc_tree_hash

# Example general purpose verifiable database
# Mutation opertions append to its log
# Its verifiable map then calls the callback (_apply_operation) to change the view.
class VerifiableDatabase(VerifiableBase):
  def __init__(self):
    VerifiableBase.__init__(self, VerifiableLog())

  # Private, call back for the underlying map when new entries are sequenced by the log
  def _apply_operation(self, idx, operation, map):
    op = json.loads(codecs.decode(operation, 'utf-8'))
    if op['operation'] == 'set':
      map.put(str(op['key']), codecs.encode(str(op['value']), 'utf-8'))
    elif op['operation'] == 'delete':
      map.put(str(op['key']), b'')

  # Example database operation
  def set(self, key, value):
    self._log.append(codecs.encode(json.dumps({'operation': 'set', 'key': key, 'value': value}, sort_keys=True), 'utf-8'))

  # Example database operation
  def delete(self, key):
    self._log.append(codecs.encode(json.dumps({'operation': 'delete', 'key': key}, sort_keys=True), 'utf-8'))

  # Return a value for a key and given tree_size (as returned by get_tree_head)
  # Also returns proof
  def get(self, key, tree_size):
    val, proof = VerifiableBase.get(self, str(key), tree_size)
    val = codecs.decode(val, 'utf-8') if len(val) else None
    return val, proof

# Test right val is returned and inclusion proof checks out
def test(db, query, tree_size, exp_val):
  val, proof = db.get(query, tree_size)
  assert val == exp_val
  assert recalc_tree_hash(query, val if val else '', proof) == db.get_tree_head(tree_size)['sha256_root_hash']


class ReplCmd(cmd.Cmd):
  def __init__(self):
    cmd.Cmd.__init__(self)
    self.prompt = '> '
    self.do_new()

  def do_sth(self, arg):
    try:
      if not len(arg.strip()):
        seq = None
      else:
        seq = int(arg)
    except:
      self.help_sth()
      return
    print(self.db.get_tree_head(seq))

  def help_sth(self):
    print('sth <integer> - Updates tree to sequence number and print STH. Leave blank for latest.')

  def do_new(self, arg=''):
    self.db = VerifiableDatabase()

  def help_new(self):
    print('new - creates a new database, called by default upon launch')

  def do_save(self, arg):
    arg = arg.strip()
    if not len(arg):
      self.help_save()
      return
    pickle.dump(self.db, file(arg, 'wb'))

  def help_save(self):
    print('save <path> - save state to a path')

  def do_load(self, arg):
    arg = arg.strip()
    if not len(arg):
      self.help_load()
      return
    self.db = pickle.load(file(arg, 'rb'))

  def help_load(self):
    print('load <path> - load state from path')

  def do_set(self, arg):
    try:
      n, v = arg.split(' ')
    except:
      self.help_set()
      return
    n = n.strip()
    v = v.strip()
    self.db.set(n, v)
    self.do_get(n)

  def help_set(self):
    print('set <key> <value> - set key (string) to the specified value (string)')

  def do_get(self, arg):
    try:
      n, v = arg.split(' ')
      n = n.strip()
      v = v.strip()
    except:
      n = arg.strip()
      v = self.db.get_tree_head(None)['tree_size']
    try:
      v = int(v)
    except:
      self.help_get()
      return
    try:
      val, proof = self.db.get(n, v)
    except ValueError:
      print('Tree size does not exist.')
      return
    print('Value:     ', val)
    print('Proof:     ', proof)
    print('Map hash:  ', self.db.get_tree_head(v)['sha256_root_hash'])
    print('Log hash:  ', self.db.get_tree_head(v)['log_tree_head']['sha256_root_hash'])
    print('Tree size: ', self.db.get_tree_head(v)['tree_size'])

  def help_get(self):
    print('get <key> <integer> - get value as of this sequence number.  Leave blank for latest.')

  def do_del(self, arg):
    n = arg.strip()
    self.db.delete(n)
    self.do_get(n)

  def help_del(self):
    print('del <key> - delete key (string) from database')

  def do_dump(self, arg=''):
    try:
      if not len(arg.strip()):
        seq = None
      else:
        seq = int(arg)
    except:
      self.help_dump()
      return
    print('Tree:')
    self.db.debug_dump(seq)

  def help_dump(self):
    print('dump <integer> - dump the tree as of this sequence number.  Leave blank for latest.')

  def do_log(self, arg=''):
    for i, x in enumerate(self.db.get_log_entries(0, self.db.get_tree_head()['tree_size'] - 1)):
      print(i, x)

  def help_log(self):
    print('log - dump all ops')

  def do_EOF(self, line):
    print('')
    return True


db = VerifiableDatabase()
db.set('foo', 'bar')
db.set('foo', 'baz')

db.delete('foo')
db.set('foo', 'bar')
db.get_tree_head()

test(db, 'foo', 0, None)
test(db, 'foo', 1, 'bar')
test(db, 'foo', 2, 'baz')
test(db, 'foo', 3, None)
test(db, 'foo', 4, 'bar')

if __name__ == '__main__':
  ReplCmd().cmdloop('Type "help" to get started.')