PyORAM/examples/path_oram_sftp_test.py

   1 #
   2 # This example demonstrates how to access an existing Path ORAM
   3 # storage space through an SSH client using the Secure File Transfer
   4 # Protocol (SFTP). This file should not be executed until the
   5 # path_oram_sftp_setup.py example has been executed. The user is
   6 # encouraged to tweak the settings for 'cached_levels',
   7 # 'concurrency_level', and 'threadpool_size' to observe their effect
   8 # on access latency.
   9 #
  10 # In order to run this example, you must provide a host
  11 # (server) address along with valid login credentials
  12 #
  13
  14 import os
  15 import random
  16 import time
  17 import pickle
  18 import multiprocessing
  19
  20 import pyoram
  21 from pyoram.util.misc import MemorySize, load_private_key
  22 from pyoram.oblivious_storage.tree.path_oram import \
  23     PathORAM
  24
  25 import paramiko
  26 import tqdm
  27
  28 pyoram.config.SHOW_PROGRESS_BAR = True
  29
  30 # Set SSH login credentials here
  31 # (by default, we pull these from the environment
  32 # for testing purposes)
  33 ssh_host = os.environ.get('PYORAM_SSH_TEST_HOST')
  34 ssh_username = os.environ.get('PYORAM_SSH_TEST_USERNAME')
  35 ssh_password = os.environ.get('PYORAM_SSH_TEST_PASSWORD')
  36
  37 # Set the storage location and size
  38 storage_name = "heap.bin"
  39
  40 def main():
  41
  42     print("Loading key from file: %s.key"
  43           % (storage_name))
  44     key = load_private_key(storage_name+".key")
  45     print("Loading stash from file: %s.stash"
  46           % (storage_name))
  47     with open(storage_name+".stash", 'rb') as fstash:
  48         stash = pickle.load(fstash)
  49     print("Loading position map from file: %s.position"
  50           % (storage_name))
  51     with open(storage_name+".position", 'rb') as fpos:
  52         position_map = pickle.load(fpos)
  53
  54     # Start an SSH client using paramiko
  55     print("Starting SSH Client")
  56     with paramiko.SSHClient() as ssh:
  57         ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
  58         ssh.load_system_host_keys()
  59         ssh.connect(ssh_host,
  60                     username=ssh_username,
  61                     password=ssh_password)
  62
  63         with PathORAM(storage_name,
  64                       stash,
  65                       position_map,
  66                       key=key,
  67                       storage_type='sftp',
  68                       cached_levels=6,
  69                       concurrency_level=3,
  70                       threadpool_size=multiprocessing.cpu_count()*2,
  71                       sshclient=ssh) as f:
  72
  73             try:
  74
  75                 test_count = 100
  76                 start_time = time.time()
  77                 for t in tqdm.tqdm(list(range(test_count)),
  78                                    desc="Running I/O Performance Test"):
  79                     f.read_block(random.randint(0,f.block_count-1))
  80                 stop_time = time.time()
  81                 print("Current Stash Size: %s"
  82                       % len(f.stash))
  83                 print("Fetch Block Avg. Latency: %.2f ms"
  84                       % ((stop_time-start_time)/float(test_count)*1000))
  85                 print("")
  86
  87             finally:
  88
  89                 print("Saving stash to file: %s.stash"
  90                       % (storage_name))
  91                 with open(storage_name+".stash", 'wb') as fstash:
  92                     pickle.dump(f.stash, fstash)
  93                 print("Saving position map to file: %s.position"
  94                       % (storage_name))
  95                 with open(storage_name+".position", 'wb') as fpos:
  96                     pickle.dump(f.position_map, fpos)
  97
  98 if __name__ == "__main__":
  99     main()                                             # pragma: no cover