PyORAM/examples/path_oram_sftp_setup.py

   1 #
   2 # This example demonstrates how to setup an instance of Path ORAM
   3 # locally and then transfer the storage to a server using a paramiko
   4 # SSHClient. After executing this file, path_oram_sftp_test.py can be
   5 # executed to run simple I/O performance tests using different caching
   6 # settings.
   7 #
   8 # In order to run this example, you must provide a host
   9 # (server) address along with valid login credentials
  10 #
  11
  12 import os
  13 import random
  14 import time
  15 import pickle
  16
  17 import pyoram
  18 from pyoram.util.misc import MemorySize, save_private_key
  19 from pyoram.oblivious_storage.tree.path_oram import \
  20     PathORAM
  21
  22 import paramiko
  23 import tqdm
  24
  25 pyoram.config.SHOW_PROGRESS_BAR = True
  26
  27 # Set SSH login credentials here
  28 # (by default, we pull these from the environment
  29 # for testing purposes)
  30 ssh_host = os.environ.get('PYORAM_SSH_TEST_HOST')
  31 ssh_username = os.environ.get('PYORAM_SSH_TEST_USERNAME')
  32 ssh_password = os.environ.get('PYORAM_SSH_TEST_PASSWORD')
  33
  34 # Set the storage location and size
  35 storage_name = "heap.bin"
  36 # 4KB block size
  37 block_size = 4000
  38 # one block per bucket in the
  39 # storage heap of height 8
  40 block_count = 2**(8+1)-1
  41
  42 def main():
  43
  44     print("Storage Name: %s" % (storage_name))
  45     print("Block Count: %s" % (block_count))
  46     print("Block Size: %s" % (MemorySize(block_size)))
  47     print("Total Memory: %s"
  48           % (MemorySize(block_size*block_count)))
  49     print("Actual Storage Required: %s"
  50           % (MemorySize(
  51               PathORAM.compute_storage_size(
  52                   block_size,
  53                   block_count,
  54                   storage_type='mmap'))))
  55     print("")
  56
  57     print("Setting Up Path ORAM Storage Locally")
  58     setup_start = time.time()
  59     with PathORAM.setup(storage_name,
  60                         block_size,
  61                         block_count,
  62                         storage_type='mmap',
  63                         ignore_existing=True) as f:
  64         print("Total Setup Time: %.2f s"
  65               % (time.time()-setup_start))
  66         print("Current Stash Size: %s"
  67               % len(f.stash))
  68         print("Total Data Transmission: %s"
  69               % (MemorySize(f.bytes_sent + f.bytes_received)))
  70         print("")
  71
  72     print("Saving key to file: %s.key"
  73           % (storage_name))
  74     save_private_key(storage_name+".key", f.key)
  75     print("Saving stash to file: %s.stash"
  76           % (storage_name))
  77     with open(storage_name+".stash", 'wb') as fstash:
  78         pickle.dump(f.stash, fstash)
  79     print("Saving position map to file: %s.position"
  80           % (storage_name))
  81     with open(storage_name+".position", 'wb') as fpos:
  82         pickle.dump(f.position_map, fpos)
  83
  84     # Start an SSH client using paramiko
  85     print("Starting SSH Client")
  86     with paramiko.SSHClient() as ssh:
  87         ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
  88         ssh.load_system_host_keys()
  89         ssh.connect(ssh_host,
  90                     username=ssh_username,
  91                     password=ssh_password)
  92
  93         sftp = ssh.open_sftp()
  94
  95         def my_hook(t):
  96             def inner(b, total):
  97                 t.total = total
  98                 t.update(b - inner.last_b)
  99                 inner.last_b = b
 100             inner.last_b = 0
 101             return inner
 102         with tqdm.tqdm(desc="Transferring Storage",
 103                        unit='B',
 104                        unit_scale=True,
 105                        miniters=1) as t:
 106             sftp.put(storage_name,
 107                      storage_name,
 108                      callback=my_hook(t))
 109         sftp.close()
 110
 111     print("Deleting Local Copy of Storage")
 112     os.remove(storage_name)
 113
 114 if __name__ == "__main__":
 115     main()                                             # pragma: no cover