PyORAM/examples/encrypted_heap_test.py

   1 import os
   2 import tempfile
   3
   4 from pyoram.util.virtual_heap import \
   5     SizedVirtualHeap
   6 from pyoram.encrypted_storage.encrypted_heap_storage import \
   7     EncryptedHeapStorage
   8
   9 def main():
  10     #
  11     # get a unique filename in the current directory
  12     #
  13     fid, tmpname = tempfile.mkstemp(dir=os.getcwd())
  14     os.close(fid)
  15     os.remove(tmpname)
  16     print("Storage Name: %s" % (tmpname))
  17
  18     key_size = 32
  19     header_data = b'a message'
  20     heap_base = 3
  21     heap_height = 2
  22     block_size = 8
  23     blocks_per_bucket=4
  24     initialize = lambda i: \
  25         bytes(bytearray([i] * block_size * blocks_per_bucket))
  26     vheap = SizedVirtualHeap(
  27         heap_base,
  28         heap_height,
  29         blocks_per_bucket=blocks_per_bucket)
  30
  31     with EncryptedHeapStorage.setup(
  32             tmpname,
  33             block_size,
  34             heap_height,
  35             key_size=key_size,
  36             header_data=header_data,
  37             heap_base=heap_base,
  38             blocks_per_bucket=blocks_per_bucket,
  39             initialize=initialize) as f:
  40         assert tmpname == f.storage_name
  41         assert f.header_data == header_data
  42         print(f.read_path(vheap.random_bucket()))
  43         key = f.key
  44     assert os.path.exists(tmpname)
  45
  46     with EncryptedHeapStorage(tmpname, key=key) as f:
  47         assert tmpname == f.storage_name
  48         assert f.header_data == header_data
  49         print(f.read_path(vheap.random_bucket()))
  50
  51     #
  52     # cleanup
  53     #
  54     os.remove(tmpname)
  55
  56 if __name__ == "__main__":
  57     main()                                             # pragma: no cover