PyORAM/examples/tree_oram_test.py

   1 import os
   2 import struct
   3 import random
   4
   5 from pyoram.util.virtual_heap import \
   6     SizedVirtualHeap
   7 from pyoram.encrypted_storage.encrypted_heap_storage import \
   8     EncryptedHeapStorage
   9 from pyoram.oblivious_storage.tree.tree_oram_helper import \
  10     TreeORAMStorageManagerPointerAddressing
  11
  12 def main():
  13     storage_name = "heap.bin"
  14     print("Storage Name: %s" % (storage_name))
  15
  16     key_size = 32
  17     heap_base = 2
  18     heap_height = 2
  19     block_size = struct.calcsize("!?LL")
  20     blocks_per_bucket = 2
  21     vheap = SizedVirtualHeap(
  22         heap_base,
  23         heap_height,
  24         blocks_per_bucket=blocks_per_bucket)
  25
  26     print("Block Size: %s" % (block_size))
  27     print("Blocks Per Bucket: %s" % (blocks_per_bucket))
  28
  29     position_map = {}
  30     def initialize(i):
  31         bucket = bytes()
  32         for j in range(blocks_per_bucket):
  33             if (i*j) % 3:
  34                 bucket += struct.pack(
  35                     "!?LL", False, 0, 0)
  36             else:
  37                 x = vheap.Node(i)
  38                 while not vheap.is_nil_node(x):
  39                     x = x.child_node(random.randint(0, heap_base-1))
  40                 x = x.parent_node()
  41                 bucket += struct.pack(
  42                     "!?LL", True, initialize.id_, x.bucket)
  43                 position_map[initialize.id_] = x.bucket
  44                 initialize.id_ += 1
  45         return bucket
  46     initialize.id_ = 1
  47
  48     with EncryptedHeapStorage.setup(
  49             storage_name,
  50             block_size,
  51             heap_height,
  52             heap_base=heap_base,
  53             key_size=key_size,
  54             blocks_per_bucket=blocks_per_bucket,
  55             initialize=initialize,
  56             ignore_existing=True) as f:
  57         assert storage_name == f.storage_name
  58         stash = {}
  59         oram = TreeORAMStorageManagerPointerAddressing(f, stash)
  60
  61         b = vheap.random_bucket()
  62         oram.load_path(b)
  63         print("")
  64         print(repr(vheap.Node(oram.path_stop_bucket)))
  65         print(oram.path_block_ids)
  66         print(oram.path_block_eviction_levels)
  67
  68         oram.push_down_path()
  69         print("")
  70         print(repr(vheap.Node(oram.path_stop_bucket)))
  71         print(oram.path_block_ids)
  72         print(oram.path_block_eviction_levels)
  73         print(oram.path_block_reordering)
  74
  75         oram.evict_path()
  76         oram.load_path(b)
  77         print("")
  78         print(repr(vheap.Node(oram.path_stop_bucket)))
  79         print(oram.path_block_ids)
  80         print(oram.path_block_eviction_levels)
  81
  82         oram.push_down_path()
  83         print("")
  84         print(repr(vheap.Node(oram.path_stop_bucket)))
  85         print(oram.path_block_ids)
  86         print(oram.path_block_eviction_levels)
  87         print(oram.path_block_reordering)
  88         assert all(x is None for x in oram.path_block_reordering)
  89
  90     os.remove(storage_name)
  91
  92 if __name__ == "__main__":
  93     main()                                             # pragma: no cover