├── .appveyor.yml ├── .gitignore ├── .travis.yml ├── .travis_before_install_jython.sh ├── CHANGELOG.rst ├── LICENSE.md ├── MANIFEST.in ├── Makefile ├── README.rst ├── examples ├── aesctr_performance.py ├── encrypted_heap_test.py ├── encrypted_storage_file.py ├── encrypted_storage_mmap.py ├── encrypted_storage_ram.py ├── encrypted_storage_s3.py ├── encrypted_storage_sftp.py ├── path_oram_file.py ├── path_oram_mmap.py ├── path_oram_ram.py ├── path_oram_s3.py ├── path_oram_sftp.py ├── path_oram_sftp_setup.py ├── path_oram_sftp_test.py └── tree_oram_test.py ├── requirements.txt ├── setup.cfg ├── setup.py └── src ├── _cffi_src ├── __init__.py └── virtual_heap_helper_build.py └── pyoram ├── __about__.py ├── __init__.py ├── crypto ├── __init__.py └── aes.py ├── encrypted_storage ├── __init__.py ├── encrypted_block_storage.py ├── encrypted_heap_storage.py └── top_cached_encrypted_heap_storage.py ├── oblivious_storage ├── __init__.py └── tree │ ├── __init__.py │ ├── path_oram.py │ └── tree_oram_helper.py ├── storage ├── __init__.py ├── block_storage.py ├── block_storage_file.py ├── block_storage_mmap.py ├── block_storage_ram.py ├── block_storage_s3.py ├── block_storage_sftp.py ├── boto3_s3_wrapper.py └── heap_storage.py ├── tests ├── __init__.py ├── baselines │ ├── exists.empty │ ├── k200_h0_b1.dot │ ├── k200_h0_b1_data.dot │ ├── k2_h3_b1.dot │ ├── k2_h3_b1_data.dot │ ├── k2_h3_b2.dot │ ├── k2_h3_b2_data.dot │ ├── k3_h3_b1.dot │ ├── k3_h3_b1_data.dot │ ├── k3_h3_b2.dot │ └── k3_h3_b2_data.dot ├── test_aes.py ├── test_block_storage.py ├── test_encrypted_block_storage.py ├── test_encrypted_heap_storage.py ├── test_examples.py ├── test_heap_storage.py ├── test_misc.py ├── test_package.py ├── test_path_oram.py ├── test_top_cached_encrypted_heap_storage.py └── test_virtual_heap.py └── util ├── __init__.py ├── misc.py └── virtual_heap.py /.appveyor.yml: -------------------------------------------------------------------------------- 1 | environment: 2 | 3 | matrix: 4 | - PYTHON: "C:\\Python27" 5 | - PYTHON: "C:\\Python34" 6 | - PYTHON: "C:\\Python35" 7 | - PYTHON: "C:\\Python36" 8 | - PYTHON: "C:\\Python27-x64" 9 | - PYTHON: "C:\\Python35-x64" 10 | - PYTHON: "C:\\Python36-x64" 11 | 12 | install: 13 | - "SET PATH=%PYTHON%;%PYTHON%\\Scripts;%PATH%" 14 | - "%PYTHON%\\python.exe -m pip install -v -U pip setuptools" 15 | - "%PYTHON%\\python.exe -m pip install -v -e ." 16 | - "%PYTHON%\\python.exe -m pip install -U unittest2 nose2 cov-core codecov coverage" 17 | 18 | build: none 19 | 20 | test_script: 21 | - "%PYTHON%\\python.exe -m nose2 -v --log-capture --with-coverage --coverage src --coverage examples -s src" 22 | 23 | on_success: 24 | - "%PYTHON%\\Scripts\\codecov.exe" 25 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | # Emacs 2 | *~ 3 | **/\#* 4 | 5 | # Python 6 | *.py[cod] 7 | 8 | # C extensions 9 | *.so 10 | 11 | # Setuptools distribution folder. 12 | /dist/ 13 | /build/ 14 | 15 | # Python egg metadata, regenerated from source files by setuptools. 16 | *.egg-info 17 | *.egg 18 | *.eggs 19 | 20 | # nose 21 | .coverage 22 | coverage.xml 23 | -------------------------------------------------------------------------------- /.travis.yml: -------------------------------------------------------------------------------- 1 | # travis CI config 2 | language: python 3 | matrix: 4 | include: 5 | - python: 2.7 6 | - python: 3.4 7 | - python: 3.5 8 | - python: 3.6 9 | - python: 3.7 10 | - python: 3.8 11 | dist: xenial 12 | - python: pypy 13 | - python: pypy3 14 | cache: false 15 | before_install: 16 | - sudo apt-get update -q 17 | - sudo apt-get install graphviz -y 18 | - python -m pip install -v -U pip setuptools virtualenv wheel 19 | install: 20 | - python -m pip install -v -e . 21 | - python -m pip install -U unittest2 nose2 cov-core codecov coverage 22 | script: python -m nose2 -v --log-capture --with-coverage --coverage src --coverage examples -s src 23 | after_success: 24 | - codecov 25 | branches: 26 | only: 27 | - master 28 | -------------------------------------------------------------------------------- /.travis_before_install_jython.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | set -e 3 | 4 | pip install jip 5 | jip install $JYTHON 6 | NON_GROUP_ID=${JYTHON#*:} 7 | _JYTHON_BASENAME=${NON_GROUP_ID/:/-} 8 | OLD_VIRTUAL_ENV=${VIRTUAL_ENV:=.} 9 | java -jar $OLD_VIRTUAL_ENV/javalib/${_JYTHON_BASENAME}.jar -s -d $HOME/jython 10 | $HOME/jython/bin/jython -c "import sys; print(sys.version_info)" 11 | virtualenv --version 12 | virtualenv -p $HOME/jython/bin/jython $HOME/myvirtualenv 13 | -------------------------------------------------------------------------------- /CHANGELOG.rst: -------------------------------------------------------------------------------- 1 | Changelog 2 | ========= 3 | 4 | 0.3.0 - `master`_ 5 | ~~~~~~~~~~~~~~~~~ 6 | 7 | 0.2.1 - 2018-01-04 8 | ~~~~~~~~~~~~~~~~~~ 9 | 10 | * fixes to support latest version of cryptography 11 | 12 | 0.2.0 - 2016-10-18 13 | ~~~~~~~~~~~~~~~~~~ 14 | 15 | * using chunking to speed up yield_blocks for SFTP 16 | * speed up clearing entries in S3 interface by chunking delete requests 17 | * adding helper property to access heap storage on path oram 18 | * use a mmap to store the top-cached heap buckets 19 | * replace the show_status_bar keywords by a global config item 20 | * express status bar units as a memory transfer rate during setup 21 | * tweaks to Path ORAM to make it easier to generalize to other schemes 22 | * changing suffix of S3 index file from txt to bin 23 | * updates to readme 24 | 25 | 0.1.2 - 2016-05-15 26 | ~~~~~~~~~~~~~~~~~~ 27 | 28 | * Initial release. 29 | 30 | .. _`master`: https://github.com/ghackebeil/PyORAM 31 | -------------------------------------------------------------------------------- /LICENSE.md: -------------------------------------------------------------------------------- 1 | The MIT License (MIT) 2 | 3 | Copyright (c) 2016 Gabriel Hackebeil 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /MANIFEST.in: -------------------------------------------------------------------------------- 1 | include README.rst 2 | include CHANGELOG.rst 3 | include LICENSE.md 4 | 5 | recursive-include src/_cffi_src *.py 6 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | all: 2 | 3 | .PHONY: clean 4 | clean: 5 | find examples -name "*.pyc" | xargs rm 6 | find src -name "*.pyc" | xargs rm 7 | find . -depth 1 -name "*.pyc" | xargs rm 8 | 9 | find examples -name "*.pyo" | xargs rm 10 | find src -name "*.pyo" | xargs rm 11 | find . -depth 1 -name "*.pyo" | xargs rm 12 | 13 | find examples -name "__pycache__" | xargs rm -r 14 | find src -name "__pycache__" | xargs rm -r 15 | find . -depth 1 -name "__pycache__" | xargs rm -r 16 | 17 | find examples -name "*~" | xargs rm 18 | find src -name "*~" | xargs rm 19 | find . -depth 1 -name "*~" | xargs rm 20 | 21 | find src -name "*.so" | xargs rm 22 | -------------------------------------------------------------------------------- /README.rst: -------------------------------------------------------------------------------- 1 | PyORAM 2 | ====== 3 | 4 | .. image:: https://travis-ci.org/ghackebeil/PyORAM.svg?branch=master 5 | :target: https://travis-ci.org/ghackebeil/PyORAM 6 | 7 | .. image:: https://ci.appveyor.com/api/projects/status/1tpnf7fr0qthrwxx/branch/master?svg=true 8 | :target: https://ci.appveyor.com/project/ghackebeil/PyORAM?branch=master 9 | 10 | .. image:: https://codecov.io/github/ghackebeil/PyORAM/coverage.svg?branch=master 11 | :target: https://codecov.io/github/ghackebeil/PyORAM?branch=master 12 | 13 | .. image:: https://img.shields.io/pypi/v/PyORAM.svg 14 | :target: https://pypi.python.org/pypi/PyORAM/ 15 | 16 | Python-based Oblivious RAM (PyORAM) is a collection of 17 | Oblivious RAM algorithms implemented in Python. This package 18 | serves to enable rapid prototyping and testing of new ORAM 19 | algorithms and ORAM-based applications tailored for the 20 | cloud-storage setting. PyORAM is written to support as many 21 | Python versions as possible, including Python 2.7+, Python 22 | 3.4+, and PyPy 2.6+. 23 | 24 | This software is copyright (c) by Gabriel A. Hackebeil (gabe.hackebeil@gmail.com). 25 | 26 | This software is released under the MIT software license. 27 | This license, including disclaimer, is available in the 'LICENSE' file. 28 | 29 | This work was funded by the Privacy Enhancing Technologies 30 | project under the guidance of Professor `Attila Yavuz 31 | `_ at Oregon State 32 | University. 33 | 34 | Why Python? 35 | ----------- 36 | 37 | This project is meant for research. It is provided mainly as 38 | a tool for other researchers studying the applicability of 39 | ORAM to the cloud-storage setting. In such a setting, we 40 | observe that network latency far outweighs any overhead 41 | introduced from switching to an interpreted language such as 42 | Python (as opposed to C++ or Java). Thus, our hope is that 43 | by providing a Python-based library of ORAM tools, we will 44 | enable researchers to spend more time prototyping new and 45 | interesting ORAM applications and less time fighting with a 46 | compiler or chasing down segmentation faults. 47 | 48 | Installation 49 | ------------ 50 | 51 | To install the latest release of PyORAM, simply execute:: 52 | 53 | $ pip install PyORAM 54 | 55 | To install the trunk version of PyORAM, first clone the repository:: 56 | 57 | $ git clone https://github.com/ghackebeil/PyORAM.git 58 | 59 | Next, enter the directory where PyORAM has been cloned and run setup:: 60 | 61 | $ python setup.py install 62 | 63 | If you are a developer, you should instead install using:: 64 | 65 | $ pip install -e . 66 | $ pip install nose2 unittest 67 | 68 | Installation Tips 69 | ----------------- 70 | 71 | * OS X users are recommended to work with the `homebrew 72 | `_ version of Python2 or Python3. If you 73 | must use the default system Python, then the best thing to 74 | do is create a virtual environment and install PyORAM into 75 | that. The process of creating a virtual environment that is 76 | stored in the PyORAM directory would look something like:: 77 | 78 | $ sudo pip install virtualenv 79 | $ cd 80 | $ virtualenv local_python2.7 81 | 82 | If you had already attempted to install PyORAM into the 83 | system Python and encountered errors, it may be necessary 84 | to delete the directories :code:`build` and :code:`dist` 85 | from the current directory using the command:: 86 | 87 | $ sudo rm -rf build dist 88 | 89 | Once this virtual environment has been successfully 90 | created, you can *activate* it using the command:: 91 | 92 | $ . local_python2.7/bin/activate 93 | 94 | Then, proceed with the normal installation steps to 95 | install PyORAM into this environment. Note that you must 96 | *activate* this environment each time you open a new 97 | terminal if PyORAM is installed in this way. Also, note 98 | that use of the :code:`sudo` command is no longer 99 | necessary (and should be avoided) once a virtual 100 | environment is activated in the current shell. 101 | 102 | * If you have trouble installing the cryptography package 103 | on OS X with PyPy: `stackoverflow `_. 104 | 105 | * If you encounter the dreaded "unable to find 106 | vcvarsall.bat" error when installing packages with C 107 | extensions through pip on Windows: `blog post `_. 108 | 109 | Tools Available (So Far) 110 | ------------------------ 111 | 112 | Encrypted block storage 113 | ~~~~~~~~~~~~~~~~~~~~~~~ 114 | 115 | * The basic building block for any ORAM implementation. 116 | 117 | * Available storage interfaces include: 118 | 119 | - local storage using a file, a memory-mapped file, or RAM 120 | 121 | + Dropbox 122 | 123 | - cloud storage using SFTP (requires SSH access to a server) 124 | 125 | + Amazon EC2 126 | 127 | + Microsoft Azure 128 | 129 | + Google Cloud Platform 130 | 131 | - cloud storage using Amazon Simple Storage Service (S3) 132 | 133 | * See Examples: 134 | 135 | - examples/encrypted_storage_ram.py 136 | 137 | - examples/encrypted_storage_mmap.py 138 | 139 | - examples/encrypted_storage_file.py 140 | 141 | - examples/encrypted_storage_sftp.py 142 | 143 | - examples/encrypted_storage_s3.py 144 | 145 | Path ORAM 146 | ~~~~~~~~~ 147 | 148 | * Reference: `Stefanov et al. `_ 149 | 150 | * Generalized to work over k-kary storage heaps. Default 151 | settings use a binary storage heap and bucket size 152 | parameter set to 4. Using a k-ary storage heap can reduce 153 | the access cost; however, stash size behavior has not been 154 | formally analyzed in this setting. 155 | 156 | * Tree-Top caching can be used to reduce data transmission 157 | per access as well as reduce access latency by exploiting 158 | parallelism across independent sub-heaps below the last 159 | cached heap level. 160 | 161 | * See Examples: 162 | 163 | - examples/path_oram_ram.py 164 | 165 | - examples/path_oram_mmap.py 166 | 167 | - examples/path_oram_file.py 168 | 169 | - examples/path_oram_sftp.py 170 | 171 | - examples/path_oram_s3.py 172 | 173 | Performance Tips 174 | ---------------- 175 | 176 | Setup Storage Locally 177 | ~~~~~~~~~~~~~~~~~~~~~ 178 | 179 | Storage schemes such as BlockStorageFile ("file"), BlockStorageMMap 180 | ("mmap"), BlockStorageRAM ("ram"), and BlockStorageSFTP ("sftp") all 181 | employ the same underlying storage format. Thus, an oblivious storage 182 | scheme can be initialized locally and then transferred to an external 183 | storage location and accessed via BlockStorageSFTP using SSH login 184 | credentials. See the following pair of files for an example of this: 185 | 186 | * examples/path_oram_sftp_setup.py 187 | 188 | * examples/path_oram_sftp_test.py 189 | 190 | BlockStorageS3 ("s3") employs a different format whereby the 191 | underlying blocks are stored in separate "file" objects. 192 | This design is due to the fact that the Amazon S3 API does 193 | not allow modifications to a specific byte range within a 194 | file, but instead requires that the entire modified file 195 | object be re-uploaded. Thus, any efficient block storage 196 | scheme must use separate "file" objects for each block. 197 | 198 | Tree-Top Caching 199 | ~~~~~~~~~~~~~~~~ 200 | 201 | For schemes that employ a storage heap (such as Path ORAM), 202 | tree-top caching provides the ability to parallelize I/O 203 | operations across the independent sub-heaps below the last 204 | cached heap level. The default behavior of this 205 | implementation of Path ORAM, for instance, caches the top 206 | three levels of the storage heap in RAM, which creates eight 207 | independent sub-heaps across which write operations can be 208 | asynchronous. 209 | 210 | If the underlying storage is being accessed through SFTP, the 211 | tree-top cached storage heap will attempt to open an 212 | independent SFTP session for each sub-heap using the same 213 | SSH connection. Typically, the maximum number of allowable 214 | sessions associated with a single SSH connection is limited 215 | by the SSH server. For instance, the default maximum number 216 | of sessions allowed by a server using OpenSSH is 10. Thus, 217 | increasing the number of cached levels beyond 3 when using 218 | a binary storage heap will attempt to generate 16 or more SFTP 219 | sessions and result in an error such as:: 220 | 221 | paramiko.ssh_exception.ChannelException: (1, 'Administratively prohibited') 222 | 223 | There are two options for avoiding this error: 224 | 225 | 1. If you have administrative privileges on the server, you 226 | can increase the maximum number of allowed sessions for a 227 | single SSH connection. For example, to set the maximum 228 | allowed sessions to 128 on a server using OpenSSH, one 229 | would set:: 230 | 231 | MaxSessions 128 232 | 233 | in :code:`/etc/ssh/sshd_config`, and then run the 234 | command :code:`sudo service ssh restart`. 235 | 236 | 2. You can limit the number of concurrent devices that will 237 | be created by setting the concurrency level to something 238 | below the last cached level using the 239 | :code:`concurrency_level` keyword. For example, the 240 | settings :code:`cached_levels=5` and 241 | :code:`concurrency_level=0` would cache the top 5 levels 242 | of the storage heap locally, but all external I/O 243 | operations would take place through a single storage 244 | device (e.g., using 1 SFTP session). 245 | -------------------------------------------------------------------------------- /examples/aesctr_performance.py: -------------------------------------------------------------------------------- 1 | import time 2 | import base64 3 | 4 | from pyoram.crypto.aes import AES 5 | 6 | def runtest(label, enc_func, dec_func): 7 | print("") 8 | print("$"*20) 9 | print("{0:^20}".format(label)) 10 | print("$"*20) 11 | for keysize in AES.key_sizes[:3]: 12 | print("") 13 | print("@@@@@@@@@@@@@@@@@@@@") 14 | print(" Key Size: %s bytes" % (keysize)) 15 | print("@@@@@@@@@@@@@@@@@@@@") 16 | print("\nTest Bulk") 17 | # 18 | # generate a key 19 | # 20 | key = AES.KeyGen(keysize) 21 | print("Key: %s" % (base64.b64encode(key))) 22 | 23 | # 24 | # generate some plaintext 25 | # 26 | nblocks = 1000000 27 | plaintext_numbytes = AES.block_size * nblocks 28 | print("Plaintext Size: %s MB" 29 | % (plaintext_numbytes * 1.0e-6)) 30 | # all zeros 31 | plaintext = bytes(bytearray(plaintext_numbytes)) 32 | 33 | # 34 | # time encryption 35 | # 36 | start_time = time.time() 37 | ciphertext = enc_func(key, plaintext) 38 | stop_time = time.time() 39 | print("Encryption Time: %.3fs (%.3f MB/s)" 40 | % (stop_time-start_time, 41 | (plaintext_numbytes * 1.0e-6) / (stop_time-start_time))) 42 | 43 | # 44 | # time decryption 45 | # 46 | start_time = time.time() 47 | plaintext_decrypted = dec_func(key, ciphertext) 48 | stop_time = time.time() 49 | print("Decryption Time: %.3fs (%.3f MB/s)" 50 | % (stop_time-start_time, 51 | (plaintext_numbytes * 1.0e-6) / (stop_time-start_time))) 52 | 53 | assert plaintext_decrypted == plaintext 54 | assert ciphertext != plaintext 55 | # IND-CPA 56 | assert enc_func(key, plaintext) != ciphertext 57 | # make sure the only difference is not in the IV 58 | assert enc_func(key, plaintext)[AES.block_size:] \ 59 | != ciphertext[AES.block_size:] 60 | if enc_func is AES.CTREnc: 61 | assert len(plaintext) == \ 62 | len(ciphertext) - AES.block_size 63 | else: 64 | assert enc_func is AES.GCMEnc 65 | assert len(plaintext) == \ 66 | len(ciphertext) - 2*AES.block_size 67 | 68 | del plaintext 69 | del plaintext_decrypted 70 | del ciphertext 71 | 72 | print("\nTest Chunks") 73 | # 74 | # generate a key 75 | # 76 | key = AES.KeyGen(keysize) 77 | print("Key: %s" % (base64.b64encode(key))) 78 | 79 | # 80 | # generate some plaintext 81 | # 82 | nblocks = 1000 83 | blocksize = 16000 84 | total_bytes = blocksize * nblocks 85 | print("Block Size: %s KB" % (blocksize * 1.0e-3)) 86 | print("Block Count: %s" % (nblocks)) 87 | print("Total: %s MB" % (total_bytes * 1.0e-6)) 88 | plaintext_blocks = [bytes(bytearray(blocksize)) 89 | for i in range(nblocks)] 90 | 91 | # 92 | # time encryption 93 | # 94 | start_time = time.time() 95 | ciphertext_blocks = [enc_func(key, b) 96 | for b in plaintext_blocks] 97 | stop_time = time.time() 98 | print("Encryption Time: %.3fs (%.3f MB/s)" 99 | % (stop_time-start_time, 100 | (total_bytes * 1.0e-6) / (stop_time-start_time))) 101 | 102 | # 103 | # time decryption 104 | # 105 | start_time = time.time() 106 | plaintext_decrypted_blocks = [dec_func(key, c) 107 | for c in ciphertext_blocks] 108 | stop_time = time.time() 109 | print("Decryption Time: %.3fs (%.3f MB/s)" 110 | % (stop_time-start_time, 111 | (total_bytes * 1.0e-6) / (stop_time-start_time))) 112 | 113 | def main(): 114 | runtest("AES - CTR Mode", AES.CTREnc, AES.CTRDec) 115 | runtest("AES - GCM Mode", AES.GCMEnc, AES.GCMDec) 116 | 117 | if __name__ == "__main__": 118 | main() # pragma: no cover 119 | -------------------------------------------------------------------------------- /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 58 | -------------------------------------------------------------------------------- /examples/encrypted_storage_file.py: -------------------------------------------------------------------------------- 1 | # 2 | # This example measures the performance of encrypted 3 | # storage access through a local file. 4 | # 5 | 6 | import os 7 | import random 8 | import time 9 | 10 | import pyoram 11 | from pyoram.util.misc import MemorySize 12 | from pyoram.encrypted_storage.encrypted_block_storage import \ 13 | EncryptedBlockStorage 14 | 15 | import tqdm 16 | 17 | pyoram.config.SHOW_PROGRESS_BAR = True 18 | 19 | # Set the storage location and size 20 | storage_name = "heap.bin" 21 | # 4KB block size 22 | block_size = 4000 23 | # one block per bucket in the 24 | # storage heap of height 8 25 | block_count = 2**(8+1)-1 26 | 27 | def main(): 28 | 29 | print("Storage Name: %s" % (storage_name)) 30 | print("Block Count: %s" % (block_count)) 31 | print("Block Size: %s" % (MemorySize(block_size))) 32 | print("Total Memory: %s" 33 | % (MemorySize(block_size*block_count))) 34 | print("Actual Storage Required: %s" 35 | % (MemorySize( 36 | EncryptedBlockStorage.compute_storage_size( 37 | block_size, 38 | block_count, 39 | storage_type='file')))) 40 | print("") 41 | 42 | print("Setting Up Encrypted Block Storage") 43 | setup_start = time.time() 44 | with EncryptedBlockStorage.setup(storage_name, 45 | block_size, 46 | block_count, 47 | storage_type='file', 48 | ignore_existing=True) as f: 49 | print("Total Setup Time: %2.f s" 50 | % (time.time()-setup_start)) 51 | print("Total Data Transmission: %s" 52 | % (MemorySize(f.bytes_sent + f.bytes_received))) 53 | print("") 54 | 55 | # We close the device and reopen it after 56 | # setup to reset the bytes sent and bytes 57 | # received stats. 58 | with EncryptedBlockStorage(storage_name, 59 | key=f.key, 60 | storage_type='file') as f: 61 | 62 | test_count = 1000 63 | start_time = time.time() 64 | for t in tqdm.tqdm(list(range(test_count)), 65 | desc="Running I/O Performance Test"): 66 | f.read_block(random.randint(0,f.block_count-1)) 67 | stop_time = time.time() 68 | print("Access Block Avg. Data Transmitted: %s (%.3fx)" 69 | % (MemorySize((f.bytes_sent + f.bytes_received)/float(test_count)), 70 | (f.bytes_sent + f.bytes_received)/float(test_count)/float(block_size))) 71 | print("Access Block Avg. Latency: %.2f ms" 72 | % ((stop_time-start_time)/float(test_count)*1000)) 73 | print("") 74 | 75 | # cleanup because this is a test example 76 | os.remove(storage_name) 77 | 78 | if __name__ == "__main__": 79 | main() # pragma: no cover 80 | -------------------------------------------------------------------------------- /examples/encrypted_storage_mmap.py: -------------------------------------------------------------------------------- 1 | # 2 | # This example measures the performance of encrypted 3 | # storage access through a local memory-mapped file. 4 | # 5 | 6 | import os 7 | import random 8 | import time 9 | 10 | import pyoram 11 | from pyoram.util.misc import MemorySize 12 | from pyoram.encrypted_storage.encrypted_block_storage import \ 13 | EncryptedBlockStorage 14 | 15 | import tqdm 16 | 17 | pyoram.config.SHOW_PROGRESS_BAR = True 18 | 19 | # Set the storage location and size 20 | storage_name = "heap.bin" 21 | # 4KB block size 22 | block_size = 4000 23 | # one block per bucket in the 24 | # storage heap of height 8 25 | block_count = 2**(8+1)-1 26 | 27 | def main(): 28 | 29 | print("Storage Name: %s" % (storage_name)) 30 | print("Block Count: %s" % (block_count)) 31 | print("Block Size: %s" % (MemorySize(block_size))) 32 | print("Total Memory: %s" 33 | % (MemorySize(block_size*block_count))) 34 | print("Actual Storage Required: %s" 35 | % (MemorySize( 36 | EncryptedBlockStorage.compute_storage_size( 37 | block_size, 38 | block_count, 39 | storage_type='mmap')))) 40 | print("") 41 | 42 | print("Setting Up Encrypted Block Storage") 43 | setup_start = time.time() 44 | with EncryptedBlockStorage.setup(storage_name, 45 | block_size, 46 | block_count, 47 | storage_type='mmap', 48 | ignore_existing=True) as f: 49 | print("Total Setup Time: %2.f s" 50 | % (time.time()-setup_start)) 51 | print("Total Data Transmission: %s" 52 | % (MemorySize(f.bytes_sent + f.bytes_received))) 53 | print("") 54 | 55 | # We close the device and reopen it after 56 | # setup to reset the bytes sent and bytes 57 | # received stats. 58 | with EncryptedBlockStorage(storage_name, 59 | key=f.key, 60 | storage_type='mmap') as f: 61 | 62 | test_count = 1000 63 | start_time = time.time() 64 | for t in tqdm.tqdm(list(range(test_count)), 65 | desc="Running I/O Performance Test"): 66 | f.read_block(random.randint(0,f.block_count-1)) 67 | stop_time = time.time() 68 | print("Access Block Avg. Data Transmitted: %s (%.3fx)" 69 | % (MemorySize((f.bytes_sent + f.bytes_received)/float(test_count)), 70 | (f.bytes_sent + f.bytes_received)/float(test_count)/float(block_size))) 71 | print("Access Block Avg. Latency: %.2f ms" 72 | % ((stop_time-start_time)/float(test_count)*1000)) 73 | print("") 74 | 75 | # cleanup because this is a test example 76 | os.remove(storage_name) 77 | 78 | if __name__ == "__main__": 79 | main() # pragma: no cover 80 | -------------------------------------------------------------------------------- /examples/encrypted_storage_ram.py: -------------------------------------------------------------------------------- 1 | # 2 | # This example measures the performance of encrypted 3 | # storage access through RAM. 4 | # 5 | 6 | import os 7 | import random 8 | import time 9 | 10 | import pyoram 11 | from pyoram.util.misc import MemorySize 12 | from pyoram.encrypted_storage.encrypted_block_storage import \ 13 | EncryptedBlockStorage 14 | from pyoram.storage.block_storage_ram import \ 15 | BlockStorageRAM 16 | 17 | import tqdm 18 | 19 | pyoram.config.SHOW_PROGRESS_BAR = True 20 | 21 | # Set the storage location and size 22 | storage_name = "heap.bin" 23 | # 4KB block size 24 | block_size = 4000 25 | # one block per bucket in the 26 | # storage heap of height 8 27 | block_count = 2**(8+1)-1 28 | 29 | def main(): 30 | 31 | print("Storage Name: %s" % (storage_name)) 32 | print("Block Count: %s" % (block_count)) 33 | print("Block Size: %s" % (MemorySize(block_size))) 34 | print("Total Memory: %s" 35 | % (MemorySize(block_size*block_count))) 36 | print("Actual Storage Required: %s" 37 | % (MemorySize( 38 | EncryptedBlockStorage.compute_storage_size( 39 | block_size, 40 | block_count, 41 | storage_type='ram')))) 42 | print("") 43 | 44 | print("Setting Up Encrypted Block Storage") 45 | setup_start = time.time() 46 | with EncryptedBlockStorage.setup(storage_name, # RAM storage ignores this argument 47 | block_size, 48 | block_count, 49 | storage_type='ram', 50 | ignore_existing=True) as f: 51 | print("Total Setup Time: %2.f s" 52 | % (time.time()-setup_start)) 53 | print("Total Data Transmission: %s" 54 | % (MemorySize(f.bytes_sent + f.bytes_received))) 55 | print("") 56 | 57 | # This must be done after closing the file to ensure the lock flag 58 | # is set to False in the saved data. The tofile method only exists 59 | # on BlockStorageRAM 60 | f.raw_storage.tofile(storage_name) 61 | 62 | # We close the device and reopen it after 63 | # setup to reset the bytes sent and bytes 64 | # received stats. 65 | with EncryptedBlockStorage(BlockStorageRAM.fromfile(storage_name), 66 | key=f.key) as f: 67 | 68 | test_count = 1000 69 | start_time = time.time() 70 | for t in tqdm.tqdm(list(range(test_count)), 71 | desc="Running I/O Performance Test"): 72 | f.read_block(random.randint(0,f.block_count-1)) 73 | stop_time = time.time() 74 | print("Access Block Avg. Data Transmitted: %s (%.3fx)" 75 | % (MemorySize((f.bytes_sent + f.bytes_received)/float(test_count)), 76 | (f.bytes_sent + f.bytes_received)/float(test_count)/float(block_size))) 77 | print("Access Block Avg. Latency: %.2f ms" 78 | % ((stop_time-start_time)/float(test_count)*1000)) 79 | print("") 80 | 81 | # cleanup because this is a test example 82 | os.remove(storage_name) 83 | 84 | if __name__ == "__main__": 85 | main() # pragma: no cover 86 | -------------------------------------------------------------------------------- /examples/encrypted_storage_s3.py: -------------------------------------------------------------------------------- 1 | # 2 | # This example measures the performance of encrypted 3 | # storage access through Amazon Simple Storage Service 4 | # (S3). 5 | # 6 | # In order to run this example, you must provide a valid 7 | # S3 bucket name and have the following variables defined 8 | # in your current environment: 9 | # - AWS_ACCESS_KEY_ID 10 | # - AWS_SECRET_ACCESS_KEY 11 | # - AWS_DEFAULT_REGION 12 | # These can also be set using keywords. 13 | # 14 | 15 | import os 16 | import random 17 | import time 18 | 19 | import pyoram 20 | from pyoram.util.misc import MemorySize 21 | from pyoram.encrypted_storage.encrypted_block_storage import \ 22 | EncryptedBlockStorage 23 | 24 | import tqdm 25 | 26 | pyoram.SHOW_PROGRESS_BAR = True 27 | 28 | # Set S3 bucket name here 29 | # (by default, we pull this from the environment 30 | # for testing purposes) 31 | bucket_name = os.environ.get('PYORAM_AWS_TEST_BUCKET') 32 | 33 | # Set the storage location and size 34 | storage_name = "heap.bin" 35 | # 4KB block size 36 | block_size = 4000 37 | # one block per bucket in the 38 | # storage heap of height 8 39 | block_count = 2**(8+1)-1 40 | 41 | def main(): 42 | 43 | print("Storage Name: %s" % (storage_name)) 44 | print("Block Count: %s" % (block_count)) 45 | print("Block Size: %s" % (MemorySize(block_size))) 46 | print("Total Memory: %s" 47 | % (MemorySize(block_size*block_count))) 48 | print("Actual Storage Required: %s" 49 | % (MemorySize( 50 | EncryptedBlockStorage.compute_storage_size( 51 | block_size, 52 | block_count, 53 | storage_type='s3')))) 54 | print("") 55 | 56 | print("Setting Up Encrypted Block Storage") 57 | setup_start = time.time() 58 | with EncryptedBlockStorage.setup(storage_name, 59 | block_size, 60 | block_count, 61 | storage_type='s3', 62 | bucket_name=bucket_name, 63 | ignore_existing=True) as f: 64 | print("Total Setup Time: %2.f s" 65 | % (time.time()-setup_start)) 66 | print("Total Data Transmission: %s" 67 | % (MemorySize(f.bytes_sent + f.bytes_received))) 68 | print("") 69 | 70 | # We close the device and reopen it after 71 | # setup to reset the bytes sent and bytes 72 | # received stats. 73 | with EncryptedBlockStorage(storage_name, 74 | key=f.key, 75 | storage_type='s3', 76 | bucket_name=bucket_name) as f: 77 | 78 | test_count = 1000 79 | start_time = time.time() 80 | for t in tqdm.tqdm(list(range(test_count)), 81 | desc="Running I/O Performance Test"): 82 | f.read_block(random.randint(0,f.block_count-1)) 83 | stop_time = time.time() 84 | print("Access Block Avg. Data Transmitted: %s (%.3fx)" 85 | % (MemorySize((f.bytes_sent + f.bytes_received)/float(test_count)), 86 | (f.bytes_sent + f.bytes_received)/float(test_count)/float(block_size))) 87 | print("Access Block Avg. Latency: %.2f ms" 88 | % ((stop_time-start_time)/float(test_count)*1000)) 89 | print("") 90 | 91 | if __name__ == "__main__": 92 | main() # pragma: no cover 93 | -------------------------------------------------------------------------------- /examples/encrypted_storage_sftp.py: -------------------------------------------------------------------------------- 1 | # 2 | # This example measures the performance of encrypted storage 3 | # access through an SSH client using the Secure File 4 | # Transfer Protocol (SFTP). 5 | # 6 | # In order to run this example, you must provide a host 7 | # (server) address along with valid login credentials 8 | # 9 | 10 | import os 11 | import random 12 | import time 13 | 14 | import pyoram 15 | from pyoram.util.misc import MemorySize 16 | from pyoram.encrypted_storage.encrypted_block_storage import \ 17 | EncryptedBlockStorage 18 | 19 | import paramiko 20 | import tqdm 21 | 22 | pyoram.config.SHOW_PROGRESS_BAR = True 23 | 24 | # Set SSH login credentials here 25 | # (by default, we pull these from the environment 26 | # for testing purposes) 27 | ssh_host = os.environ.get('PYORAM_SSH_TEST_HOST') 28 | ssh_username = os.environ.get('PYORAM_SSH_TEST_USERNAME') 29 | ssh_password = os.environ.get('PYORAM_SSH_TEST_PASSWORD') 30 | 31 | # Set the storage location and size 32 | storage_name = "heap.bin" 33 | # 4KB block size 34 | block_size = 4000 35 | # one block per bucket in the 36 | # storage heap of height 8 37 | block_count = 2**(8+1)-1 38 | 39 | def main(): 40 | 41 | print("Storage Name: %s" % (storage_name)) 42 | print("Block Count: %s" % (block_count)) 43 | print("Block Size: %s" % (MemorySize(block_size))) 44 | print("Total Memory: %s" 45 | % (MemorySize(block_size*block_count))) 46 | print("Actual Storage Required: %s" 47 | % (MemorySize( 48 | EncryptedBlockStorage.compute_storage_size( 49 | block_size, 50 | block_count, 51 | storage_type='sftp')))) 52 | print("") 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 | print("Setting Up Encrypted Block Storage") 64 | setup_start = time.time() 65 | with EncryptedBlockStorage.setup(storage_name, 66 | block_size, 67 | block_count, 68 | storage_type='sftp', 69 | sshclient=ssh, 70 | ignore_existing=True) as f: 71 | print("Total Setup Time: %2.f s" 72 | % (time.time()-setup_start)) 73 | print("Total Data Transmission: %s" 74 | % (MemorySize(f.bytes_sent + f.bytes_received))) 75 | print("") 76 | 77 | # We close the device and reopen it after 78 | # setup to reset the bytes sent and bytes 79 | # received stats. 80 | with EncryptedBlockStorage(storage_name, 81 | key=f.key, 82 | storage_type='sftp', 83 | sshclient=ssh) as f: 84 | 85 | test_count = 1000 86 | start_time = time.time() 87 | for t in tqdm.tqdm(list(range(test_count)), 88 | desc="Running I/O Performance Test"): 89 | f.read_block(random.randint(0,f.block_count-1)) 90 | stop_time = time.time() 91 | print("Access Block Avg. Data Transmitted: %s (%.3fx)" 92 | % (MemorySize((f.bytes_sent + f.bytes_received)/float(test_count)), 93 | (f.bytes_sent + f.bytes_received)/float(test_count)/float(block_size))) 94 | print("Access Block Avg. Latency: %.2f ms" 95 | % ((stop_time-start_time)/float(test_count)*1000)) 96 | print("") 97 | 98 | if __name__ == "__main__": 99 | main() # pragma: no cover 100 | -------------------------------------------------------------------------------- /examples/path_oram_file.py: -------------------------------------------------------------------------------- 1 | # 2 | # This example measures the performance of Path ORAM 3 | # when storage is accessed through a local file. 4 | # 5 | 6 | import os 7 | import random 8 | import time 9 | 10 | import pyoram 11 | from pyoram.util.misc import MemorySize 12 | from pyoram.oblivious_storage.tree.path_oram import \ 13 | PathORAM 14 | 15 | import tqdm 16 | 17 | pyoram.config.SHOW_PROGRESS_BAR = True 18 | 19 | # Set the storage location and size 20 | storage_name = "heap.bin" 21 | # 4KB block size 22 | block_size = 4000 23 | # one block per bucket in the 24 | # storage heap of height 8 25 | block_count = 2**(8+1)-1 26 | 27 | def main(): 28 | 29 | print("Storage Name: %s" % (storage_name)) 30 | print("Block Count: %s" % (block_count)) 31 | print("Block Size: %s" % (MemorySize(block_size))) 32 | print("Total Memory: %s" 33 | % (MemorySize(block_size*block_count))) 34 | print("Actual Storage Required: %s" 35 | % (MemorySize( 36 | PathORAM.compute_storage_size( 37 | block_size, 38 | block_count, 39 | storage_type='file')))) 40 | print("") 41 | 42 | print("Setting Up Path ORAM Storage") 43 | setup_start = time.time() 44 | with PathORAM.setup(storage_name, 45 | block_size, 46 | block_count, 47 | storage_type='file', 48 | ignore_existing=True) as f: 49 | print("Total Setup Time: %2.f s" 50 | % (time.time()-setup_start)) 51 | print("Current Stash Size: %s" 52 | % len(f.stash)) 53 | print("Total Data Transmission: %s" 54 | % (MemorySize(f.bytes_sent + f.bytes_received))) 55 | print("") 56 | 57 | # We close the device and reopen it after 58 | # setup to reset the bytes sent and bytes 59 | # received stats. 60 | with PathORAM(storage_name, 61 | f.stash, 62 | f.position_map, 63 | key=f.key, 64 | storage_type='file') as f: 65 | 66 | test_count = 100 67 | start_time = time.time() 68 | for t in tqdm.tqdm(list(range(test_count)), 69 | desc="Running I/O Performance Test"): 70 | f.read_block(random.randint(0,f.block_count-1)) 71 | stop_time = time.time() 72 | print("Current Stash Size: %s" 73 | % len(f.stash)) 74 | print("Access Block Avg. Data Transmitted: %s (%.3fx)" 75 | % (MemorySize((f.bytes_sent + f.bytes_received)/float(test_count)), 76 | (f.bytes_sent + f.bytes_received)/float(test_count)/float(block_size))) 77 | print("Access Block Avg. Latency: %.2f ms" 78 | % ((stop_time-start_time)/float(test_count)*1000)) 79 | print("") 80 | 81 | # cleanup because this is a test example 82 | os.remove(storage_name) 83 | 84 | if __name__ == "__main__": 85 | main() # pragma: no cover 86 | -------------------------------------------------------------------------------- /examples/path_oram_mmap.py: -------------------------------------------------------------------------------- 1 | # 2 | # This example measures the performance of Path ORAM 3 | # when storage is accessed through a local memory-mapped 4 | # file (mmap). 5 | # 6 | 7 | import os 8 | import random 9 | import time 10 | 11 | import pyoram 12 | from pyoram.util.misc import MemorySize 13 | from pyoram.oblivious_storage.tree.path_oram import \ 14 | PathORAM 15 | 16 | import tqdm 17 | 18 | pyoram.config.SHOW_PROGRESS_BAR = True 19 | 20 | # Set the storage location and size 21 | storage_name = "heap.bin" 22 | # 4KB block size 23 | block_size = 4000 24 | # one block per bucket in the 25 | # storage heap of height 8 26 | block_count = 2**(8+1)-1 27 | 28 | def main(): 29 | 30 | print("Storage Name: %s" % (storage_name)) 31 | print("Block Count: %s" % (block_count)) 32 | print("Block Size: %s" % (MemorySize(block_size))) 33 | print("Total Memory: %s" 34 | % (MemorySize(block_size*block_count))) 35 | print("Actual Storage Required: %s" 36 | % (MemorySize( 37 | PathORAM.compute_storage_size( 38 | block_size, 39 | block_count, 40 | storage_type='mmap')))) 41 | print("") 42 | 43 | print("Setting Up Path ORAM Storage") 44 | setup_start = time.time() 45 | with PathORAM.setup(storage_name, 46 | block_size, 47 | block_count, 48 | storage_type='mmap', 49 | ignore_existing=True) as f: 50 | print("Total Setup Time: %2.f s" 51 | % (time.time()-setup_start)) 52 | print("Current Stash Size: %s" 53 | % len(f.stash)) 54 | print("Total Data Transmission: %s" 55 | % (MemorySize(f.bytes_sent + f.bytes_received))) 56 | print("") 57 | 58 | # We close the device and reopen it after 59 | # setup to reset the bytes sent and bytes 60 | # received stats. 61 | with PathORAM(storage_name, 62 | f.stash, 63 | f.position_map, 64 | key=f.key, 65 | storage_type='mmap') as f: 66 | 67 | test_count = 100 68 | start_time = time.time() 69 | for t in tqdm.tqdm(list(range(test_count)), 70 | desc="Running I/O Performance Test"): 71 | f.read_block(random.randint(0,f.block_count-1)) 72 | stop_time = time.time() 73 | print("Current Stash Size: %s" 74 | % len(f.stash)) 75 | print("Access Block Avg. Data Transmitted: %s (%.3fx)" 76 | % (MemorySize((f.bytes_sent + f.bytes_received)/float(test_count)), 77 | (f.bytes_sent + f.bytes_received)/float(test_count)/float(block_size))) 78 | print("Access Block Avg. Latency: %.2f ms" 79 | % ((stop_time-start_time)/float(test_count)*1000)) 80 | print("") 81 | 82 | # cleanup because this is a test example 83 | os.remove(storage_name) 84 | 85 | if __name__ == "__main__": 86 | main() # pragma: no cover 87 | -------------------------------------------------------------------------------- /examples/path_oram_ram.py: -------------------------------------------------------------------------------- 1 | # 2 | # This example measures the performance of Path ORAM 3 | # when storage is accessed through RAM. 4 | # 5 | 6 | import os 7 | import random 8 | import time 9 | 10 | import pyoram 11 | from pyoram.util.misc import MemorySize 12 | from pyoram.oblivious_storage.tree.path_oram import \ 13 | PathORAM 14 | from pyoram.storage.block_storage_ram import \ 15 | BlockStorageRAM 16 | 17 | import tqdm 18 | 19 | pyoram.config.SHOW_PROGRESS_BAR = True 20 | 21 | # Set the storage location and size 22 | storage_name = "heap.bin" 23 | # 4KB block size 24 | block_size = 4000 25 | # one block per bucket in the 26 | # storage heap of height 8 27 | block_count = 2**(8+1)-1 28 | 29 | def main(): 30 | 31 | print("Storage Name: %s" % (storage_name)) 32 | print("Block Count: %s" % (block_count)) 33 | print("Block Size: %s" % (MemorySize(block_size))) 34 | print("Total Memory: %s" 35 | % (MemorySize(block_size*block_count))) 36 | print("Actual Storage Required: %s" 37 | % (MemorySize( 38 | PathORAM.compute_storage_size( 39 | block_size, 40 | block_count, 41 | storage_type='ram')))) 42 | print("") 43 | 44 | print("Setting Up Path ORAM Storage") 45 | setup_start = time.time() 46 | with PathORAM.setup(storage_name, # RAM storage ignores this argument 47 | block_size, 48 | block_count, 49 | storage_type='ram', 50 | ignore_existing=True) as f: 51 | print("Total Setup Time: %2.f s" 52 | % (time.time()-setup_start)) 53 | print("Current Stash Size: %s" 54 | % len(f.stash)) 55 | print("Total Data Transmission: %s" 56 | % (MemorySize(f.bytes_sent + f.bytes_received))) 57 | print("") 58 | 59 | # This must be done after closing the file to ensure the lock flag 60 | # is set to False in the saved data. The tofile method only exists 61 | # on BlockStorageRAM 62 | f.raw_storage.tofile(storage_name) 63 | 64 | # We close the device and reopen it after 65 | # setup to reset the bytes sent and bytes 66 | # received stats. 67 | with PathORAM(BlockStorageRAM.fromfile(storage_name), 68 | f.stash, 69 | f.position_map, 70 | key=f.key) as f: 71 | 72 | test_count = 100 73 | start_time = time.time() 74 | for t in tqdm.tqdm(list(range(test_count)), 75 | desc="Running I/O Performance Test"): 76 | f.read_block(random.randint(0,f.block_count-1)) 77 | stop_time = time.time() 78 | print("Current Stash Size: %s" 79 | % len(f.stash)) 80 | print("Access Block Avg. Data Transmitted: %s (%.3fx)" 81 | % (MemorySize((f.bytes_sent + f.bytes_received)/float(test_count)), 82 | (f.bytes_sent + f.bytes_received)/float(test_count)/float(block_size))) 83 | print("Access Block Avg. Latency: %.2f ms" 84 | % ((stop_time-start_time)/float(test_count)*1000)) 85 | print("") 86 | 87 | # cleanup because this is a test example 88 | os.remove(storage_name) 89 | 90 | if __name__ == "__main__": 91 | main() # pragma: no cover 92 | -------------------------------------------------------------------------------- /examples/path_oram_s3.py: -------------------------------------------------------------------------------- 1 | # 2 | # This example measures the performance of Path ORAM when 3 | # storage is accessed through Amazon Simple Storage Service 4 | # (S3). 5 | # 6 | # In order to run this example, you must provide a valid 7 | # S3 bucket name and have the following variables defined 8 | # in your current environment: 9 | # - AWS_ACCESS_KEY_ID 10 | # - AWS_SECRET_ACCESS_KEY 11 | # - AWS_DEFAULT_REGION 12 | # These can also be set using keywords. 13 | # 14 | 15 | import os 16 | import random 17 | import time 18 | 19 | import pyoram 20 | from pyoram.util.misc import MemorySize 21 | from pyoram.oblivious_storage.tree.path_oram import \ 22 | PathORAM 23 | 24 | import tqdm 25 | 26 | pyoram.config.SHOW_PROGRESS_BAR = True 27 | 28 | # Set S3 bucket name here 29 | # (by default, we pull this from the environment 30 | # for testing purposes) 31 | bucket_name = os.environ.get('PYORAM_AWS_TEST_BUCKET') 32 | 33 | # Set the storage location and size 34 | storage_name = "heap.bin" 35 | # 4KB block size 36 | block_size = 4000 37 | # one block per bucket in the 38 | # storage heap of height 8 39 | block_count = 2**(8+1)-1 40 | 41 | def main(): 42 | 43 | print("Storage Name: %s" % (storage_name)) 44 | print("Block Count: %s" % (block_count)) 45 | print("Block Size: %s" % (MemorySize(block_size))) 46 | print("Total Memory: %s" 47 | % (MemorySize(block_size*block_count))) 48 | print("Actual Storage Required: %s" 49 | % (MemorySize( 50 | PathORAM.compute_storage_size( 51 | block_size, 52 | block_count, 53 | storage_type='s3')))) 54 | print("") 55 | 56 | print("Setting Up Path ORAM Storage") 57 | setup_start = time.time() 58 | with PathORAM.setup(storage_name, 59 | block_size, 60 | block_count, 61 | storage_type='s3', 62 | bucket_name=bucket_name, 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 | # We close the device and reopen it after 73 | # setup to reset the bytes sent and bytes 74 | # received stats. 75 | with PathORAM(storage_name, 76 | f.stash, 77 | f.position_map, 78 | key=f.key, 79 | storage_type='s3', 80 | bucket_name=bucket_name) as f: 81 | 82 | test_count = 100 83 | start_time = time.time() 84 | for t in tqdm.tqdm(list(range(test_count)), 85 | desc="Running I/O Performance Test"): 86 | f.read_block(random.randint(0,f.block_count-1)) 87 | stop_time = time.time() 88 | print("Current Stash Size: %s" 89 | % len(f.stash)) 90 | print("Access Block Avg. Data Transmitted: %s (%.3fx)" 91 | % (MemorySize((f.bytes_sent + f.bytes_received)/float(test_count)), 92 | (f.bytes_sent + f.bytes_received)/float(test_count)/float(block_size))) 93 | print("Fetch Block Avg. Latency: %.2f ms" 94 | % ((stop_time-start_time)/float(test_count)*1000)) 95 | print("") 96 | 97 | if __name__ == "__main__": 98 | main() # pragma: no cover 99 | -------------------------------------------------------------------------------- /examples/path_oram_sftp.py: -------------------------------------------------------------------------------- 1 | # 2 | # This example measures the performance of Path ORAM when 3 | # storage is accessed through an SSH client using the Secure 4 | # File Transfer Protocol (SFTP). 5 | # 6 | # In order to run this example, you must provide a host 7 | # (server) address along with valid login credentials 8 | # 9 | 10 | import os 11 | import random 12 | import time 13 | 14 | import pyoram 15 | from pyoram.util.misc import MemorySize 16 | from pyoram.oblivious_storage.tree.path_oram import \ 17 | PathORAM 18 | 19 | import paramiko 20 | import tqdm 21 | 22 | pyoram.config.SHOW_PROGRESS_BAR = True 23 | 24 | # Set SSH login credentials here 25 | # (by default, we pull these from the environment 26 | # for testing purposes) 27 | ssh_host = os.environ.get('PYORAM_SSH_TEST_HOST') 28 | ssh_username = os.environ.get('PYORAM_SSH_TEST_USERNAME') 29 | ssh_password = os.environ.get('PYORAM_SSH_TEST_PASSWORD') 30 | 31 | # Set the storage location and size 32 | storage_name = "heap.bin" 33 | # 4KB block size 34 | block_size = 4000 35 | # one block per bucket in the 36 | # storage heap of height 8 37 | block_count = 2**(8+1)-1 38 | 39 | def main(): 40 | 41 | print("Storage Name: %s" % (storage_name)) 42 | print("Block Count: %s" % (block_count)) 43 | print("Block Size: %s" % (MemorySize(block_size))) 44 | print("Total Memory: %s" 45 | % (MemorySize(block_size*block_count))) 46 | print("Actual Storage Required: %s" 47 | % (MemorySize( 48 | PathORAM.compute_storage_size( 49 | block_size, 50 | block_count, 51 | storage_type='sftp')))) 52 | print("") 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 | print("Setting Up Path ORAM Storage") 64 | setup_start = time.time() 65 | with PathORAM.setup(storage_name, 66 | block_size, 67 | block_count, 68 | storage_type='sftp', 69 | sshclient=ssh, 70 | ignore_existing=True) as f: 71 | print("Total Setup Time: %.2f s" 72 | % (time.time()-setup_start)) 73 | print("Current Stash Size: %s" 74 | % len(f.stash)) 75 | print("Total Data Transmission: %s" 76 | % (MemorySize(f.bytes_sent + f.bytes_received))) 77 | print("") 78 | 79 | # We close the device and reopen it after 80 | # setup to reset the bytes sent and bytes 81 | # received stats. 82 | with PathORAM(storage_name, 83 | f.stash, 84 | f.position_map, 85 | key=f.key, 86 | storage_type='sftp', 87 | sshclient=ssh) as f: 88 | 89 | test_count = 100 90 | start_time = time.time() 91 | for t in tqdm.tqdm(list(range(test_count)), 92 | desc="Running I/O Performance Test"): 93 | f.read_block(random.randint(0,f.block_count-1)) 94 | stop_time = time.time() 95 | print("Current Stash Size: %s" 96 | % len(f.stash)) 97 | print("Access Block Avg. Data Transmitted: %s (%.3fx)" 98 | % (MemorySize((f.bytes_sent + f.bytes_received)/float(test_count)), 99 | (f.bytes_sent + f.bytes_received)/float(test_count)/float(block_size))) 100 | print("Fetch Block Avg. Latency: %.2f ms" 101 | % ((stop_time-start_time)/float(test_count)*1000)) 102 | print("") 103 | 104 | if __name__ == "__main__": 105 | main() # pragma: no cover 106 | -------------------------------------------------------------------------------- /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 116 | -------------------------------------------------------------------------------- /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 100 | -------------------------------------------------------------------------------- /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 94 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | cryptography 2 | paramiko 3 | boto3 4 | cffi>=1.0.0 5 | six 6 | tqdm 7 | -------------------------------------------------------------------------------- /setup.cfg: -------------------------------------------------------------------------------- 1 | [metadata] 2 | description-file = README.rst 3 | 4 | [bdist_wheel] 5 | # supports python3 6 | universal=1 7 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | import platform 4 | from setuptools import setup, find_packages 5 | from codecs import open 6 | 7 | here = os.path.abspath(os.path.dirname(__file__)) 8 | 9 | about = {} 10 | with open(os.path.join("src", "pyoram", "__about__.py")) as f: 11 | exec(f.read(), about) 12 | 13 | # Get the long description from the README file 14 | def _readme(): 15 | with open(os.path.join(here, 'README.rst'), encoding='utf-8') as f: 16 | return f.read() 17 | 18 | setup_requirements = [] 19 | requirements = ['cryptography', 20 | 'paramiko', 21 | 'boto3', 22 | 'six', 23 | 'tqdm'] 24 | 25 | if platform.python_implementation() == "PyPy": 26 | if sys.pypy_version_info < (2, 6): 27 | raise RuntimeError( 28 | "PyORAM is not compatible with PyPy < 2.6. Please " 29 | "upgrade PyPy to use this library.") 30 | else: 31 | if sys.version_info <= (2, 6): 32 | raise RuntimeError( 33 | "PyORAM is not compatible with Python < 2.7. Please " 34 | "upgrade Python to use this library.") 35 | requirements.append("cffi>=1.0.0") 36 | setup_requirements.append("cffi>=1.0.0") 37 | 38 | setup( 39 | name=about['__title__'], 40 | version=about['__version__'], 41 | description=about['__summary__'], 42 | long_description=_readme(), 43 | url=about['__uri__'], 44 | author=about['__author__'], 45 | author_email=about['__email__'], 46 | license=about['__license__'], 47 | # https://pypi.python.org/pypi?%3Aaction=list_classifiers 48 | classifiers=[ 49 | 'Development Status :: 4 - Beta', 50 | 'Intended Audience :: Science/Research', 51 | 'Topic :: Security :: Cryptography', 52 | "Natural Language :: English", 53 | 'License :: OSI Approved :: MIT License', 54 | 'Programming Language :: Python :: 2.7', 55 | 'Programming Language :: Python :: 3', 56 | 'Programming Language :: Python :: 3.3', 57 | 'Programming Language :: Python :: 3.4', 58 | 'Programming Language :: Python :: 3.5', 59 | 'Programming Language :: Python :: 3.6', 60 | 'Programming Language :: Python :: 3.7', 61 | 'Programming Language :: Python :: 3.8', 62 | 'Programming Language :: Python :: Implementation :: CPython', 63 | 'Programming Language :: Python :: Implementation :: PyPy', 64 | ], 65 | keywords='oram, storage, privacy, cryptography, cloud storage', 66 | package_dir={'': 'src'}, 67 | packages=find_packages(where="src", exclude=["_cffi_src", "_cffi_src.*"]), 68 | setup_requires=setup_requirements, 69 | install_requires=requirements, 70 | cffi_modules=["src/_cffi_src/virtual_heap_helper_build.py:ffi"], 71 | # use MANIFEST.in 72 | include_package_data=True, 73 | test_suite='nose2.collector.collector', 74 | tests_require=['unittest','nose2'] 75 | ) 76 | -------------------------------------------------------------------------------- /src/_cffi_src/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ghackebeil/PyORAM/53e109dfb1ecec52348a70ddc64fae65eea7490a/src/_cffi_src/__init__.py -------------------------------------------------------------------------------- /src/_cffi_src/virtual_heap_helper_build.py: -------------------------------------------------------------------------------- 1 | import cffi 2 | 3 | # 4 | # C functions that speed up commonly 5 | # executed heap calculations in tree-based 6 | # orams 7 | # 8 | 9 | ffi = cffi.FFI() 10 | ffi.cdef( 11 | """ 12 | int calculate_bucket_level(unsigned int k, 13 | unsigned long long b); 14 | int calculate_last_common_level(unsigned int k, 15 | unsigned long long b1, 16 | unsigned long long b2); 17 | """) 18 | 19 | ffi.set_source("pyoram.util._virtual_heap_helper", 20 | """ 21 | #include 22 | #include 23 | 24 | int calculate_bucket_level(unsigned int k, 25 | unsigned long long b) 26 | { 27 | unsigned int h; 28 | unsigned long long pow; 29 | if (k == 2) { 30 | // This is simply log2floor(b+1) 31 | h = 0; 32 | b += 1; 33 | while (b >>= 1) {++h;} 34 | return h; 35 | } 36 | b = (k - 1) * (b + 1) + 1; 37 | h = 0; 38 | pow = k; 39 | while (pow < b) {++h; pow *= k;} 40 | return h; 41 | } 42 | 43 | int calculate_last_common_level(unsigned int k, 44 | unsigned long long b1, 45 | unsigned long long b2) 46 | { 47 | int level1, level2; 48 | level1 = calculate_bucket_level(k, b1); 49 | level2 = calculate_bucket_level(k, b2); 50 | if (level1 != level2) { 51 | if (level1 > level2) { 52 | while (level1 != level2) { 53 | b1 = (b1 - 1)/k; 54 | --level1; 55 | } 56 | } 57 | else { 58 | while (level2 != level1) { 59 | b2 = (b2 - 1)/k; 60 | --level2; 61 | } 62 | } 63 | } 64 | while (b1 != b2) { 65 | b1 = (b1 - 1)/k; 66 | b2 = (b2 - 1)/k; 67 | --level1; 68 | } 69 | return level1; 70 | } 71 | """) 72 | 73 | if __name__ == "__main__": 74 | ffi.compile() 75 | -------------------------------------------------------------------------------- /src/pyoram/__about__.py: -------------------------------------------------------------------------------- 1 | __all__ = ('__title__', 2 | '__summary__', 3 | '__uri__', 4 | '__version__', 5 | '__author__', 6 | '__email__', 7 | '__license__', 8 | '__copyright__') 9 | 10 | __title__ = 'PyORAM' 11 | __summary__ = 'Python-based Oblivious RAM' 12 | __uri__ = 'https://github.com/ghackebeil/PyORAM' 13 | __version__ = '0.3.0' 14 | __author__ = 'Gabriel A. Hackebeil' 15 | __email__ = 'gabe.hackebeil@gmail.com' 16 | __license__ = 'MIT' 17 | __copyright__ = 'Copyright {0}'.format(__author__) 18 | -------------------------------------------------------------------------------- /src/pyoram/__init__.py: -------------------------------------------------------------------------------- 1 | from pyoram.__about__ import __version__ 2 | 3 | def _configure_logging(): 4 | import os 5 | import logging 6 | 7 | log = logging.getLogger("pyoram") 8 | formatter = logging.Formatter( 9 | fmt=("[%(asctime)s.%(msecs)03d," 10 | "%(name)s,%(levelname)s] %(threadName)s %(message)s"), 11 | datefmt="%Y-%m-%d %H:%M:%S") 12 | 13 | level = os.environ.get("PYORAM_LOGLEVEL", "WARNING") 14 | logfilename = os.environ.get("PYORAM_LOGFILE", None) 15 | if len(logging.root.handlers) == 0: 16 | # configure the logging with some sensible 17 | # defaults. 18 | try: 19 | import tempfile 20 | tempfile = tempfile.TemporaryFile(dir=".") 21 | tempfile.close() 22 | except OSError: 23 | # cannot write in current directory, use the 24 | # console logger 25 | handler = logging.StreamHandler() 26 | else: 27 | if logfilename is None: 28 | handler = logging.StreamHandler() 29 | else: 30 | # set up a basic logfile in current directory 31 | handler = logging.FileHandler(logfilename) 32 | handler.setFormatter(formatter) 33 | handler.setLevel(level) 34 | log.addHandler(handler) 35 | log.setLevel(level) 36 | log.info("PyORAM log configured using built-in " 37 | "defaults, level=%s", level) 38 | 39 | _configure_logging() 40 | del _configure_logging 41 | 42 | def _configure_pyoram(): 43 | class _Configure(object): 44 | __slots__ = ("SHOW_PROGRESS_BAR",) 45 | def __init__(self): 46 | self.SHOW_PROGRESS_BAR = False 47 | return _Configure() 48 | config = _configure_pyoram() 49 | del _configure_pyoram 50 | 51 | import pyoram.util 52 | import pyoram.crypto 53 | import pyoram.storage 54 | import pyoram.encrypted_storage 55 | import pyoram.oblivious_storage 56 | -------------------------------------------------------------------------------- /src/pyoram/crypto/__init__.py: -------------------------------------------------------------------------------- 1 | import pyoram.crypto.aes 2 | -------------------------------------------------------------------------------- /src/pyoram/crypto/aes.py: -------------------------------------------------------------------------------- 1 | __all__ = ("AES",) 2 | 3 | import os 4 | import cryptography.hazmat.primitives.ciphers 5 | import cryptography.hazmat.backends 6 | 7 | _backend = cryptography.hazmat.backends.default_backend() 8 | _aes = cryptography.hazmat.primitives.ciphers.algorithms.AES 9 | _cipher = cryptography.hazmat.primitives.ciphers.Cipher 10 | _ctrmode = cryptography.hazmat.primitives.ciphers.modes.CTR 11 | _gcmmode = cryptography.hazmat.primitives.ciphers.modes.GCM 12 | 13 | class AES(object): 14 | 15 | key_sizes = [k//8 for k in sorted(_aes.key_sizes)] 16 | block_size = _aes.block_size//8 17 | 18 | @staticmethod 19 | def KeyGen(size_bytes): 20 | assert size_bytes in AES.key_sizes 21 | return os.urandom(size_bytes) 22 | 23 | @staticmethod 24 | def CTREnc(key, plaintext): 25 | iv = os.urandom(AES.block_size) 26 | cipher = _cipher(_aes(key), _ctrmode(iv), backend=_backend).encryptor() 27 | return iv + cipher.update(plaintext) + cipher.finalize() 28 | 29 | @staticmethod 30 | def CTRDec(key, ciphertext): 31 | iv = ciphertext[:AES.block_size] 32 | cipher = _cipher(_aes(key), _ctrmode(iv), backend=_backend).decryptor() 33 | return cipher.update(ciphertext[AES.block_size:]) + \ 34 | cipher.finalize() 35 | 36 | @staticmethod 37 | def GCMEnc(key, plaintext): 38 | iv = os.urandom(AES.block_size) 39 | cipher = _cipher(_aes(key), _gcmmode(iv), backend=_backend).encryptor() 40 | return iv + cipher.update(plaintext) + cipher.finalize() + cipher.tag 41 | 42 | @staticmethod 43 | def GCMDec(key, ciphertext): 44 | iv = ciphertext[:AES.block_size] 45 | tag = ciphertext[-AES.block_size:] 46 | cipher = _cipher(_aes(key), _gcmmode(iv, tag), backend=_backend).decryptor() 47 | return cipher.update(ciphertext[AES.block_size:-AES.block_size]) + \ 48 | cipher.finalize() 49 | -------------------------------------------------------------------------------- /src/pyoram/encrypted_storage/__init__.py: -------------------------------------------------------------------------------- 1 | import pyoram.encrypted_storage.encrypted_block_storage 2 | import pyoram.encrypted_storage.encrypted_heap_storage 3 | import pyoram.encrypted_storage.top_cached_encrypted_heap_storage 4 | -------------------------------------------------------------------------------- /src/pyoram/encrypted_storage/encrypted_block_storage.py: -------------------------------------------------------------------------------- 1 | __all__ = ('EncryptedBlockStorage',) 2 | 3 | import struct 4 | import hmac 5 | import hashlib 6 | 7 | from pyoram.storage.block_storage import (BlockStorageInterface, 8 | BlockStorageTypeFactory) 9 | from pyoram.crypto.aes import AES 10 | 11 | import six 12 | 13 | class EncryptedBlockStorageInterface(BlockStorageInterface): 14 | 15 | # 16 | # Abstract Interface 17 | # 18 | 19 | @property 20 | def key(self, *args, **kwds): 21 | raise NotImplementedError # pragma: no cover 22 | @property 23 | def raw_storage(self, *args, **kwds): 24 | raise NotImplementedError # pragma: no cover 25 | 26 | class EncryptedBlockStorage(EncryptedBlockStorageInterface): 27 | 28 | _index_struct_string = "!"+("x"*hashlib.sha384().digest_size)+"?" 29 | _index_offset = struct.calcsize(_index_struct_string) 30 | _verify_struct_string = "!LLL" 31 | _verify_size = struct.calcsize(_verify_struct_string) 32 | 33 | def __init__(self, storage, **kwds): 34 | self._key = kwds.pop('key', None) 35 | if self._key is None: 36 | raise ValueError( 37 | "An encryption key is required using " 38 | "the 'key' keyword.") 39 | if isinstance(storage, BlockStorageInterface): 40 | storage_owned = False 41 | self._storage = storage 42 | if len(kwds): 43 | raise ValueError( 44 | "Keywords not used when initializing " 45 | "with a storage device: %s" 46 | % (str(kwds))) 47 | else: 48 | storage_owned = True 49 | storage_type = kwds.pop('storage_type', 'file') 50 | self._storage = \ 51 | BlockStorageTypeFactory(storage_type)(storage, **kwds) 52 | 53 | try: 54 | header_data = AES.GCMDec(self._key, 55 | self._storage.header_data) 56 | (self._ismodegcm,) = struct.unpack( 57 | self._index_struct_string, 58 | header_data[:self._index_offset]) 59 | self._verify_digest = header_data[:hashlib.sha384().digest_size] 60 | 61 | verify = hmac.HMAC( 62 | key=self.key, 63 | msg=struct.pack(self._verify_struct_string, 64 | self._storage.block_size, 65 | self._storage.block_count, 66 | len(self._storage.header_data)), 67 | digestmod=hashlib.sha384) 68 | if verify.digest() != self._verify_digest: 69 | raise ValueError( 70 | "HMAC of plaintext index data does not match") 71 | if self._ismodegcm: 72 | self._encrypt_block_func = AES.GCMEnc 73 | self._decrypt_block_func = AES.GCMDec 74 | else: 75 | self._encrypt_block_func = AES.CTREnc 76 | self._decrypt_block_func = AES.CTRDec 77 | except: 78 | if storage_owned: 79 | self._storage.close() 80 | raise 81 | 82 | # 83 | # Define EncryptedBlockStorageInterface Methods 84 | # 85 | 86 | @property 87 | def key(self): 88 | return self._key 89 | 90 | @property 91 | def raw_storage(self): 92 | return self._storage 93 | 94 | # 95 | # Define BlockStorageInterface Methods 96 | # 97 | 98 | def clone_device(self): 99 | return EncryptedBlockStorage(self._storage.clone_device(), 100 | key=self.key) 101 | 102 | @classmethod 103 | def compute_storage_size(cls, 104 | block_size, 105 | block_count, 106 | aes_mode='ctr', 107 | storage_type='file', 108 | ignore_header=False, 109 | **kwds): 110 | assert (block_size > 0) and (block_size == int(block_size)) 111 | assert (block_count > 0) and (block_count == int(block_count)) 112 | assert aes_mode in ('ctr', 'gcm') 113 | if not isinstance(storage_type, BlockStorageInterface): 114 | storage_type = BlockStorageTypeFactory(storage_type) 115 | 116 | if aes_mode == 'ctr': 117 | extra_block_data = AES.block_size 118 | else: 119 | assert aes_mode == 'gcm' 120 | extra_block_data = 2 * AES.block_size 121 | if ignore_header: 122 | return (extra_block_data * block_count) + \ 123 | storage_type.compute_storage_size( 124 | block_size, 125 | block_count, 126 | ignore_header=True, 127 | **kwds) 128 | else: 129 | return cls._index_offset + \ 130 | 2 * AES.block_size + \ 131 | (extra_block_data * block_count) + \ 132 | storage_type.compute_storage_size( 133 | block_size, 134 | block_count, 135 | ignore_header=False, 136 | **kwds) 137 | 138 | @classmethod 139 | def setup(cls, 140 | storage_name, 141 | block_size, 142 | block_count, 143 | aes_mode='ctr', 144 | key_size=None, 145 | key=None, 146 | storage_type='file', 147 | initialize=None, 148 | **kwds): 149 | 150 | if (key is not None) and (key_size is not None): 151 | raise ValueError( 152 | "Only one of 'key' or 'keysize' keywords can " 153 | "be specified at a time") 154 | if key is None: 155 | if key_size is None: 156 | key_size = 32 157 | if key_size not in AES.key_sizes: 158 | raise ValueError( 159 | "Invalid key size: %s" % (key_size)) 160 | key = AES.KeyGen(key_size) 161 | else: 162 | if len(key) not in AES.key_sizes: 163 | raise ValueError( 164 | "Invalid key size: %s" % (len(key))) 165 | 166 | if (block_size <= 0) or (block_size != int(block_size)): 167 | raise ValueError( 168 | "Block size (bytes) must be a positive integer: %s" 169 | % (block_size)) 170 | 171 | ismodegcm = None 172 | encrypt_block_func = None 173 | encrypted_block_size = block_size 174 | if aes_mode == 'ctr': 175 | ismodegcm = False 176 | encrypt_block_func = AES.CTREnc 177 | encrypted_block_size += AES.block_size 178 | elif aes_mode == 'gcm': 179 | ismodegcm = True 180 | encrypt_block_func = AES.GCMEnc 181 | encrypted_block_size += (2 * AES.block_size) 182 | else: 183 | raise ValueError( 184 | "AES encryption mode must be one of 'ctr' or 'gcm'. " 185 | "Invalid value: %s" % (aes_mode)) 186 | assert ismodegcm is not None 187 | assert encrypt_block_func is not None 188 | 189 | if not isinstance(storage_type, BlockStorageInterface): 190 | storage_type = BlockStorageTypeFactory(storage_type) 191 | 192 | if initialize is None: 193 | zeros = bytes(bytearray(block_size)) 194 | initialize = lambda i: zeros 195 | def encrypted_initialize(i): 196 | return encrypt_block_func(key, initialize(i)) 197 | kwds['initialize'] = encrypted_initialize 198 | 199 | user_header_data = kwds.get('header_data', bytes()) 200 | if type(user_header_data) is not bytes: 201 | raise TypeError( 202 | "'header_data' must be of type bytes. " 203 | "Invalid type: %s" % (type(user_header_data))) 204 | # we generate the first time simply to 205 | # compute the length 206 | tmp = hmac.HMAC( 207 | key=key, 208 | msg=struct.pack(cls._verify_struct_string, 209 | encrypted_block_size, 210 | block_count, 211 | 0), 212 | digestmod=hashlib.sha384).digest() 213 | header_data = bytearray(struct.pack(cls._index_struct_string, 214 | ismodegcm)) 215 | header_data[:hashlib.sha384().digest_size] = tmp 216 | header_data = header_data + user_header_data 217 | header_data = AES.GCMEnc(key, bytes(header_data)) 218 | # now that we know the length of the header data 219 | # being sent to the underlying storage we can 220 | # compute the real hmac 221 | verify_digest = hmac.HMAC( 222 | key=key, 223 | msg=struct.pack(cls._verify_struct_string, 224 | encrypted_block_size, 225 | block_count, 226 | len(header_data)), 227 | digestmod=hashlib.sha384).digest() 228 | header_data = bytearray(struct.pack(cls._index_struct_string, 229 | ismodegcm)) 230 | header_data[:hashlib.sha384().digest_size] = verify_digest 231 | header_data = header_data + user_header_data 232 | kwds['header_data'] = AES.GCMEnc(key, bytes(header_data)) 233 | 234 | return EncryptedBlockStorage( 235 | storage_type.setup(storage_name, 236 | encrypted_block_size, 237 | block_count, 238 | **kwds), 239 | key=key) 240 | 241 | @property 242 | def header_data(self): 243 | return AES.GCMDec(self._key, 244 | self._storage.header_data)\ 245 | [self._index_offset:] 246 | 247 | @property 248 | def block_count(self): 249 | return self._storage.block_count 250 | 251 | @property 252 | def block_size(self): 253 | if self._ismodegcm: 254 | return self._storage.block_size - 2 * AES.block_size 255 | else: 256 | return self._storage.block_size - AES.block_size 257 | 258 | @property 259 | def storage_name(self): 260 | return self._storage.storage_name 261 | 262 | def update_header_data(self, new_header_data): 263 | self._storage.update_header_data( 264 | AES.GCMEnc( 265 | self.key, 266 | AES.GCMDec(self._key, 267 | self._storage.header_data)\ 268 | [:self._index_offset] + \ 269 | new_header_data)) 270 | 271 | def close(self): 272 | self._storage.close() 273 | 274 | def read_block(self, i): 275 | return self._decrypt_block_func( 276 | self._key, 277 | self._storage.read_block(i)) 278 | 279 | def read_blocks(self, indices, *args, **kwds): 280 | return [self._decrypt_block_func(self._key, b) 281 | for b in self._storage.read_blocks(indices, *args, **kwds)] 282 | 283 | def yield_blocks(self, indices, *args, **kwds): 284 | for b in self._storage.yield_blocks(indices, *args, **kwds): 285 | yield self._decrypt_block_func(self._key, b) 286 | 287 | def write_block(self, i, block, *args, **kwds): 288 | self._storage.write_block( 289 | i, 290 | self._encrypt_block_func(self._key, block), 291 | *args, **kwds) 292 | 293 | def write_blocks(self, indices, blocks, *args, **kwds): 294 | enc_blocks = [] 295 | for i, b in zip(indices, blocks): 296 | enc_blocks.append( 297 | self._encrypt_block_func(self._key, b)) 298 | self._storage.write_blocks(indices, enc_blocks, *args, **kwds) 299 | 300 | @property 301 | def bytes_sent(self): 302 | return self._storage.bytes_sent 303 | 304 | @property 305 | def bytes_received(self): 306 | return self._storage.bytes_received 307 | -------------------------------------------------------------------------------- /src/pyoram/encrypted_storage/encrypted_heap_storage.py: -------------------------------------------------------------------------------- 1 | __all__ = ('EncryptedHeapStorage',) 2 | 3 | import struct 4 | 5 | from pyoram.util.virtual_heap import SizedVirtualHeap 6 | from pyoram.storage.heap_storage import \ 7 | (HeapStorageInterface, 8 | HeapStorage) 9 | from pyoram.encrypted_storage.encrypted_block_storage import \ 10 | (EncryptedBlockStorageInterface, 11 | EncryptedBlockStorage) 12 | 13 | class EncryptedHeapStorageInterface(HeapStorageInterface): 14 | 15 | # 16 | # Abstract Interface 17 | # 18 | 19 | @property 20 | def key(self, *args, **kwds): 21 | raise NotImplementedError # pragma: no cover 22 | @property 23 | def raw_storage(self, *args, **kwds): 24 | raise NotImplementedError # pragma: no cover 25 | 26 | class EncryptedHeapStorage(HeapStorage, 27 | EncryptedHeapStorageInterface): 28 | 29 | def __init__(self, storage, **kwds): 30 | 31 | if isinstance(storage, EncryptedBlockStorageInterface): 32 | if len(kwds): 33 | raise ValueError( 34 | "Keywords not used when initializing " 35 | "with a storage device: %s" 36 | % (str(kwds))) 37 | else: 38 | storage = EncryptedBlockStorage(storage, **kwds) 39 | 40 | super(EncryptedHeapStorage, self).__init__(storage) 41 | 42 | # 43 | # Define EncryptedHeapStorageInterface Methods 44 | # 45 | 46 | @property 47 | def key(self): 48 | return self._storage.key 49 | 50 | @property 51 | def raw_storage(self): 52 | return self._storage.raw_storage 53 | 54 | # 55 | # Define HeapStorageInterface Methods 56 | # (override what is defined on HeapStorage) 57 | 58 | def clone_device(self): 59 | return EncryptedHeapStorage(self._storage.clone_device()) 60 | 61 | @classmethod 62 | def compute_storage_size(cls, 63 | block_size, 64 | heap_height, 65 | blocks_per_bucket=1, 66 | heap_base=2, 67 | ignore_header=False, 68 | **kwds): 69 | assert (block_size > 0) and (block_size == int(block_size)) 70 | assert heap_height >= 0 71 | assert blocks_per_bucket >= 1 72 | assert heap_base >= 2 73 | assert 'block_count' not in kwds 74 | vheap = SizedVirtualHeap( 75 | heap_base, 76 | heap_height, 77 | blocks_per_bucket=blocks_per_bucket) 78 | if ignore_header: 79 | return EncryptedBlockStorage.compute_storage_size( 80 | vheap.blocks_per_bucket * block_size, 81 | vheap.bucket_count(), 82 | ignore_header=True, 83 | **kwds) 84 | else: 85 | return cls._header_offset + \ 86 | EncryptedBlockStorage.compute_storage_size( 87 | vheap.blocks_per_bucket * block_size, 88 | vheap.bucket_count(), 89 | ignore_header=False, 90 | **kwds) 91 | 92 | @classmethod 93 | def setup(cls, 94 | storage_name, 95 | block_size, 96 | heap_height, 97 | blocks_per_bucket=1, 98 | heap_base=2, 99 | **kwds): 100 | if 'block_count' in kwds: 101 | raise ValueError("'block_count' keyword is not accepted") 102 | if heap_height < 0: 103 | raise ValueError( 104 | "heap height must be 0 or greater. Invalid value: %s" 105 | % (heap_height)) 106 | if blocks_per_bucket < 1: 107 | raise ValueError( 108 | "blocks_per_bucket must be 1 or greater. " 109 | "Invalid value: %s" % (blocks_per_bucket)) 110 | if heap_base < 2: 111 | raise ValueError( 112 | "heap base must be 2 or greater. Invalid value: %s" 113 | % (heap_base)) 114 | 115 | vheap = SizedVirtualHeap( 116 | heap_base, 117 | heap_height, 118 | blocks_per_bucket=blocks_per_bucket) 119 | 120 | user_header_data = kwds.pop('header_data', bytes()) 121 | if type(user_header_data) is not bytes: 122 | raise TypeError( 123 | "'header_data' must be of type bytes. " 124 | "Invalid type: %s" % (type(user_header_data))) 125 | kwds['header_data'] = \ 126 | struct.pack(cls._header_struct_string, 127 | heap_base, 128 | heap_height, 129 | blocks_per_bucket) + \ 130 | user_header_data 131 | 132 | return EncryptedHeapStorage( 133 | EncryptedBlockStorage.setup( 134 | storage_name, 135 | vheap.blocks_per_bucket * block_size, 136 | vheap.bucket_count(), 137 | **kwds)) 138 | 139 | #@property 140 | #def header_data(...) 141 | 142 | #@property 143 | #def bucket_count(...) 144 | 145 | #@property 146 | #def bucket_size(...) 147 | 148 | #@property 149 | #def blocks_per_bucket(...) 150 | 151 | #@property 152 | #def storage_name(...) 153 | 154 | #@property 155 | #def virtual_heap(...) 156 | 157 | #@property 158 | #def bucket_storage(...) 159 | 160 | #def update_header_data(...) 161 | 162 | #def close(...) 163 | 164 | #def read_path(...) 165 | 166 | #def write_path(...) 167 | 168 | #@property 169 | #def bytes_sent(...) 170 | 171 | #@property 172 | #def bytes_received(...) 173 | -------------------------------------------------------------------------------- /src/pyoram/encrypted_storage/top_cached_encrypted_heap_storage.py: -------------------------------------------------------------------------------- 1 | __all__ = ('TopCachedEncryptedHeapStorage',) 2 | 3 | import logging 4 | import tempfile 5 | import mmap 6 | 7 | import pyoram 8 | from pyoram.util.virtual_heap import SizedVirtualHeap 9 | from pyoram.encrypted_storage.encrypted_heap_storage import \ 10 | (EncryptedHeapStorageInterface, 11 | EncryptedHeapStorage) 12 | 13 | import tqdm 14 | import six 15 | from six.moves import xrange 16 | 17 | log = logging.getLogger("pyoram") 18 | 19 | class TopCachedEncryptedHeapStorage(EncryptedHeapStorageInterface): 20 | """ 21 | An encrypted block storage device for accessing memory 22 | organized as a heap, where the top 1 or more levels can 23 | be cached in local memory. This achieves two things: 24 | 25 | (1) Reduces the number of buckets that need to be read 26 | from or written to external storage for a given 27 | path I/O operation. 28 | (2) Allows certain block storage devices to achieve 29 | concurrency across path writes by partioning the 30 | storage space into independent subheaps starting 31 | below the cache line. 32 | 33 | This devices takes as input an existing encrypted heap 34 | storage device. This class should not be cloned or used 35 | to setup storage, but rather used as a wrapper class for 36 | an existing heap storage device to speed up a bulk set 37 | of I/O requests. The original heap storage device should 38 | not be used after it is wrapped by this class. This 39 | class will close the original device when closing 40 | itself. 41 | 42 | The number of cached levels (starting from the root 43 | bucket at level 0) can be set with the 'cached_levels' 44 | keyword (>= 1). 45 | 46 | By default, this will create an independent storage 47 | device capable of reading from and writing to the 48 | original storage devices memory for each independent 49 | subheap (if any) below the last cached level. The 50 | 'concurrency_level' keyword can be used to limit the 51 | number of concurrent devices to some level below the 52 | cache line (>= 0, <= 'cached_levels'). 53 | 54 | Values for 'cached_levels' and 'concurrency_level' will 55 | be automatically reduced when they are larger than what 56 | is allowed by the heap size. 57 | """ 58 | 59 | def __new__(cls, *args, **kwds): 60 | if kwds.get("cached_levels", 1) == 0: 61 | assert len(args) == 1 62 | storage = args[0] 63 | storage.cached_bucket_data = bytes() 64 | return storage 65 | else: 66 | return super(TopCachedEncryptedHeapStorage, cls).\ 67 | __new__(cls) 68 | 69 | def __init__(self, 70 | heap_storage, 71 | cached_levels=1, 72 | concurrency_level=None): 73 | assert isinstance(heap_storage, EncryptedHeapStorage) 74 | assert cached_levels != 0 75 | vheap = heap_storage.virtual_heap 76 | if cached_levels < 0: 77 | cached_levels = vheap.levels 78 | if concurrency_level is None: 79 | concurrency_level = cached_levels 80 | assert concurrency_level >= 0 81 | cached_levels = min(vheap.levels, cached_levels) 82 | concurrency_level = min(cached_levels, concurrency_level) 83 | self._external_level = cached_levels 84 | total_buckets = sum(vheap.bucket_count_at_level(l) 85 | for l in xrange(cached_levels)) 86 | 87 | self._root_device = heap_storage 88 | # clone before we download the cache so that we can 89 | # track bytes transferred during read/write requests 90 | # (separate from the cached download) 91 | self._concurrent_devices = \ 92 | {vheap.first_bucket_at_level(0): self._root_device.clone_device()} 93 | 94 | self._cached_bucket_count = total_buckets 95 | self._cached_buckets_tempfile = tempfile.TemporaryFile() 96 | self._cached_buckets_tempfile.seek(0) 97 | with tqdm.tqdm(desc=("Downloading %s Cached Heap Buckets" 98 | % (self._cached_bucket_count)), 99 | total=self._cached_bucket_count*self._root_device.bucket_size, 100 | unit="B", 101 | unit_scale=True, 102 | disable=not pyoram.config.SHOW_PROGRESS_BAR) as progress_bar: 103 | for b, bucket in enumerate( 104 | self._root_device.bucket_storage.yield_blocks( 105 | xrange(vheap.first_bucket_at_level(cached_levels)))): 106 | self._cached_buckets_tempfile.write(bucket) 107 | progress_bar.update(self._root_device.bucket_size) 108 | self._cached_buckets_tempfile.flush() 109 | self._cached_buckets_mmap = mmap.mmap( 110 | self._cached_buckets_tempfile.fileno(), 0) 111 | 112 | log.info("%s: Cloning %s sub-heap devices" 113 | % (self.__class__.__name__, vheap.bucket_count_at_level(concurrency_level))) 114 | # Avoid cloning devices when the cache line is at the root 115 | # bucket or when the entire heap is cached 116 | if (concurrency_level > 0) and \ 117 | (concurrency_level <= vheap.last_level): 118 | for b in xrange(vheap.first_bucket_at_level(concurrency_level), 119 | vheap.first_bucket_at_level(concurrency_level+1)): 120 | try: 121 | self._concurrent_devices[b] = self._root_device.clone_device() 122 | except: # pragma: no cover 123 | log.error( # pragma: no cover 124 | "%s: Exception encountered " # pragma: no cover 125 | "while cloning device. " # pragma: no cover 126 | "Closing storage." # pragma: no cover 127 | % (self.__class__.__name__)) # pragma: no cover 128 | self.close() # pragma: no cover 129 | raise # pragma: no cover 130 | 131 | self._subheap_storage = {} 132 | # Avoid populating this dictionary when the entire 133 | # heap is cached 134 | if self._external_level <= vheap.last_level: 135 | for b in xrange(vheap.first_bucket_at_level(self._external_level), 136 | vheap.first_bucket_at_level(self._external_level+1)): 137 | node = vheap.Node(b) 138 | while node.bucket not in self._concurrent_devices: 139 | node = node.parent_node() 140 | assert node.bucket >= 0 141 | assert node.level == concurrency_level 142 | self._subheap_storage[b] = self._concurrent_devices[node.bucket] 143 | 144 | # 145 | # Additional Methods 146 | # 147 | 148 | @property 149 | def cached_bucket_data(self): 150 | return self._cached_buckets_mmap 151 | 152 | # 153 | # Define EncryptedHeapStorageInterface Methods 154 | # 155 | 156 | @property 157 | def key(self): 158 | return self._root_device.key 159 | 160 | @property 161 | def raw_storage(self): 162 | return self._root_device.raw_storage 163 | 164 | # 165 | # Define HeapStorageInterface Methods 166 | # 167 | 168 | def clone_device(self, *args, **kwds): 169 | raise NotImplementedError( # pragma: no cover 170 | "Class is not designed for cloning") # pragma: no cover 171 | 172 | @classmethod 173 | def compute_storage_size(cls, *args, **kwds): 174 | return EncryptedHeapStorage.compute_storage_size(*args, **kwds) 175 | 176 | @classmethod 177 | def setup(cls, *args, **kwds): 178 | raise NotImplementedError( # pragma: no cover 179 | "Class is not designed to setup storage") # pragma: no cover 180 | 181 | @property 182 | def header_data(self): 183 | return self._root_device.header_data 184 | 185 | @property 186 | def bucket_count(self): 187 | return self._root_device.bucket_count 188 | 189 | @property 190 | def bucket_size(self): 191 | return self._root_device.bucket_size 192 | 193 | @property 194 | def blocks_per_bucket(self): 195 | return self._root_device.blocks_per_bucket 196 | 197 | @property 198 | def storage_name(self): 199 | return self._root_device.storage_name 200 | 201 | @property 202 | def virtual_heap(self): 203 | return self._root_device.virtual_heap 204 | 205 | @property 206 | def bucket_storage(self): 207 | return self._root_device.bucket_storage 208 | 209 | def update_header_data(self, new_header_data): 210 | self._root_device.update_header_data(new_header_data) 211 | 212 | def close(self): 213 | log.info("%s: Uploading %s cached bucket data before closing" 214 | % (self.__class__.__name__, self._cached_bucket_count)) 215 | with tqdm.tqdm(desc=("Uploading %s Cached Heap Buckets" 216 | % (self._cached_bucket_count)), 217 | total=self._cached_bucket_count*self.bucket_size, 218 | unit="B", 219 | unit_scale=True, 220 | disable=not pyoram.config.SHOW_PROGRESS_BAR) as progress_bar: 221 | self.bucket_storage.\ 222 | write_blocks( 223 | xrange(self._cached_bucket_count), 224 | (self._cached_buckets_mmap[(b*self.bucket_size): 225 | ((b+1)*self.bucket_size)] 226 | for b in xrange(self._cached_bucket_count)), 227 | callback=lambda i: progress_bar.update(self._root_device.bucket_size)) 228 | for b in self._concurrent_devices: 229 | self._concurrent_devices[b].close() 230 | self._root_device.close() 231 | # forces the bar to become full at close 232 | # even if te write_blocks action was faster 233 | # the the mininterval time 234 | progress_bar.mininterval = 0 235 | 236 | self._cached_buckets_mmap.close() 237 | self._cached_buckets_tempfile.close() 238 | 239 | def read_path(self, b, level_start=0): 240 | assert 0 <= b < self.virtual_heap.bucket_count() 241 | bucket_list = self.virtual_heap.Node(b).bucket_path_from_root() 242 | if len(bucket_list) <= self._external_level: 243 | return [self._cached_buckets_mmap[(bb*self.bucket_size): 244 | ((bb+1)*self.bucket_size)] 245 | for bb in bucket_list[level_start:]] 246 | elif level_start >= self._external_level: 247 | return self._subheap_storage[bucket_list[self._external_level]].\ 248 | bucket_storage.read_blocks(bucket_list[level_start:]) 249 | else: 250 | local_buckets = bucket_list[:self._external_level] 251 | external_buckets = bucket_list[self._external_level:] 252 | buckets = [] 253 | for bb in local_buckets[level_start:]: 254 | buckets.append( 255 | self._cached_buckets_mmap[(bb*self.bucket_size): 256 | ((bb+1)*self.bucket_size)]) 257 | if len(external_buckets) > 0: 258 | buckets.extend( 259 | self._subheap_storage[external_buckets[0]].\ 260 | bucket_storage.read_blocks(external_buckets)) 261 | assert len(buckets) == len(bucket_list[level_start:]) 262 | return buckets 263 | 264 | def write_path(self, b, buckets, level_start=0): 265 | assert 0 <= b < self.virtual_heap.bucket_count() 266 | bucket_list = self.virtual_heap.Node(b).bucket_path_from_root() 267 | if len(bucket_list) <= self._external_level: 268 | for bb, bucket in zip(bucket_list[level_start:], buckets): 269 | self._cached_buckets_mmap[(bb*self.bucket_size): 270 | ((bb+1)*self.bucket_size)] = bucket 271 | elif level_start >= self._external_level: 272 | self._subheap_storage[bucket_list[self._external_level]].\ 273 | bucket_storage.write_blocks(bucket_list[level_start:], buckets) 274 | else: 275 | buckets = list(buckets) 276 | assert len(buckets) == len(bucket_list[level_start:]) 277 | local_buckets = bucket_list[:self._external_level] 278 | external_buckets = bucket_list[self._external_level:] 279 | ndx = -1 280 | for ndx, bb in enumerate(local_buckets[level_start:]): 281 | self._cached_buckets_mmap[(bb*self.bucket_size): 282 | ((bb+1)*self.bucket_size)] = buckets[ndx] 283 | if len(external_buckets) > 0: 284 | self._subheap_storage[external_buckets[0]].\ 285 | bucket_storage.write_blocks(external_buckets, 286 | buckets[(ndx+1):]) 287 | @property 288 | def bytes_sent(self): 289 | return sum(device.bytes_sent for device 290 | in self._concurrent_devices.values()) 291 | 292 | @property 293 | def bytes_received(self): 294 | return sum(device.bytes_received for device 295 | in self._concurrent_devices.values()) 296 | -------------------------------------------------------------------------------- /src/pyoram/oblivious_storage/__init__.py: -------------------------------------------------------------------------------- 1 | import pyoram.oblivious_storage.tree 2 | -------------------------------------------------------------------------------- /src/pyoram/oblivious_storage/tree/__init__.py: -------------------------------------------------------------------------------- 1 | import pyoram.oblivious_storage.tree.tree_oram_helper 2 | import pyoram.oblivious_storage.tree.path_oram 3 | -------------------------------------------------------------------------------- /src/pyoram/oblivious_storage/tree/tree_oram_helper.py: -------------------------------------------------------------------------------- 1 | __all__ = ('TreeORAMStorageManagerExplicitAddressing', 2 | 'TreeORAMStorageManagerPointerAddressing') 3 | 4 | import struct 5 | import copy 6 | 7 | from pyoram.util.virtual_heap import \ 8 | SizedVirtualHeap 9 | 10 | from six.moves import xrange 11 | 12 | class TreeORAMStorage(object): 13 | 14 | empty_block_id = -1 15 | 16 | block_status_storage_string = "!?" 17 | block_id_storage_string = "!L" 18 | block_info_storage_string = "!?L" 19 | 20 | block_status_storage_size = \ 21 | struct.calcsize(block_status_storage_string) 22 | block_info_storage_size = \ 23 | struct.calcsize(block_info_storage_string) 24 | 25 | empty_block_bytes_tag = \ 26 | struct.pack(block_status_storage_string, False) 27 | real_block_bytes_tag = \ 28 | struct.pack(block_status_storage_string, True) 29 | 30 | def __init__(self, 31 | storage_heap, 32 | stash): 33 | self.storage_heap = storage_heap 34 | self.stash = stash 35 | 36 | vheap = self.storage_heap.virtual_heap 37 | self.bucket_size = self.storage_heap.bucket_size 38 | self.block_size = self.bucket_size // vheap.blocks_per_bucket 39 | assert self.block_size * vheap.blocks_per_bucket == \ 40 | self.bucket_size 41 | 42 | self.path_stop_bucket = None 43 | self.path_bucket_count = 0 44 | self.path_byte_dataview = \ 45 | bytearray(self.bucket_size * vheap.levels) 46 | dataview = memoryview(self.path_byte_dataview) 47 | self.path_bucket_dataview = \ 48 | [dataview[(i*self.bucket_size):((i+1)*self.bucket_size)] 49 | for i in xrange(vheap.levels)] 50 | 51 | self.path_block_dataview = [] 52 | for i in xrange(vheap.levels): 53 | bucketview = self.path_bucket_dataview[i] 54 | for j in xrange(vheap.blocks_per_bucket): 55 | self.path_block_dataview.append( 56 | bucketview[(j*self.block_size):((j+1)*self.block_size)]) 57 | 58 | max_blocks_on_path = vheap.levels * vheap.blocks_per_bucket 59 | assert len(self.path_block_dataview) == max_blocks_on_path 60 | self.path_block_ids = [-1] * max_blocks_on_path 61 | self.path_block_eviction_levels = [None] * max_blocks_on_path 62 | self.path_block_reordering = [None] * max_blocks_on_path 63 | self.path_blocks_inserted = [] 64 | 65 | def load_path(self, b): 66 | vheap = self.storage_heap.virtual_heap 67 | Z = vheap.blocks_per_bucket 68 | lcl = vheap.clib.calculate_last_common_level 69 | k = vheap.k 70 | 71 | read_level_start = 0 72 | if self.path_stop_bucket is not None: 73 | # don't download the root and any other buckets 74 | # that are common between the previous bucket path 75 | # and the new one 76 | read_level_start = lcl(k, self.path_stop_bucket, b) 77 | assert 0 <= b < vheap.bucket_count() 78 | self.path_stop_bucket = b 79 | new_buckets = self.storage_heap.read_path( 80 | self.path_stop_bucket, 81 | level_start=read_level_start) 82 | 83 | self.path_bucket_count = read_level_start + len(new_buckets) 84 | pos = 0 85 | for i in xrange(self.path_bucket_count): 86 | if i >= read_level_start: 87 | self.path_bucket_dataview[i][:] = \ 88 | new_buckets[i-read_level_start][:] 89 | for j in xrange(Z): 90 | block_id, block_addr = \ 91 | self.get_block_info(self.path_block_dataview[pos]) 92 | self.path_block_ids[pos] = block_id 93 | if block_id != self.empty_block_id: 94 | self.path_block_eviction_levels[pos] = \ 95 | lcl(k, self.path_stop_bucket, block_addr) 96 | else: 97 | self.path_block_eviction_levels[pos] = None 98 | self.path_block_reordering[pos] = None 99 | pos += 1 100 | 101 | max_blocks_on_path = vheap.levels * Z 102 | while pos != max_blocks_on_path: 103 | self.path_block_ids[pos] = None 104 | self.path_block_eviction_levels[pos] = None 105 | self.path_block_reordering[pos] = None 106 | pos += 1 107 | 108 | self.path_blocks_inserted = [] 109 | 110 | def push_down_path(self): 111 | vheap = self.storage_heap.virtual_heap 112 | Z = vheap.blocks_per_bucket 113 | 114 | bucket_count = self.path_bucket_count 115 | block_ids = self.path_block_ids 116 | block_eviction_levels = self.path_block_eviction_levels 117 | block_reordering = self.path_block_reordering 118 | def _do_swap(write_pos, read_pos): 119 | block_ids[write_pos], block_eviction_levels[write_pos] = \ 120 | block_ids[read_pos], block_eviction_levels[read_pos] 121 | block_ids[read_pos], block_eviction_levels[read_pos] = \ 122 | self.empty_block_id, None 123 | block_reordering[write_pos] = read_pos 124 | block_reordering[read_pos] = -1 125 | 126 | def _new_write_pos(current): 127 | current -= 1 128 | if current < 0: 129 | return None, None 130 | while (block_eviction_levels[current] is not None): 131 | current -= 1 132 | if current < 0: 133 | return None, None 134 | assert block_ids[current] == \ 135 | self.empty_block_id 136 | return current, current // Z 137 | 138 | def _new_read_pos(current): 139 | current -= 1 140 | if current < 0: 141 | return None 142 | while (block_eviction_levels[current] is None): 143 | current -= 1 144 | if current < 0: 145 | return None 146 | assert block_ids[current] != \ 147 | self.empty_block_id 148 | return current 149 | 150 | write_pos, write_level = _new_write_pos(bucket_count * Z) 151 | while write_pos is not None: 152 | read_pos = _new_read_pos(write_pos) 153 | if read_pos is None: 154 | break 155 | while ((read_pos // Z) == write_level) or \ 156 | (write_level > block_eviction_levels[read_pos]): 157 | read_pos = _new_read_pos(read_pos) 158 | if read_pos is None: 159 | break 160 | if read_pos is not None: 161 | _do_swap(write_pos, read_pos) 162 | else: 163 | # Jump directly to the start of this 164 | # bucket. There is not point in checking 165 | # for other empty slots because no blocks 166 | # can be evicted to this level. 167 | write_pos = Z * (write_pos//Z) 168 | write_pos, write_level = _new_write_pos(write_pos) 169 | 170 | def fill_path_from_stash(self): 171 | vheap = self.storage_heap.virtual_heap 172 | lcl = vheap.clib.calculate_last_common_level 173 | k = vheap.k 174 | Z = vheap.blocks_per_bucket 175 | 176 | bucket_count = self.path_bucket_count 177 | stop_bucket = self.path_stop_bucket 178 | block_ids = self.path_block_ids 179 | block_eviction_levels = self.path_block_eviction_levels 180 | blocks_inserted = self.path_blocks_inserted 181 | 182 | stash_eviction_levels = {} 183 | largest_write_position = (bucket_count * Z) - 1 184 | for write_pos in xrange(largest_write_position,-1,-1): 185 | write_level = write_pos // Z 186 | if block_ids[write_pos] == self.empty_block_id: 187 | del_id = None 188 | for id_ in self.stash: 189 | if id_ not in stash_eviction_levels: 190 | block_id, block_addr = \ 191 | self.get_block_info(self.stash[id_]) 192 | assert id_ != self.empty_block_id 193 | eviction_level = stash_eviction_levels[id_] = \ 194 | lcl(k, stop_bucket, block_addr) 195 | else: 196 | eviction_level = stash_eviction_levels[id_] 197 | if write_level <= eviction_level: 198 | block_ids[write_pos] = id_ 199 | block_eviction_levels[write_pos] = \ 200 | eviction_level 201 | blocks_inserted.append( 202 | (write_pos, self.stash[id_])) 203 | del_id = id_ 204 | break 205 | if del_id is not None: 206 | del self.stash[del_id] 207 | 208 | def evict_path(self): 209 | vheap = self.storage_heap.virtual_heap 210 | Z = vheap.blocks_per_bucket 211 | 212 | bucket_count = self.path_bucket_count 213 | stop_bucket = self.path_stop_bucket 214 | bucket_dataview = self.path_bucket_dataview 215 | block_dataview = self.path_block_dataview 216 | block_reordering = self.path_block_reordering 217 | blocks_inserted = self.path_blocks_inserted 218 | 219 | for i, read_pos in enumerate( 220 | reversed(block_reordering)): 221 | if (read_pos is not None) and \ 222 | (read_pos != -1): 223 | write_pos = len(block_reordering) - 1 - i 224 | block_dataview[write_pos][:] = block_dataview[read_pos][:] 225 | 226 | for write_pos, read_pos in enumerate(block_reordering): 227 | if read_pos == -1: 228 | self.tag_block_as_empty(block_dataview[write_pos]) 229 | 230 | for write_pos, block in blocks_inserted: 231 | block_dataview[write_pos][:] = block[:] 232 | 233 | self.storage_heap.write_path( 234 | stop_bucket, 235 | (bucket_dataview[i].tobytes() 236 | for i in xrange(bucket_count))) 237 | 238 | def extract_block_from_path(self, id_): 239 | block_ids = self.path_block_ids 240 | block_dataview = self.path_block_dataview 241 | try: 242 | pos = block_ids.index(id_) 243 | # make a copy 244 | block = bytearray(block_dataview[pos]) 245 | self._set_path_position_to_empty(pos) 246 | return block 247 | except ValueError: 248 | return None 249 | 250 | def _set_path_position_to_empty(self, pos): 251 | self.path_block_ids[pos] = self.empty_block_id 252 | self.path_block_eviction_levels[pos] = None 253 | self.path_block_reordering[pos] = -1 254 | 255 | @staticmethod 256 | def tag_block_as_empty(block): 257 | block[:TreeORAMStorage.block_status_storage_size] = \ 258 | TreeORAMStorage.empty_block_bytes_tag[:] 259 | 260 | @staticmethod 261 | def tag_block_with_id(block, id_): 262 | assert id_ >= 0 263 | struct.pack_into(TreeORAMStorage.block_info_storage_string, 264 | block, 265 | 0, 266 | True, 267 | id_) 268 | 269 | def get_block_info(self, block): 270 | raise NotImplementedError # pragma: no cover 271 | 272 | class TreeORAMStorageManagerExplicitAddressing( 273 | TreeORAMStorage): 274 | """ 275 | This class should be used to implement tree-based ORAMs 276 | that use an explicit position map. Blocks are assumed to 277 | begin with bytes representing the block id. 278 | """ 279 | 280 | block_info_storage_string = \ 281 | TreeORAMStorage.block_info_storage_string 282 | block_info_storage_size = \ 283 | struct.calcsize(block_info_storage_string) 284 | 285 | def __init__(self, 286 | storage_heap, 287 | stash, 288 | position_map): 289 | super(TreeORAMStorageManagerExplicitAddressing, self).\ 290 | __init__(storage_heap, stash) 291 | self.position_map = position_map 292 | 293 | def get_block_info(self, block): 294 | real, id_ = struct.unpack_from( 295 | self.block_info_storage_string, block) 296 | if real: 297 | return id_, self.position_map[id_] 298 | else: 299 | return self.empty_block_id, None 300 | 301 | class TreeORAMStorageManagerPointerAddressing( 302 | TreeORAMStorage): 303 | """ 304 | This class should be used to implement tree-based ORAMs 305 | that use a pointer-based position map stored with the 306 | blocks. Blocks are assumed to begin with bytes 307 | representing the block id followed by bytes representing 308 | the blocks current heap bucket address. 309 | """ 310 | 311 | block_info_storage_string = \ 312 | TreeORAMStorage.block_info_storage_string + "L" 313 | block_info_storage_size = \ 314 | struct.calcsize(block_info_storage_string) 315 | 316 | def __init__(self, 317 | storage_heap, 318 | stash): 319 | super(TreeORAMStorageManagerPointerAddressing, self).\ 320 | __init__(storage_heap, stash) 321 | self.position_map = None 322 | 323 | def get_block_info(self, block): 324 | real, id_, addr = struct.unpack_from( 325 | self.block_info_storage_string, block) 326 | if not real: 327 | return self.empty_block_id, 0 328 | else: 329 | return id_, addr 330 | -------------------------------------------------------------------------------- /src/pyoram/storage/__init__.py: -------------------------------------------------------------------------------- 1 | import pyoram.storage.block_storage 2 | import pyoram.storage.block_storage_file 3 | import pyoram.storage.block_storage_mmap 4 | import pyoram.storage.block_storage_ram 5 | import pyoram.storage.block_storage_sftp 6 | import pyoram.storage.block_storage_s3 7 | import pyoram.storage.heap_storage 8 | -------------------------------------------------------------------------------- /src/pyoram/storage/block_storage.py: -------------------------------------------------------------------------------- 1 | __all__ = ('BlockStorageTypeFactory',) 2 | 3 | import logging 4 | 5 | log = logging.getLogger("pyoram") 6 | 7 | def BlockStorageTypeFactory(storage_type_name): 8 | if storage_type_name in BlockStorageTypeFactory._registered_devices: 9 | return BlockStorageTypeFactory.\ 10 | _registered_devices[storage_type_name] 11 | else: 12 | raise ValueError( 13 | "BlockStorageTypeFactory: Unsupported storage " 14 | "type: %s" % (storage_type_name)) 15 | BlockStorageTypeFactory._registered_devices = {} 16 | 17 | def _register_device(name, type_): 18 | if name in BlockStorageTypeFactory._registered_devices: 19 | raise ValueError("Can not register block storage device type " 20 | "with name '%s'. A device type is already " 21 | "registered with that name." % (name)) 22 | if not issubclass(type_, BlockStorageInterface): 23 | raise TypeError("Can not register block storage device type " 24 | "'%s'. The device must be a subclass of " 25 | "BlockStorageInterface" % (type_)) 26 | BlockStorageTypeFactory._registered_devices[name] = type_ 27 | BlockStorageTypeFactory.register_device = _register_device 28 | 29 | class BlockStorageInterface(object): 30 | 31 | def __enter__(self): 32 | return self 33 | def __exit__(self, *args): 34 | self.close() 35 | 36 | # 37 | # Abstract Interface 38 | # 39 | 40 | def clone_device(self, *args, **kwds): 41 | raise NotImplementedError # pragma: no cover 42 | 43 | @classmethod 44 | def compute_storage_size(cls, *args, **kwds): 45 | raise NotImplementedError # pragma: no cover 46 | @classmethod 47 | def setup(cls, *args, **kwds): 48 | raise NotImplementedError # pragma: no cover 49 | 50 | @property 51 | def header_data(self, *args, **kwds): 52 | raise NotImplementedError # pragma: no cover 53 | @property 54 | def block_count(self, *args, **kwds): 55 | raise NotImplementedError # pragma: no cover 56 | @property 57 | def block_size(self, *args, **kwds): 58 | raise NotImplementedError # pragma: no cover 59 | @property 60 | def storage_name(self, *args, **kwds): 61 | raise NotImplementedError # pragma: no cover 62 | 63 | def update_header_data(self, *args, **kwds): 64 | raise NotImplementedError # pragma: no cover 65 | def close(self, *args, **kwds): 66 | raise NotImplementedError # pragma: no cover 67 | def read_blocks(self, *args, **kwds): 68 | raise NotImplementedError # pragma: no cover 69 | def yield_blocks(self, *args, **kwds): 70 | raise NotImplementedError # pragma: no cover 71 | def read_block(self, *args, **kwds): 72 | raise NotImplementedError # pragma: no cover 73 | def write_blocks(self, *args, **kwds): 74 | raise NotImplementedError # pragma: no cover 75 | def write_block(self, *args, **kwds): 76 | raise NotImplementedError # pragma: no cover 77 | 78 | @property 79 | def bytes_sent(self): 80 | raise NotImplementedError # pragma: no cover 81 | @property 82 | def bytes_received(self): 83 | raise NotImplementedError # pragma: no cover 84 | -------------------------------------------------------------------------------- /src/pyoram/storage/block_storage_file.py: -------------------------------------------------------------------------------- 1 | __all__ = ('BlockStorageFile',) 2 | 3 | import os 4 | import struct 5 | import logging 6 | import errno 7 | from multiprocessing.pool import ThreadPool 8 | 9 | import pyoram 10 | from pyoram.storage.block_storage import \ 11 | (BlockStorageInterface, 12 | BlockStorageTypeFactory) 13 | 14 | import tqdm 15 | import six 16 | from six.moves import xrange 17 | 18 | log = logging.getLogger("pyoram") 19 | 20 | class default_filesystem(object): 21 | open = open 22 | remove = os.remove 23 | stat = os.stat 24 | 25 | class BlockStorageFile(BlockStorageInterface): 26 | """ 27 | A class implementing the block storage interface 28 | using a local file. 29 | """ 30 | 31 | _index_struct_string = "!LLL?" 32 | _index_offset = struct.calcsize(_index_struct_string) 33 | 34 | def __init__(self, 35 | storage_name, 36 | threadpool_size=None, 37 | ignore_lock=False, 38 | _filesystem=default_filesystem): 39 | 40 | self._bytes_sent = 0 41 | self._bytes_received = 0 42 | self._filesystem = _filesystem 43 | self._ignore_lock = ignore_lock 44 | self._f = None 45 | self._pool = None 46 | self._close_pool = True 47 | self._async_write = None 48 | self._storage_name = storage_name 49 | self._f = self._filesystem.open(self.storage_name, "r+b") 50 | self._f.seek(0) 51 | self._block_size, self._block_count, user_header_size, locked = \ 52 | struct.unpack( 53 | BlockStorageFile._index_struct_string, 54 | self._f.read(BlockStorageFile._index_offset)) 55 | 56 | if locked and (not self._ignore_lock): 57 | self._f.close() 58 | self._f = None 59 | raise IOError( 60 | "Can not open block storage device because it is " 61 | "locked by another process. To ignore this check, " 62 | "initialize this class with the keyword 'ignore_lock' " 63 | "set to True.") 64 | self._user_header_data = bytes() 65 | if user_header_size > 0: 66 | self._user_header_data = \ 67 | self._f.read(user_header_size) 68 | self._header_offset = BlockStorageFile._index_offset + \ 69 | len(self._user_header_data) 70 | 71 | # TODO: Figure out why this is required for Python3 72 | # in order to prevent issues with the 73 | # TopCachedEncryptedHeapStorage class. The 74 | # problem has something to do with bufferedio, 75 | # but it makes no sense why this fixes it (all 76 | # we've done is read above these lines). As 77 | # part of this, investigate whethor or not we 78 | # need the call to flush after write_block(s), 79 | # or if its simply connected to some Python3 80 | # bug in bufferedio. 81 | self._f.flush() 82 | 83 | if not self._ignore_lock: 84 | # turn on the locked flag 85 | self._f.seek(0) 86 | self._f.write( 87 | struct.pack(BlockStorageFile._index_struct_string, 88 | self.block_size, 89 | self.block_count, 90 | len(self._user_header_data), 91 | True)) 92 | self._f.flush() 93 | 94 | if threadpool_size != 0: 95 | self._pool = ThreadPool(threadpool_size) 96 | 97 | def _check_async(self): 98 | if self._async_write is not None: 99 | self._async_write.get() 100 | self._async_write = None 101 | # TODO: Figure out why tests fail on Python3 without this 102 | if six.PY3: 103 | if self._f is None: 104 | return 105 | self._f.flush() 106 | 107 | def _schedule_async_write(self, args, callback=None): 108 | assert self._async_write is None 109 | if self._pool is not None: 110 | self._async_write = \ 111 | self._pool.apply_async(self._writev, (args, callback)) 112 | else: 113 | self._writev(args, callback) 114 | 115 | # This method is usually executed in another thread, so 116 | # do not attempt to handle exceptions because it will 117 | # not work. 118 | def _writev(self, chunks, callback): 119 | for i, block in chunks: 120 | self._f.seek(self._header_offset + i * self.block_size) 121 | self._f.write(block) 122 | if callback is not None: 123 | callback(i) 124 | 125 | def _prep_for_close(self): 126 | self._check_async() 127 | if self._close_pool and (self._pool is not None): 128 | self._pool.close() 129 | self._pool.join() 130 | self._pool = None 131 | if self._f is not None: 132 | if not self._ignore_lock: 133 | # turn off the locked flag 134 | self._f.seek(0) 135 | self._f.write( 136 | struct.pack(BlockStorageFile._index_struct_string, 137 | self.block_size, 138 | self.block_count, 139 | len(self._user_header_data), 140 | False)) 141 | self._f.flush() 142 | 143 | # 144 | # Define BlockStorageInterface Methods 145 | # 146 | 147 | def clone_device(self): 148 | f = BlockStorageFile(self.storage_name, 149 | threadpool_size=0, 150 | ignore_lock=True) 151 | f._pool = self._pool 152 | f._close_pool = False 153 | return f 154 | 155 | @classmethod 156 | def compute_storage_size(cls, 157 | block_size, 158 | block_count, 159 | header_data=None, 160 | ignore_header=False): 161 | assert (block_size > 0) and (block_size == int(block_size)) 162 | assert (block_count > 0) and (block_count == int(block_count)) 163 | if header_data is None: 164 | header_data = bytes() 165 | if ignore_header: 166 | return block_size * block_count 167 | else: 168 | return BlockStorageFile._index_offset + \ 169 | len(header_data) + \ 170 | block_size * block_count 171 | 172 | @classmethod 173 | def setup(cls, 174 | storage_name, 175 | block_size, 176 | block_count, 177 | initialize=None, 178 | header_data=None, 179 | ignore_existing=False, 180 | threadpool_size=None, 181 | _filesystem=default_filesystem): 182 | 183 | if (not ignore_existing): 184 | _exists = True 185 | try: 186 | _filesystem.stat(storage_name) 187 | except OSError as e: 188 | if e.errno == errno.ENOENT: 189 | _exists = False 190 | if _exists: 191 | raise IOError( 192 | "Storage location already exists: %s" 193 | % (storage_name)) 194 | if (block_size <= 0) or (block_size != int(block_size)): 195 | raise ValueError( 196 | "Block size (bytes) must be a positive integer: %s" 197 | % (block_size)) 198 | if (block_count <= 0) or (block_count != int(block_count)): 199 | raise ValueError( 200 | "Block count must be a positive integer: %s" 201 | % (block_count)) 202 | if (header_data is not None) and \ 203 | (type(header_data) is not bytes): 204 | raise TypeError( 205 | "'header_data' must be of type bytes. " 206 | "Invalid type: %s" % (type(header_data))) 207 | 208 | if initialize is None: 209 | zeros = bytes(bytearray(block_size)) 210 | initialize = lambda i: zeros 211 | try: 212 | with _filesystem.open(storage_name, "wb") as f: 213 | # create_index 214 | if header_data is None: 215 | f.write(struct.pack(BlockStorageFile._index_struct_string, 216 | block_size, 217 | block_count, 218 | 0, 219 | False)) 220 | else: 221 | f.write(struct.pack(BlockStorageFile._index_struct_string, 222 | block_size, 223 | block_count, 224 | len(header_data), 225 | False)) 226 | f.write(header_data) 227 | with tqdm.tqdm(total=block_count*block_size, 228 | desc="Initializing File Block Storage Space", 229 | unit="B", 230 | unit_scale=True, 231 | disable=not pyoram.config.SHOW_PROGRESS_BAR) as progress_bar: 232 | for i in xrange(block_count): 233 | block = initialize(i) 234 | assert len(block) == block_size, \ 235 | ("%s != %s" % (len(block), block_size)) 236 | f.write(block) 237 | progress_bar.update(n=block_size) 238 | except: # pragma: no cover 239 | _filesystem.remove(storage_name) # pragma: no cover 240 | raise # pragma: no cover 241 | 242 | return BlockStorageFile(storage_name, 243 | threadpool_size=threadpool_size, 244 | _filesystem=_filesystem) 245 | 246 | @property 247 | def header_data(self): 248 | return self._user_header_data 249 | 250 | @property 251 | def block_count(self): 252 | return self._block_count 253 | 254 | @property 255 | def block_size(self): 256 | return self._block_size 257 | 258 | @property 259 | def storage_name(self): 260 | return self._storage_name 261 | 262 | def update_header_data(self, new_header_data): 263 | self._check_async() 264 | if len(new_header_data) != len(self.header_data): 265 | raise ValueError( 266 | "The size of header data can not change.\n" 267 | "Original bytes: %s\n" 268 | "New bytes: %s" % (len(self.header_data), 269 | len(new_header_data))) 270 | self._user_header_data = bytes(new_header_data) 271 | self._f.seek(BlockStorageFile._index_offset) 272 | self._f.write(self._user_header_data) 273 | 274 | def close(self): 275 | self._prep_for_close() 276 | if self._f is not None: 277 | try: 278 | self._f.close() 279 | except OSError: # pragma: no cover 280 | pass # pragma: no cover 281 | self._f = None 282 | 283 | def read_blocks(self, indices): 284 | self._check_async() 285 | blocks = [] 286 | for i in indices: 287 | assert 0 <= i < self.block_count 288 | self._bytes_received += self.block_size 289 | self._f.seek(self._header_offset + i * self.block_size) 290 | blocks.append(self._f.read(self.block_size)) 291 | return blocks 292 | 293 | def yield_blocks(self, indices): 294 | self._check_async() 295 | for i in indices: 296 | assert 0 <= i < self.block_count 297 | self._bytes_received += self.block_size 298 | self._f.seek(self._header_offset + i * self.block_size) 299 | yield self._f.read(self.block_size) 300 | 301 | def read_block(self, i): 302 | self._check_async() 303 | assert 0 <= i < self.block_count 304 | self._bytes_received += self.block_size 305 | self._f.seek(self._header_offset + i * self.block_size) 306 | return self._f.read(self.block_size) 307 | 308 | def write_blocks(self, indices, blocks, callback=None): 309 | self._check_async() 310 | chunks = [] 311 | for i, block in zip(indices, blocks): 312 | assert 0 <= i < self.block_count 313 | assert len(block) == self.block_size, \ 314 | ("%s != %s" % (len(block), self.block_size)) 315 | self._bytes_sent += self.block_size 316 | chunks.append((i, block)) 317 | self._schedule_async_write(chunks, callback=callback) 318 | 319 | def write_block(self, i, block): 320 | self._check_async() 321 | assert 0 <= i < self.block_count 322 | assert len(block) == self.block_size 323 | self._bytes_sent += self.block_size 324 | self._schedule_async_write(((i, block),)) 325 | 326 | @property 327 | def bytes_sent(self): 328 | return self._bytes_sent 329 | 330 | @property 331 | def bytes_received(self): 332 | return self._bytes_received 333 | 334 | BlockStorageTypeFactory.register_device("file", BlockStorageFile) 335 | -------------------------------------------------------------------------------- /src/pyoram/storage/block_storage_mmap.py: -------------------------------------------------------------------------------- 1 | __all__ = ('BlockStorageMMap',) 2 | 3 | import logging 4 | import mmap 5 | 6 | from pyoram.storage.block_storage import \ 7 | BlockStorageTypeFactory 8 | from pyoram.storage.block_storage_file import \ 9 | BlockStorageFile 10 | 11 | log = logging.getLogger("pyoram") 12 | 13 | class _BlockStorageMemoryImpl(object): 14 | """ 15 | This class implementents the BlockStorageInterface read/write 16 | methods for classes with a private attribute _f that can be 17 | accessed using __getslice__/__setslice__ notation. 18 | """ 19 | 20 | def read_blocks(self, indices): 21 | blocks = [] 22 | for i in indices: 23 | assert 0 <= i < self.block_count 24 | self._bytes_received += self.block_size 25 | pos_start = self._header_offset + i * self.block_size 26 | pos_stop = pos_start + self.block_size 27 | blocks.append(self._f[pos_start:pos_stop]) 28 | return blocks 29 | 30 | def yield_blocks(self, indices): 31 | for i in indices: 32 | assert 0 <= i < self.block_count 33 | self._bytes_received += self.block_size 34 | pos_start = self._header_offset + i * self.block_size 35 | pos_stop = pos_start + self.block_size 36 | yield self._f[pos_start:pos_stop] 37 | 38 | def read_block(self, i): 39 | assert 0 <= i < self.block_count 40 | self._bytes_received += self.block_size 41 | pos_start = self._header_offset + i * self.block_size 42 | pos_stop = pos_start + self.block_size 43 | return self._f[pos_start:pos_stop] 44 | 45 | def write_blocks(self, indices, blocks, callback=None): 46 | for i, block in zip(indices, blocks): 47 | assert 0 <= i < self.block_count 48 | self._bytes_sent += self.block_size 49 | pos_start = self._header_offset + i * self.block_size 50 | pos_stop = pos_start + self.block_size 51 | self._f[pos_start:pos_stop] = block 52 | if callback is not None: 53 | callback(i) 54 | 55 | def write_block(self, i, block): 56 | assert 0 <= i < self.block_count 57 | self._bytes_sent += self.block_size 58 | pos_start = self._header_offset + i * self.block_size 59 | pos_stop = pos_start + self.block_size 60 | self._f[pos_start:pos_stop] = block 61 | 62 | class BlockStorageMMap(_BlockStorageMemoryImpl, 63 | BlockStorageFile): 64 | """ 65 | A class implementing the block storage interface by creating a 66 | memory map over a local file. This class uses the same storage 67 | format as BlockStorageFile. Thus, a block storage space can be 68 | created using this class and then, after saving the raw storage 69 | data to disk, reopened with any other class compatible with 70 | BlockStorageFile (and visa versa). 71 | """ 72 | 73 | def __init__(self, *args, **kwds): 74 | mm = kwds.pop('mm', None) 75 | self._mmap_owned = True 76 | super(BlockStorageMMap, self).__init__(*args, **kwds) 77 | if mm is None: 78 | self._f.flush() 79 | mm = mmap.mmap(self._f.fileno(), 0) 80 | else: 81 | self._mmap_owned = False 82 | self._f.close() 83 | self._f = mm 84 | 85 | # 86 | # Define BlockStorageInterface Methods 87 | # (override what is defined on BlockStorageFile) 88 | # 89 | 90 | #@classmethod 91 | #def compute_storage_size(...) 92 | 93 | def clone_device(self): 94 | f = BlockStorageMMap(self.storage_name, 95 | threadpool_size=0, 96 | mm=self._f, 97 | ignore_lock=True) 98 | f._pool = self._pool 99 | f._close_pool = False 100 | return f 101 | 102 | @classmethod 103 | def setup(cls, 104 | storage_name, 105 | block_size, 106 | block_count, 107 | **kwds): 108 | f = BlockStorageFile.setup(storage_name, 109 | block_size, 110 | block_count, 111 | **kwds) 112 | f.close() 113 | return BlockStorageMMap(storage_name) 114 | 115 | #def update_header_data(...) 116 | 117 | def close(self): 118 | self._prep_for_close() 119 | if self._f is not None: 120 | if self._mmap_owned: 121 | try: 122 | self._f.close() 123 | except OSError: # pragma: no cover 124 | pass # pragma: no cover 125 | self._f = None 126 | 127 | #def read_blocks(...) 128 | 129 | #def yield_blocks(...) 130 | 131 | #def read_block(...) 132 | 133 | #def write_blocks(...) 134 | 135 | #def write_block(...) 136 | 137 | #@property 138 | #def bytes_sent(...) 139 | 140 | #@property 141 | #def bytes_received(...) 142 | 143 | BlockStorageTypeFactory.register_device("mmap", BlockStorageMMap) 144 | -------------------------------------------------------------------------------- /src/pyoram/storage/block_storage_ram.py: -------------------------------------------------------------------------------- 1 | __all__ = ('BlockStorageRAM',) 2 | 3 | import os 4 | import struct 5 | import logging 6 | import errno 7 | from multiprocessing.pool import ThreadPool 8 | 9 | import pyoram 10 | from pyoram.storage.block_storage import \ 11 | (BlockStorageInterface, 12 | BlockStorageTypeFactory) 13 | from pyoram.storage.block_storage_mmap import \ 14 | (BlockStorageMMap, 15 | _BlockStorageMemoryImpl) 16 | 17 | import tqdm 18 | import six 19 | from six.moves import xrange 20 | 21 | log = logging.getLogger("pyoram") 22 | 23 | class BlockStorageRAM(_BlockStorageMemoryImpl, 24 | BlockStorageInterface): 25 | """ 26 | A class implementing the block storage interface where all data is 27 | kept in RAM. This class uses the same storage format as 28 | BlockStorageFile. Thus, a block storage space can be created using 29 | this class and then, after saving the raw storage data to disk, 30 | reopened with any other class compatible with BlockStorageFile 31 | (and visa versa). 32 | """ 33 | 34 | _index_struct_string = BlockStorageMMap._index_struct_string 35 | _index_offset = struct.calcsize(_index_struct_string) 36 | 37 | def __init__(self, 38 | storage_data, 39 | threadpool_size=None, 40 | ignore_lock=False): 41 | 42 | self._bytes_sent = 0 43 | self._bytes_received = 0 44 | self._ignore_lock = ignore_lock 45 | self._f = None 46 | self._pool = None 47 | self._close_pool = True 48 | if type(storage_data) is not bytearray: 49 | raise TypeError( 50 | "BlockStorageRAM requires input argument of type " 51 | "'bytearray'. Invalid input type: %s" 52 | % (type(storage_data))) 53 | self._f = storage_data 54 | self._block_size, self._block_count, user_header_size, locked = \ 55 | struct.unpack( 56 | BlockStorageRAM._index_struct_string, 57 | self._f[:BlockStorageRAM._index_offset]) 58 | 59 | if locked and (not self._ignore_lock): 60 | raise IOError( 61 | "Can not open block storage device because it is " 62 | "locked by another process. To ignore this check, " 63 | "initialize this class with the keyword 'ignore_lock' " 64 | "set to True.") 65 | self._user_header_data = bytes() 66 | if user_header_size > 0: 67 | self._user_header_data = \ 68 | bytes(self._f[BlockStorageRAM._index_offset:\ 69 | (BlockStorageRAM._index_offset+user_header_size)]) 70 | assert len(self._user_header_data) == user_header_size 71 | self._header_offset = BlockStorageRAM._index_offset + \ 72 | len(self._user_header_data) 73 | 74 | if not self._ignore_lock: 75 | # turn on the locked flag 76 | self._f[:BlockStorageRAM._index_offset] = \ 77 | struct.pack(BlockStorageRAM._index_struct_string, 78 | self.block_size, 79 | self.block_count, 80 | len(self._user_header_data), 81 | True) 82 | 83 | # Although we do not use the threadpool we still 84 | # create just in case we are the first 85 | if threadpool_size != 0: 86 | self._pool = ThreadPool(threadpool_size) 87 | 88 | # 89 | # Add some methods specific to BlockStorageRAM 90 | # 91 | 92 | @staticmethod 93 | def fromfile(file_, 94 | threadpool_size=None, 95 | ignore_lock=False): 96 | """ 97 | Instantiate BlockStorageRAM device from a file saved in block 98 | storage format. The file_ argument can be a file object or a 99 | string that represents a filename. If called with a file 100 | object, it should be opened in binary mode, and the caller is 101 | responsible for closing the file. 102 | 103 | This method returns a BlockStorageRAM instance. 104 | """ 105 | close_file = False 106 | if not hasattr(file_, 'read'): 107 | file_ = open(file_, 'rb') 108 | close_file = True 109 | try: 110 | header_data = file_.read(BlockStorageRAM._index_offset) 111 | block_size, block_count, user_header_size, locked = \ 112 | struct.unpack( 113 | BlockStorageRAM._index_struct_string, 114 | header_data) 115 | if locked and (not ignore_lock): 116 | raise IOError( 117 | "Can not open block storage device because it is " 118 | "locked by another process. To ignore this check, " 119 | "call this method with the keyword 'ignore_lock' " 120 | "set to True.") 121 | header_offset = len(header_data) + \ 122 | user_header_size 123 | f = bytearray(header_offset + \ 124 | (block_size * block_count)) 125 | f[:header_offset] = header_data + file_.read(user_header_size) 126 | f[header_offset:] = file_.read(block_size * block_count) 127 | finally: 128 | if close_file: 129 | file_.close() 130 | 131 | return BlockStorageRAM(f, 132 | threadpool_size=threadpool_size, 133 | ignore_lock=ignore_lock) 134 | 135 | def tofile(self, file_): 136 | """ 137 | Dump all storage data to a file. The file_ argument can be a 138 | file object or a string that represents a filename. If called 139 | with a file object, it should be opened in binary mode, and 140 | the caller is responsible for closing the file. 141 | 142 | The method should only be called after the storage device has 143 | been closed to ensure that the locked flag has been set to 144 | False. 145 | """ 146 | close_file = False 147 | if not hasattr(file_, 'write'): 148 | file_ = open(file_, 'wb') 149 | close_file = True 150 | file_.write(self._f) 151 | if close_file: 152 | file_.close() 153 | 154 | @property 155 | def data(self): 156 | """Access the raw bytearray""" 157 | return self._f 158 | 159 | # 160 | # Define BlockStorageInterface Methods 161 | # 162 | 163 | def clone_device(self): 164 | f = BlockStorageRAM(self._f, 165 | threadpool_size=0, 166 | ignore_lock=True) 167 | f._pool = self._pool 168 | f._close_pool = False 169 | return f 170 | 171 | @classmethod 172 | def compute_storage_size(cls, *args, **kwds): 173 | return BlockStorageMMap.compute_storage_size(*args, **kwds) 174 | 175 | @classmethod 176 | def setup(cls, 177 | storage_name, 178 | block_size, 179 | block_count, 180 | initialize=None, 181 | header_data=None, 182 | ignore_existing=False, 183 | threadpool_size=None): 184 | 185 | # We ignore the 'storage_name' argument 186 | # We ignore the 'ignore_existing' flag 187 | if (block_size <= 0) or (block_size != int(block_size)): 188 | raise ValueError( 189 | "Block size (bytes) must be a positive integer: %s" 190 | % (block_size)) 191 | if (block_count <= 0) or (block_count != int(block_count)): 192 | raise ValueError( 193 | "Block count must be a positive integer: %s" 194 | % (block_count)) 195 | if (header_data is not None) and \ 196 | (type(header_data) is not bytes): 197 | raise TypeError( 198 | "'header_data' must be of type bytes. " 199 | "Invalid type: %s" % (type(header_data))) 200 | 201 | if initialize is None: 202 | zeros = bytes(bytearray(block_size)) 203 | initialize = lambda i: zeros 204 | 205 | # create_index 206 | index_data = None 207 | if header_data is None: 208 | index_data = struct.pack(BlockStorageRAM._index_struct_string, 209 | block_size, 210 | block_count, 211 | 0, 212 | False) 213 | header_data = bytes() 214 | else: 215 | index_data = struct.pack(BlockStorageRAM._index_struct_string, 216 | block_size, 217 | block_count, 218 | len(header_data), 219 | False) 220 | header_offset = len(index_data) + len(header_data) 221 | f = bytearray(header_offset + \ 222 | (block_size * block_count)) 223 | f[:header_offset] = index_data + header_data 224 | progress_bar = tqdm.tqdm(total=block_count*block_size, 225 | desc="Initializing File Block Storage Space", 226 | unit="B", 227 | unit_scale=True, 228 | disable=not pyoram.config.SHOW_PROGRESS_BAR) 229 | for i in xrange(block_count): 230 | block = initialize(i) 231 | assert len(block) == block_size, \ 232 | ("%s != %s" % (len(block), block_size)) 233 | pos_start = header_offset + i * block_size 234 | pos_start = header_offset + i * block_size 235 | pos_stop = pos_start + block_size 236 | f[pos_start:pos_stop] = block[:] 237 | progress_bar.update(n=block_size) 238 | progress_bar.close() 239 | 240 | return BlockStorageRAM(f, threadpool_size=threadpool_size) 241 | 242 | @property 243 | def header_data(self): 244 | return self._user_header_data 245 | 246 | @property 247 | def block_count(self): 248 | return self._block_count 249 | 250 | @property 251 | def block_size(self): 252 | return self._block_size 253 | 254 | @property 255 | def storage_name(self): 256 | return None 257 | 258 | def update_header_data(self, new_header_data): 259 | if len(new_header_data) != len(self.header_data): 260 | raise ValueError( 261 | "The size of header data can not change.\n" 262 | "Original bytes: %s\n" 263 | "New bytes: %s" % (len(self.header_data), 264 | len(new_header_data))) 265 | self._user_header_data = bytes(new_header_data) 266 | self._f[BlockStorageRAM._index_offset:\ 267 | (BlockStorageRAM._index_offset+len(new_header_data))] = \ 268 | self._user_header_data 269 | 270 | def close(self): 271 | if self._close_pool and (self._pool is not None): 272 | self._pool.close() 273 | self._pool.join() 274 | self._pool = None 275 | if not self._ignore_lock: 276 | # turn off the locked flag 277 | self._f[:BlockStorageRAM._index_offset] = \ 278 | struct.pack(BlockStorageRAM._index_struct_string, 279 | self.block_size, 280 | self.block_count, 281 | len(self._user_header_data), 282 | False) 283 | self._ignore_lock = True 284 | 285 | # 286 | # We must cast from bytearray to bytes 287 | # when reading from a bytearray so that this 288 | # class works with the encryption layer. 289 | # 290 | 291 | def read_blocks(self, indices): 292 | return [bytes(block) for block 293 | in super(BlockStorageRAM, self).read_blocks(indices)] 294 | 295 | def yield_blocks(self, indices): 296 | for block in super(BlockStorageRAM, self).yield_blocks(indices): 297 | yield bytes(block) 298 | 299 | def read_block(self, i): 300 | return bytes(super(BlockStorageRAM, self).read_block(i)) 301 | 302 | #def write_blocks(...) 303 | 304 | #def write_block(...) 305 | 306 | @property 307 | def bytes_sent(self): 308 | return self._bytes_sent 309 | 310 | @property 311 | def bytes_received(self): 312 | return self._bytes_received 313 | 314 | BlockStorageTypeFactory.register_device("ram", BlockStorageRAM) 315 | -------------------------------------------------------------------------------- /src/pyoram/storage/block_storage_sftp.py: -------------------------------------------------------------------------------- 1 | __all__ = ('BlockStorageSFTP',) 2 | 3 | import logging 4 | 5 | from pyoram.util.misc import chunkiter 6 | from pyoram.storage.block_storage import \ 7 | BlockStorageTypeFactory 8 | from pyoram.storage.block_storage_file import \ 9 | BlockStorageFile 10 | 11 | log = logging.getLogger("pyoram") 12 | 13 | class BlockStorageSFTP(BlockStorageFile): 14 | """ 15 | A block storage device for accessing file data through 16 | an SSH portal using Secure File Transfer Protocol (SFTP). 17 | """ 18 | 19 | def __init__(self, 20 | storage_name, 21 | sshclient=None, 22 | **kwds): 23 | if sshclient is None: 24 | raise ValueError( 25 | "Can not open sftp block storage device " 26 | "without an ssh client.") 27 | super(BlockStorageSFTP, self).__init__( 28 | storage_name, 29 | _filesystem=sshclient.open_sftp(), 30 | **kwds) 31 | self._sshclient = sshclient 32 | self._f.set_pipelined() 33 | 34 | # 35 | # Define BlockStorageInterface Methods 36 | # 37 | 38 | def clone_device(self): 39 | f = BlockStorageSFTP(self.storage_name, 40 | sshclient=self._sshclient, 41 | threadpool_size=0, 42 | ignore_lock=True) 43 | f._pool = self._pool 44 | f._close_pool = False 45 | return f 46 | 47 | #@classmethod 48 | #def compute_storage_size(...) 49 | 50 | @classmethod 51 | def setup(cls, 52 | storage_name, 53 | block_size, 54 | block_count, 55 | sshclient=None, 56 | threadpool_size=None, 57 | **kwds): 58 | if sshclient is None: 59 | raise ValueError( 60 | "Can not setup sftp block storage device " 61 | "without an ssh client.") 62 | 63 | with BlockStorageFile.setup(storage_name, 64 | block_size, 65 | block_count, 66 | _filesystem=sshclient.open_sftp(), 67 | threadpool_size=threadpool_size, 68 | **kwds) as f: 69 | pass 70 | f._filesystem.close() 71 | 72 | return BlockStorageSFTP(storage_name, 73 | sshclient=sshclient, 74 | threadpool_size=threadpool_size) 75 | 76 | #@property 77 | #def header_data(...) 78 | 79 | #@property 80 | #def block_count(...) 81 | 82 | #@property 83 | #def block_size(...) 84 | 85 | #@property 86 | #def storage_name(...) 87 | 88 | #def update_header_data(...) 89 | 90 | def close(self): 91 | super(BlockStorageSFTP, self).close() 92 | self._filesystem.close() 93 | 94 | def read_blocks(self, indices): 95 | self._check_async() 96 | args = [] 97 | for i in indices: 98 | assert 0 <= i < self.block_count 99 | self._bytes_received += self.block_size 100 | args.append((self._header_offset + i * self.block_size, 101 | self.block_size)) 102 | return self._f.readv(args) 103 | 104 | def yield_blocks(self, indices, chunksize=100): 105 | for chunk in chunkiter(indices, n=chunksize): 106 | assert all(0 <= i <= self.block_count for i in chunk) 107 | self._bytes_received += self.block_size * len(chunk) 108 | args = [(self._header_offset + i * self.block_size, 109 | self.block_size) 110 | for i in chunk] 111 | for block in self._f.readv(args): 112 | yield block 113 | 114 | #def read_block(...) 115 | 116 | #def write_blocks(...) 117 | 118 | #def write_block(...) 119 | 120 | #@property 121 | #def bytes_sent(...) 122 | 123 | #@property 124 | #def bytes_received(...) 125 | 126 | BlockStorageTypeFactory.register_device("sftp", BlockStorageSFTP) 127 | -------------------------------------------------------------------------------- /src/pyoram/storage/boto3_s3_wrapper.py: -------------------------------------------------------------------------------- 1 | __all__ = ("Boto3S3Wrapper", 2 | "MockBoto3S3Wrapper") 3 | import os 4 | import shutil 5 | 6 | import pyoram 7 | 8 | import tqdm 9 | try: 10 | import boto3 11 | import botocore 12 | boto3_available = True 13 | except: # pragma: no cover 14 | boto3_available = False # pragma: no cover 15 | 16 | import six 17 | from six.moves import xrange, map 18 | 19 | class Boto3S3Wrapper(object): 20 | """ 21 | A wrapper class for the boto3 S3 service. 22 | """ 23 | 24 | def __init__(self, 25 | bucket_name, 26 | aws_access_key_id=None, 27 | aws_secret_access_key=None, 28 | region_name=None): 29 | if not boto3_available: 30 | raise ImportError( # pragma: no cover 31 | "boto3 module is required to " # pragma: no cover 32 | "use BlockStorageS3 device") # pragma: no cover 33 | 34 | self._s3 = boto3.session.Session( 35 | aws_access_key_id=aws_access_key_id, 36 | aws_secret_access_key=aws_secret_access_key, 37 | region_name=region_name).resource('s3') 38 | self._bucket = self._s3.Bucket(bucket_name) 39 | 40 | def exists(self, key): 41 | try: 42 | self._bucket.Object(key).load() 43 | except botocore.exceptions.ClientError as e: 44 | if e.response['Error']['Code'] == "404": 45 | pass 46 | else: 47 | raise e 48 | else: 49 | return True 50 | # It's not a file. Check if it's a "directory". 51 | for obj in self._bucket.objects.filter( 52 | Prefix=key+"/"): 53 | return True 54 | return False 55 | 56 | def download(self, key): 57 | try: 58 | return self._s3.meta.client.get_object( 59 | Bucket=self._bucket.name, 60 | Key=key)['Body'].read() 61 | except botocore.exceptions.ClientError: 62 | raise IOError("Can not download key: %s" 63 | % (key)) 64 | 65 | def upload(self, key_block): 66 | key, block = key_block 67 | self._bucket.put_object(Key=key, Body=block) 68 | 69 | # Chunk a streamed iterator of which we do not know 70 | # the size 71 | def _chunks(self, objs, n=100): 72 | assert 1 <= n <= 1000 # required by boto3 73 | objs = iter(objs) 74 | try: 75 | while (1): 76 | chunk = [] 77 | while len(chunk) < n: 78 | chunk.append({'Key': six.next(objs).key}) 79 | yield {'Objects': chunk} 80 | except StopIteration: 81 | pass 82 | if len(chunk): 83 | yield {'Objects': chunk} 84 | 85 | def _del(self, chunk): 86 | self._bucket.delete_objects(Delete=chunk) 87 | return len(chunk['Objects']) 88 | 89 | def clear(self, key, threadpool=None): 90 | objs = self._bucket.objects.filter(Prefix=key+"/") 91 | if threadpool is not None: 92 | deliter = threadpool.imap(self._del, self._chunks(objs)) 93 | else: 94 | deliter = map(self._del, self._chunks(objs)) 95 | with tqdm.tqdm(total=None, #len(objs), 96 | desc="Clearing S3 Blocks", 97 | unit=" objects", 98 | disable=not pyoram.config.SHOW_PROGRESS_BAR) as progress_bar: 99 | progress_bar.update(n=0) 100 | for chunksize in deliter: 101 | progress_bar.update(n=chunksize) 102 | 103 | class MockBoto3S3Wrapper(object): 104 | """ 105 | A mock class for Boto3S3Wrapper that uses the local filesystem and 106 | treats the bucket name as a directory. 107 | 108 | This class is mainly used for testing, but could potentially be 109 | used to setup storage locally that is then uploaded to S3 through 110 | the AWS web portal. 111 | """ 112 | 113 | def __init__(self, 114 | bucket_name, 115 | aws_access_key_id=None, 116 | aws_secret_access_key=None, 117 | region_name=None): 118 | 119 | self._bucket_name = os.path.abspath( 120 | os.path.normpath(bucket_name)) 121 | 122 | # called within upload to create directory 123 | # heirarchy on the fly 124 | def _makedirs_if_needed(self, key): 125 | if not os.path.exists( 126 | os.path.dirname( 127 | os.path.join(self._bucket_name, key))): 128 | os.makedirs( 129 | os.path.dirname( 130 | os.path.join(self._bucket_name, key))) 131 | assert not os.path.isdir( 132 | os.path.join(self._bucket_name, key)) 133 | 134 | def exists(self, key): 135 | return os.path.exists( 136 | os.path.join(self._bucket_name, key)) 137 | 138 | def download(self, key): 139 | with open(os.path.join(self._bucket_name, key), 'rb') as f: 140 | return f.read() 141 | 142 | def upload(self, key_block): 143 | key, block = key_block 144 | self._makedirs_if_needed(key) 145 | with open(os.path.join(self._bucket_name, key), 'wb') as f: 146 | f.write(block) 147 | 148 | def clear(self, key, threadpool=None): 149 | if os.path.exists( 150 | os.path.join(self._bucket_name, key)): 151 | if os.path.isdir( 152 | os.path.join(self._bucket_name, key)): 153 | shutil.rmtree( 154 | os.path.join(self._bucket_name, key), 155 | ignore_errors=True) 156 | else: 157 | os.remove( 158 | os.path.join(self._bucket_name, key)) 159 | -------------------------------------------------------------------------------- /src/pyoram/storage/heap_storage.py: -------------------------------------------------------------------------------- 1 | __all__ = ('HeapStorage',) 2 | 3 | import struct 4 | 5 | from pyoram.util.virtual_heap import SizedVirtualHeap 6 | from pyoram.storage.block_storage import (BlockStorageInterface, 7 | BlockStorageTypeFactory) 8 | 9 | class HeapStorageInterface(object): 10 | 11 | def __enter__(self): 12 | return self 13 | def __exit__(self, *args): 14 | self.close() 15 | 16 | # 17 | # Abstract Interface 18 | # 19 | 20 | def clone_device(self, *args, **kwds): 21 | raise NotImplementedError # pragma: no cover 22 | 23 | @classmethod 24 | def compute_storage_size(cls, *args, **kwds): 25 | raise NotImplementedError # pragma: no cover 26 | @classmethod 27 | def setup(cls, *args, **kwds): 28 | raise NotImplementedError # pragma: no cover 29 | 30 | @property 31 | def header_data(self, *args, **kwds): 32 | raise NotImplementedError # pragma: no cover 33 | @property 34 | def bucket_count(self, *args, **kwds): 35 | raise NotImplementedError # pragma: no cover 36 | @property 37 | def bucket_size(self, *args, **kwds): 38 | raise NotImplementedError # pragma: no cover 39 | @property 40 | def blocks_per_bucket(self, *args, **kwds): 41 | raise NotImplementedError # pragma: no cover 42 | @property 43 | def storage_name(self, *args, **kwds): 44 | raise NotImplementedError # pragma: no cover 45 | @property 46 | def virtual_heap(self, *args, **kwds): 47 | raise NotImplementedError # pragma: no cover 48 | @property 49 | def bucket_storage(self, *args, **kwds): 50 | raise NotImplementedError # pragma: no cover 51 | 52 | def update_header_data(self, *args, **kwds): 53 | raise NotImplementedError # pragma: no cover 54 | def close(self, *args, **kwds): 55 | raise NotImplementedError # pragma: no cover 56 | def read_path(self, *args, **kwds): 57 | raise NotImplementedError # pragma: no cover 58 | def write_path(self, *args, **kwds): 59 | raise NotImplementedError # pragma: no cover 60 | 61 | @property 62 | def bytes_sent(self): 63 | raise NotImplementedError # pragma: no cover 64 | @property 65 | def bytes_received(self): 66 | raise NotImplementedError # pragma: no cover 67 | 68 | class HeapStorage(HeapStorageInterface): 69 | 70 | _header_struct_string = "!LLL" 71 | _header_offset = struct.calcsize(_header_struct_string) 72 | 73 | def _new_storage(self, storage, **kwds): 74 | storage_type = kwds.pop('storage_type', 'file') 75 | 76 | 77 | def __init__(self, storage, **kwds): 78 | if isinstance(storage, BlockStorageInterface): 79 | self._storage = storage 80 | if len(kwds): 81 | raise ValueError( 82 | "Keywords not used when initializing " 83 | "with a storage device: %s" 84 | % (str(kwds))) 85 | else: 86 | storage_type = kwds.pop('storage_type', 'file') 87 | self._storage = BlockStorageTypeFactory(storage_type)\ 88 | (storage, **kwds) 89 | 90 | heap_base, heap_height, blocks_per_bucket = \ 91 | struct.unpack( 92 | self._header_struct_string, 93 | self._storage.header_data[:self._header_offset]) 94 | self._vheap = SizedVirtualHeap( 95 | heap_base, 96 | heap_height, 97 | blocks_per_bucket=blocks_per_bucket) 98 | 99 | # 100 | # Define HeapStorageInterface Methods 101 | # 102 | 103 | def clone_device(self): 104 | return HeapStorage(self._storage.clone_device()) 105 | 106 | @classmethod 107 | def compute_storage_size(cls, 108 | block_size, 109 | heap_height, 110 | blocks_per_bucket=1, 111 | heap_base=2, 112 | ignore_header=False, 113 | storage_type='file', 114 | **kwds): 115 | assert (block_size > 0) and (block_size == int(block_size)) 116 | assert heap_height >= 0 117 | assert blocks_per_bucket >= 1 118 | assert heap_base >= 2 119 | assert 'block_count' not in kwds 120 | vheap = SizedVirtualHeap( 121 | heap_base, 122 | heap_height, 123 | blocks_per_bucket=blocks_per_bucket) 124 | if ignore_header: 125 | return BlockStorageTypeFactory(storage_type).\ 126 | compute_storage_size( 127 | vheap.blocks_per_bucket * block_size, 128 | vheap.bucket_count(), 129 | ignore_header=True, 130 | **kwds) 131 | else: 132 | return cls._header_offset + \ 133 | BlockStorageTypeFactory(storage_type).\ 134 | compute_storage_size( 135 | vheap.blocks_per_bucket * block_size, 136 | vheap.bucket_count(), 137 | ignore_header=False, 138 | **kwds) 139 | 140 | @classmethod 141 | def setup(cls, 142 | storage_name, 143 | block_size, 144 | heap_height, 145 | blocks_per_bucket=1, 146 | heap_base=2, 147 | storage_type='file', 148 | **kwds): 149 | if 'block_count' in kwds: 150 | raise ValueError("'block_count' keyword is not accepted") 151 | if heap_height < 0: 152 | raise ValueError( 153 | "heap height must be 0 or greater. Invalid value: %s" 154 | % (heap_height)) 155 | if blocks_per_bucket < 1: 156 | raise ValueError( 157 | "blocks_per_bucket must be 1 or greater. " 158 | "Invalid value: %s" % (blocks_per_bucket)) 159 | if heap_base < 2: 160 | raise ValueError( 161 | "heap base must be 2 or greater. Invalid value: %s" 162 | % (heap_base)) 163 | 164 | vheap = SizedVirtualHeap( 165 | heap_base, 166 | heap_height, 167 | blocks_per_bucket=blocks_per_bucket) 168 | 169 | user_header_data = kwds.pop('header_data', bytes()) 170 | if type(user_header_data) is not bytes: 171 | raise TypeError( 172 | "'header_data' must be of type bytes. " 173 | "Invalid type: %s" % (type(user_header_data))) 174 | kwds['header_data'] = \ 175 | struct.pack(cls._header_struct_string, 176 | heap_base, 177 | heap_height, 178 | blocks_per_bucket) + \ 179 | user_header_data 180 | 181 | return HeapStorage( 182 | BlockStorageTypeFactory(storage_type).setup( 183 | storage_name, 184 | vheap.blocks_per_bucket * block_size, 185 | vheap.bucket_count(), 186 | **kwds)) 187 | 188 | @property 189 | def header_data(self): 190 | return self._storage.header_data[self._header_offset:] 191 | 192 | @property 193 | def bucket_count(self): 194 | return self._storage.block_count 195 | 196 | @property 197 | def bucket_size(self): 198 | return self._storage.block_size 199 | 200 | @property 201 | def blocks_per_bucket(self): 202 | return self._vheap.blocks_per_bucket 203 | 204 | @property 205 | def storage_name(self): 206 | return self._storage.storage_name 207 | 208 | @property 209 | def virtual_heap(self): 210 | return self._vheap 211 | 212 | @property 213 | def bucket_storage(self): 214 | return self._storage 215 | 216 | def update_header_data(self, new_header_data): 217 | self._storage.update_header_data( 218 | self._storage.header_data[:self._header_offset] + \ 219 | new_header_data) 220 | 221 | def close(self): 222 | self._storage.close() 223 | 224 | def read_path(self, b, level_start=0): 225 | assert 0 <= b < self._vheap.bucket_count() 226 | bucket_list = self._vheap.Node(b).bucket_path_from_root() 227 | assert 0 <= level_start < len(bucket_list) 228 | return self._storage.read_blocks(bucket_list[level_start:]) 229 | 230 | def write_path(self, b, buckets, level_start=0): 231 | assert 0 <= b < self._vheap.bucket_count() 232 | bucket_list = self._vheap.Node(b).bucket_path_from_root() 233 | assert 0 <= level_start < len(bucket_list) 234 | self._storage.write_blocks(bucket_list[level_start:], 235 | buckets) 236 | 237 | @property 238 | def bytes_sent(self): 239 | return self._storage.bytes_sent 240 | 241 | @property 242 | def bytes_received(self): 243 | return self._storage.bytes_received 244 | -------------------------------------------------------------------------------- /src/pyoram/tests/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ghackebeil/PyORAM/53e109dfb1ecec52348a70ddc64fae65eea7490a/src/pyoram/tests/__init__.py -------------------------------------------------------------------------------- /src/pyoram/tests/baselines/exists.empty: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ghackebeil/PyORAM/53e109dfb1ecec52348a70ddc64fae65eea7490a/src/pyoram/tests/baselines/exists.empty -------------------------------------------------------------------------------- /src/pyoram/tests/baselines/k200_h0_b1.dot: -------------------------------------------------------------------------------- 1 | // Created by SizedVirtualHeap.write_as_dot(...) 2 | digraph heaptree { 3 | node [shape=record] 4 | 0 [penwidth=1,label="(0, 0)}"]; 5 | } 6 | -------------------------------------------------------------------------------- /src/pyoram/tests/baselines/k200_h0_b1_data.dot: -------------------------------------------------------------------------------- 1 | // Created by SizedVirtualHeap.write_as_dot(...) 2 | digraph heaptree { 3 | node [shape=record] 4 | 0 [penwidth=1,label="{{0}}"]; 5 | } 6 | -------------------------------------------------------------------------------- /src/pyoram/tests/baselines/k2_h3_b1.dot: -------------------------------------------------------------------------------- 1 | // Created by SizedVirtualHeap.write_as_dot(...) 2 | digraph heaptree { 3 | node [shape=record] 4 | 0 [penwidth=1,label="''}"]; 5 | 1 [penwidth=1,label="'0'}"]; 6 | 3 [penwidth=1,label="'00'}"]; 7 | 7 [penwidth=1,label="'000'}"]; 8 | 3 -> 7 ; 9 | 8 [penwidth=1,label="'001'}"]; 10 | 3 -> 8 ; 11 | 1 -> 3 ; 12 | 4 [penwidth=1,label="'01'}"]; 13 | 9 [penwidth=1,label="'010'}"]; 14 | 4 -> 9 ; 15 | 10 [penwidth=1,label="'011'}"]; 16 | 4 -> 10 ; 17 | 1 -> 4 ; 18 | 0 -> 1 ; 19 | 2 [penwidth=1,label="'1'}"]; 20 | 5 [penwidth=1,label="'10'}"]; 21 | 11 [penwidth=1,label="'100'}"]; 22 | 5 -> 11 ; 23 | 12 [penwidth=1,label="'101'}"]; 24 | 5 -> 12 ; 25 | 2 -> 5 ; 26 | 6 [penwidth=1,label="'11'}"]; 27 | 13 [penwidth=1,label="'110'}"]; 28 | 6 -> 13 ; 29 | 14 [penwidth=1,label="'111'}"]; 30 | 6 -> 14 ; 31 | 2 -> 6 ; 32 | 0 -> 2 ; 33 | } 34 | -------------------------------------------------------------------------------- /src/pyoram/tests/baselines/k2_h3_b1_data.dot: -------------------------------------------------------------------------------- 1 | // Created by SizedVirtualHeap.write_as_dot(...) 2 | digraph heaptree { 3 | node [shape=record] 4 | 0 [penwidth=1,label="{{0}}"]; 5 | 1 [penwidth=1,label="{{1}}"]; 6 | 3 [penwidth=1,label="{{3}}"]; 7 | 7 [penwidth=1,label="{{7}}"]; 8 | 3 -> 7 ; 9 | 8 [penwidth=1,label="{{8}}"]; 10 | 3 -> 8 ; 11 | 1 -> 3 ; 12 | 4 [penwidth=1,label="{{4}}"]; 13 | 9 [penwidth=1,label="{{9}}"]; 14 | 4 -> 9 ; 15 | 10 [penwidth=1,label="{{10}}"]; 16 | 4 -> 10 ; 17 | 1 -> 4 ; 18 | 0 -> 1 ; 19 | 2 [penwidth=1,label="{{2}}"]; 20 | 5 [penwidth=1,label="{{5}}"]; 21 | 11 [penwidth=1,label="{{11}}"]; 22 | 5 -> 11 ; 23 | 12 [penwidth=1,label="{{12}}"]; 24 | 5 -> 12 ; 25 | 2 -> 5 ; 26 | 6 [penwidth=1,label="{{6}}"]; 27 | 13 [penwidth=1,label="{{13}}"]; 28 | 6 -> 13 ; 29 | 14 [penwidth=1,label="{{14}}"]; 30 | 6 -> 14 ; 31 | 2 -> 6 ; 32 | 0 -> 2 ; 33 | } 34 | -------------------------------------------------------------------------------- /src/pyoram/tests/baselines/k2_h3_b2.dot: -------------------------------------------------------------------------------- 1 | // Created by SizedVirtualHeap.write_as_dot(...) 2 | digraph heaptree { 3 | node [shape=record] 4 | 0 [penwidth=1,label="''}"]; 5 | 1 [penwidth=1,label="'0'}"]; 6 | 3 [penwidth=1,label="'00'}"]; 7 | 7 [penwidth=1,label="'000'}"]; 8 | 3 -> 7 ; 9 | 8 [penwidth=1,label="'001'}"]; 10 | 3 -> 8 ; 11 | 1 -> 3 ; 12 | 4 [penwidth=1,label="'01'}"]; 13 | 9 [penwidth=1,label="'010'}"]; 14 | 4 -> 9 ; 15 | 10 [penwidth=1,label="'011'}"]; 16 | 4 -> 10 ; 17 | 1 -> 4 ; 18 | 0 -> 1 ; 19 | 2 [penwidth=1,label="'1'}"]; 20 | 5 [penwidth=1,label="'10'}"]; 21 | 11 [penwidth=1,label="'100'}"]; 22 | 5 -> 11 ; 23 | 12 [penwidth=1,label="'101'}"]; 24 | 5 -> 12 ; 25 | 2 -> 5 ; 26 | 6 [penwidth=1,label="'11'}"]; 27 | 13 [penwidth=1,label="'110'}"]; 28 | 6 -> 13 ; 29 | 14 [penwidth=1,label="'111'}"]; 30 | 6 -> 14 ; 31 | 2 -> 6 ; 32 | 0 -> 2 ; 33 | } 34 | -------------------------------------------------------------------------------- /src/pyoram/tests/baselines/k2_h3_b2_data.dot: -------------------------------------------------------------------------------- 1 | // Created by SizedVirtualHeap.write_as_dot(...) 2 | digraph heaptree { 3 | node [shape=record] 4 | 0 [penwidth=1,label="{{0}|{1}}"]; 5 | 1 [penwidth=1,label="{{2}|{3}}"]; 6 | 3 [penwidth=1,label="{{6}|{7}}"]; 7 | 7 [penwidth=1,label="{{14}|{15}}"]; 8 | 3 -> 7 ; 9 | 8 [penwidth=1,label="{{16}|{17}}"]; 10 | 3 -> 8 ; 11 | 1 -> 3 ; 12 | 4 [penwidth=1,label="{{8}|{9}}"]; 13 | 9 [penwidth=1,label="{{18}|{19}}"]; 14 | 4 -> 9 ; 15 | 10 [penwidth=1,label="{{20}|{21}}"]; 16 | 4 -> 10 ; 17 | 1 -> 4 ; 18 | 0 -> 1 ; 19 | 2 [penwidth=1,label="{{4}|{5}}"]; 20 | 5 [penwidth=1,label="{{10}|{11}}"]; 21 | 11 [penwidth=1,label="{{22}|{23}}"]; 22 | 5 -> 11 ; 23 | 12 [penwidth=1,label="{{24}|{25}}"]; 24 | 5 -> 12 ; 25 | 2 -> 5 ; 26 | 6 [penwidth=1,label="{{12}|{13}}"]; 27 | 13 [penwidth=1,label="{{26}|{27}}"]; 28 | 6 -> 13 ; 29 | 14 [penwidth=1,label="{{28}|{29}}"]; 30 | 6 -> 14 ; 31 | 2 -> 6 ; 32 | 0 -> 2 ; 33 | } 34 | -------------------------------------------------------------------------------- /src/pyoram/tests/baselines/k3_h3_b1.dot: -------------------------------------------------------------------------------- 1 | // Created by SizedVirtualHeap.write_as_dot(...) 2 | digraph heaptree { 3 | node [shape=record] 4 | 0 [penwidth=1,label="''}"]; 5 | 1 [penwidth=1,label="'0'}"]; 6 | 4 [penwidth=1,label="'00'}"]; 7 | 13 [penwidth=1,label="'000'}"]; 8 | 4 -> 13 ; 9 | 14 [penwidth=1,label="'001'}"]; 10 | 4 -> 14 ; 11 | 15 [penwidth=1,label="'002'}"]; 12 | 4 -> 15 ; 13 | 1 -> 4 ; 14 | 5 [penwidth=1,label="'01'}"]; 15 | 16 [penwidth=1,label="'010'}"]; 16 | 5 -> 16 ; 17 | 17 [penwidth=1,label="'011'}"]; 18 | 5 -> 17 ; 19 | 18 [penwidth=1,label="'012'}"]; 20 | 5 -> 18 ; 21 | 1 -> 5 ; 22 | 6 [penwidth=1,label="'02'}"]; 23 | 19 [penwidth=1,label="'020'}"]; 24 | 6 -> 19 ; 25 | 20 [penwidth=1,label="'021'}"]; 26 | 6 -> 20 ; 27 | 21 [penwidth=1,label="'022'}"]; 28 | 6 -> 21 ; 29 | 1 -> 6 ; 30 | 0 -> 1 ; 31 | 2 [penwidth=1,label="'1'}"]; 32 | 7 [penwidth=1,label="'10'}"]; 33 | 22 [penwidth=1,label="'100'}"]; 34 | 7 -> 22 ; 35 | 23 [penwidth=1,label="'101'}"]; 36 | 7 -> 23 ; 37 | 24 [penwidth=1,label="'102'}"]; 38 | 7 -> 24 ; 39 | 2 -> 7 ; 40 | 8 [penwidth=1,label="'11'}"]; 41 | 25 [penwidth=1,label="'110'}"]; 42 | 8 -> 25 ; 43 | 26 [penwidth=1,label="'111'}"]; 44 | 8 -> 26 ; 45 | 27 [penwidth=1,label="'112'}"]; 46 | 8 -> 27 ; 47 | 2 -> 8 ; 48 | 9 [penwidth=1,label="'12'}"]; 49 | 28 [penwidth=1,label="'120'}"]; 50 | 9 -> 28 ; 51 | 29 [penwidth=1,label="'121'}"]; 52 | 9 -> 29 ; 53 | 30 [penwidth=1,label="'122'}"]; 54 | 9 -> 30 ; 55 | 2 -> 9 ; 56 | 0 -> 2 ; 57 | 3 [penwidth=1,label="'2'}"]; 58 | 10 [penwidth=1,label="'20'}"]; 59 | 31 [penwidth=1,label="'200'}"]; 60 | 10 -> 31 ; 61 | 32 [penwidth=1,label="'201'}"]; 62 | 10 -> 32 ; 63 | 33 [penwidth=1,label="'202'}"]; 64 | 10 -> 33 ; 65 | 3 -> 10 ; 66 | 11 [penwidth=1,label="'21'}"]; 67 | 34 [penwidth=1,label="'210'}"]; 68 | 11 -> 34 ; 69 | 35 [penwidth=1,label="'211'}"]; 70 | 11 -> 35 ; 71 | 36 [penwidth=1,label="'212'}"]; 72 | 11 -> 36 ; 73 | 3 -> 11 ; 74 | 12 [penwidth=1,label="'22'}"]; 75 | 37 [penwidth=1,label="'220'}"]; 76 | 12 -> 37 ; 77 | 38 [penwidth=1,label="'221'}"]; 78 | 12 -> 38 ; 79 | 39 [penwidth=1,label="'222'}"]; 80 | 12 -> 39 ; 81 | 3 -> 12 ; 82 | 0 -> 3 ; 83 | } 84 | -------------------------------------------------------------------------------- /src/pyoram/tests/baselines/k3_h3_b1_data.dot: -------------------------------------------------------------------------------- 1 | // Created by SizedVirtualHeap.write_as_dot(...) 2 | digraph heaptree { 3 | node [shape=record] 4 | 0 [penwidth=1,label="{{0}}"]; 5 | 1 [penwidth=1,label="{{1}}"]; 6 | 4 [penwidth=1,label="{{4}}"]; 7 | 13 [penwidth=1,label="{{13}}"]; 8 | 4 -> 13 ; 9 | 14 [penwidth=1,label="{{14}}"]; 10 | 4 -> 14 ; 11 | 15 [penwidth=1,label="{{15}}"]; 12 | 4 -> 15 ; 13 | 1 -> 4 ; 14 | 5 [penwidth=1,label="{{5}}"]; 15 | 16 [penwidth=1,label="{{16}}"]; 16 | 5 -> 16 ; 17 | 17 [penwidth=1,label="{{17}}"]; 18 | 5 -> 17 ; 19 | 18 [penwidth=1,label="{{18}}"]; 20 | 5 -> 18 ; 21 | 1 -> 5 ; 22 | 6 [penwidth=1,label="{{6}}"]; 23 | 19 [penwidth=1,label="{{19}}"]; 24 | 6 -> 19 ; 25 | 20 [penwidth=1,label="{{20}}"]; 26 | 6 -> 20 ; 27 | 21 [penwidth=1,label="{{21}}"]; 28 | 6 -> 21 ; 29 | 1 -> 6 ; 30 | 0 -> 1 ; 31 | 2 [penwidth=1,label="{{2}}"]; 32 | 7 [penwidth=1,label="{{7}}"]; 33 | 22 [penwidth=1,label="{{22}}"]; 34 | 7 -> 22 ; 35 | 23 [penwidth=1,label="{{23}}"]; 36 | 7 -> 23 ; 37 | 24 [penwidth=1,label="{{24}}"]; 38 | 7 -> 24 ; 39 | 2 -> 7 ; 40 | 8 [penwidth=1,label="{{8}}"]; 41 | 25 [penwidth=1,label="{{25}}"]; 42 | 8 -> 25 ; 43 | 26 [penwidth=1,label="{{26}}"]; 44 | 8 -> 26 ; 45 | 27 [penwidth=1,label="{{27}}"]; 46 | 8 -> 27 ; 47 | 2 -> 8 ; 48 | 9 [penwidth=1,label="{{9}}"]; 49 | 28 [penwidth=1,label="{{28}}"]; 50 | 9 -> 28 ; 51 | 29 [penwidth=1,label="{{29}}"]; 52 | 9 -> 29 ; 53 | 30 [penwidth=1,label="{{30}}"]; 54 | 9 -> 30 ; 55 | 2 -> 9 ; 56 | 0 -> 2 ; 57 | 3 [penwidth=1,label="{{3}}"]; 58 | 10 [penwidth=1,label="{{10}}"]; 59 | 31 [penwidth=1,label="{{31}}"]; 60 | 10 -> 31 ; 61 | 32 [penwidth=1,label="{{32}}"]; 62 | 10 -> 32 ; 63 | 33 [penwidth=1,label="{{33}}"]; 64 | 10 -> 33 ; 65 | 3 -> 10 ; 66 | 11 [penwidth=1,label="{{11}}"]; 67 | 34 [penwidth=1,label="{{34}}"]; 68 | 11 -> 34 ; 69 | 35 [penwidth=1,label="{{35}}"]; 70 | 11 -> 35 ; 71 | 36 [penwidth=1,label="{{36}}"]; 72 | 11 -> 36 ; 73 | 3 -> 11 ; 74 | 12 [penwidth=1,label="{{12}}"]; 75 | 37 [penwidth=1,label="{{37}}"]; 76 | 12 -> 37 ; 77 | 38 [penwidth=1,label="{{38}}"]; 78 | 12 -> 38 ; 79 | 39 [penwidth=1,label="{{39}}"]; 80 | 12 -> 39 ; 81 | 3 -> 12 ; 82 | 0 -> 3 ; 83 | } 84 | -------------------------------------------------------------------------------- /src/pyoram/tests/baselines/k3_h3_b2.dot: -------------------------------------------------------------------------------- 1 | // Created by SizedVirtualHeap.write_as_dot(...) 2 | digraph heaptree { 3 | node [shape=record] 4 | 0 [penwidth=1,label="''}"]; 5 | 1 [penwidth=1,label="'0'}"]; 6 | 4 [penwidth=1,label="'00'}"]; 7 | 13 [penwidth=1,label="'000'}"]; 8 | 4 -> 13 ; 9 | 14 [penwidth=1,label="'001'}"]; 10 | 4 -> 14 ; 11 | 15 [penwidth=1,label="'002'}"]; 12 | 4 -> 15 ; 13 | 1 -> 4 ; 14 | 5 [penwidth=1,label="'01'}"]; 15 | 16 [penwidth=1,label="'010'}"]; 16 | 5 -> 16 ; 17 | 17 [penwidth=1,label="'011'}"]; 18 | 5 -> 17 ; 19 | 18 [penwidth=1,label="'012'}"]; 20 | 5 -> 18 ; 21 | 1 -> 5 ; 22 | 6 [penwidth=1,label="'02'}"]; 23 | 19 [penwidth=1,label="'020'}"]; 24 | 6 -> 19 ; 25 | 20 [penwidth=1,label="'021'}"]; 26 | 6 -> 20 ; 27 | 21 [penwidth=1,label="'022'}"]; 28 | 6 -> 21 ; 29 | 1 -> 6 ; 30 | 0 -> 1 ; 31 | 2 [penwidth=1,label="'1'}"]; 32 | 7 [penwidth=1,label="'10'}"]; 33 | 22 [penwidth=1,label="'100'}"]; 34 | 7 -> 22 ; 35 | 23 [penwidth=1,label="'101'}"]; 36 | 7 -> 23 ; 37 | 24 [penwidth=1,label="'102'}"]; 38 | 7 -> 24 ; 39 | 2 -> 7 ; 40 | 8 [penwidth=1,label="'11'}"]; 41 | 25 [penwidth=1,label="'110'}"]; 42 | 8 -> 25 ; 43 | 26 [penwidth=1,label="'111'}"]; 44 | 8 -> 26 ; 45 | 27 [penwidth=1,label="'112'}"]; 46 | 8 -> 27 ; 47 | 2 -> 8 ; 48 | 9 [penwidth=1,label="'12'}"]; 49 | 28 [penwidth=1,label="'120'}"]; 50 | 9 -> 28 ; 51 | 29 [penwidth=1,label="'121'}"]; 52 | 9 -> 29 ; 53 | 30 [penwidth=1,label="'122'}"]; 54 | 9 -> 30 ; 55 | 2 -> 9 ; 56 | 0 -> 2 ; 57 | 3 [penwidth=1,label="'2'}"]; 58 | 10 [penwidth=1,label="'20'}"]; 59 | 31 [penwidth=1,label="'200'}"]; 60 | 10 -> 31 ; 61 | 32 [penwidth=1,label="'201'}"]; 62 | 10 -> 32 ; 63 | 33 [penwidth=1,label="'202'}"]; 64 | 10 -> 33 ; 65 | 3 -> 10 ; 66 | 11 [penwidth=1,label="'21'}"]; 67 | 34 [penwidth=1,label="'210'}"]; 68 | 11 -> 34 ; 69 | 35 [penwidth=1,label="'211'}"]; 70 | 11 -> 35 ; 71 | 36 [penwidth=1,label="'212'}"]; 72 | 11 -> 36 ; 73 | 3 -> 11 ; 74 | 12 [penwidth=1,label="'22'}"]; 75 | 37 [penwidth=1,label="'220'}"]; 76 | 12 -> 37 ; 77 | 38 [penwidth=1,label="'221'}"]; 78 | 12 -> 38 ; 79 | 39 [penwidth=1,label="'222'}"]; 80 | 12 -> 39 ; 81 | 3 -> 12 ; 82 | 0 -> 3 ; 83 | } 84 | -------------------------------------------------------------------------------- /src/pyoram/tests/baselines/k3_h3_b2_data.dot: -------------------------------------------------------------------------------- 1 | // Created by SizedVirtualHeap.write_as_dot(...) 2 | digraph heaptree { 3 | node [shape=record] 4 | 0 [penwidth=1,label="{{0}|{1}}"]; 5 | 1 [penwidth=1,label="{{2}|{3}}"]; 6 | 4 [penwidth=1,label="{{8}|{9}}"]; 7 | 13 [penwidth=1,label="{{26}|{27}}"]; 8 | 4 -> 13 ; 9 | 14 [penwidth=1,label="{{28}|{29}}"]; 10 | 4 -> 14 ; 11 | 15 [penwidth=1,label="{{30}|{31}}"]; 12 | 4 -> 15 ; 13 | 1 -> 4 ; 14 | 5 [penwidth=1,label="{{10}|{11}}"]; 15 | 16 [penwidth=1,label="{{32}|{33}}"]; 16 | 5 -> 16 ; 17 | 17 [penwidth=1,label="{{34}|{35}}"]; 18 | 5 -> 17 ; 19 | 18 [penwidth=1,label="{{36}|{37}}"]; 20 | 5 -> 18 ; 21 | 1 -> 5 ; 22 | 6 [penwidth=1,label="{{12}|{13}}"]; 23 | 19 [penwidth=1,label="{{38}|{39}}"]; 24 | 6 -> 19 ; 25 | 20 [penwidth=1,label="{{40}|{41}}"]; 26 | 6 -> 20 ; 27 | 21 [penwidth=1,label="{{42}|{43}}"]; 28 | 6 -> 21 ; 29 | 1 -> 6 ; 30 | 0 -> 1 ; 31 | 2 [penwidth=1,label="{{4}|{5}}"]; 32 | 7 [penwidth=1,label="{{14}|{15}}"]; 33 | 22 [penwidth=1,label="{{44}|{45}}"]; 34 | 7 -> 22 ; 35 | 23 [penwidth=1,label="{{46}|{47}}"]; 36 | 7 -> 23 ; 37 | 24 [penwidth=1,label="{{48}|{49}}"]; 38 | 7 -> 24 ; 39 | 2 -> 7 ; 40 | 8 [penwidth=1,label="{{16}|{17}}"]; 41 | 25 [penwidth=1,label="{{50}|{51}}"]; 42 | 8 -> 25 ; 43 | 26 [penwidth=1,label="{{52}|{53}}"]; 44 | 8 -> 26 ; 45 | 27 [penwidth=1,label="{{54}|{55}}"]; 46 | 8 -> 27 ; 47 | 2 -> 8 ; 48 | 9 [penwidth=1,label="{{18}|{19}}"]; 49 | 28 [penwidth=1,label="{{56}|{57}}"]; 50 | 9 -> 28 ; 51 | 29 [penwidth=1,label="{{58}|{59}}"]; 52 | 9 -> 29 ; 53 | 30 [penwidth=1,label="{{60}|{61}}"]; 54 | 9 -> 30 ; 55 | 2 -> 9 ; 56 | 0 -> 2 ; 57 | 3 [penwidth=1,label="{{6}|{7}}"]; 58 | 10 [penwidth=1,label="{{20}|{21}}"]; 59 | 31 [penwidth=1,label="{{62}|{63}}"]; 60 | 10 -> 31 ; 61 | 32 [penwidth=1,label="{{64}|{65}}"]; 62 | 10 -> 32 ; 63 | 33 [penwidth=1,label="{{66}|{67}}"]; 64 | 10 -> 33 ; 65 | 3 -> 10 ; 66 | 11 [penwidth=1,label="{{22}|{23}}"]; 67 | 34 [penwidth=1,label="{{68}|{69}}"]; 68 | 11 -> 34 ; 69 | 35 [penwidth=1,label="{{70}|{71}}"]; 70 | 11 -> 35 ; 71 | 36 [penwidth=1,label="{{72}|{73}}"]; 72 | 11 -> 36 ; 73 | 3 -> 11 ; 74 | 12 [penwidth=1,label="{{24}|{25}}"]; 75 | 37 [penwidth=1,label="{{74}|{75}}"]; 76 | 12 -> 37 ; 77 | 38 [penwidth=1,label="{{76}|{77}}"]; 78 | 12 -> 38 ; 79 | 39 [penwidth=1,label="{{78}|{79}}"]; 80 | 12 -> 39 ; 81 | 3 -> 12 ; 82 | 0 -> 3 ; 83 | } 84 | -------------------------------------------------------------------------------- /src/pyoram/tests/test_aes.py: -------------------------------------------------------------------------------- 1 | import unittest 2 | 3 | from pyoram.crypto.aes import AES 4 | 5 | class TestAES(unittest.TestCase): 6 | 7 | def test_KeyGen(self): 8 | self.assertTrue(len(AES.key_sizes) in (3,4)) 9 | self.assertTrue(len(set(AES.key_sizes)) in (3,4)) 10 | for keysize in AES.key_sizes: 11 | key_list = [] 12 | key_set = set() 13 | for i in range(10): 14 | k = AES.KeyGen(keysize) 15 | self.assertEqual(len(k), keysize) 16 | key_list.append(k) 17 | key_set.add(k) 18 | self.assertEqual(len(key_list), 10) 19 | # make sure every key is unique 20 | self.assertEqual(len(key_list), len(key_set)) 21 | 22 | def test_CTR(self): 23 | self._test_Enc_Dec( 24 | AES.CTREnc, 25 | AES.CTRDec, 26 | lambda i, size: bytes(bytearray([i]) * size), 27 | [16,24,32]) 28 | 29 | def test_GCM(self): 30 | self._test_Enc_Dec( 31 | AES.GCMEnc, 32 | AES.GCMDec, 33 | lambda i, size: bytes(bytearray([i]) * size), 34 | [16,24,32]) 35 | 36 | def _test_Enc_Dec(self, 37 | enc_func, 38 | dec_func, 39 | get_plaintext, 40 | keysizes): 41 | keysizes = list(keysizes) 42 | self.assertTrue(len(keysizes) > 0) 43 | blocksize_factor = [0.5, 1, 1.5, 2, 2.5] 44 | plaintext_blocks = [] 45 | for i, f in enumerate(blocksize_factor): 46 | size = AES.block_size * f 47 | size = int(round(size)) 48 | if int(f) != f: 49 | assert (size % AES.block_size) != 0 50 | plaintext_blocks.append(get_plaintext(i, size)) 51 | 52 | assert len(AES.key_sizes) > 0 53 | ciphertext_blocks = {} 54 | keys = {} 55 | for keysize in keysizes: 56 | key = AES.KeyGen(keysize) 57 | keys[keysize] = key 58 | ciphertext_blocks[keysize] = [] 59 | for block in plaintext_blocks: 60 | ciphertext_blocks[keysize].append( 61 | enc_func(key, block)) 62 | 63 | self.assertEqual(len(ciphertext_blocks), 64 | len(keysizes)) 65 | self.assertEqual(len(keys), 66 | len(keysizes)) 67 | 68 | plaintext_decrypted_blocks = {} 69 | for keysize in keys: 70 | key = keys[keysize] 71 | plaintext_decrypted_blocks[keysize] = [] 72 | for block in ciphertext_blocks[keysize]: 73 | plaintext_decrypted_blocks[keysize].append( 74 | dec_func(key, block)) 75 | 76 | self.assertEqual(len(plaintext_decrypted_blocks), 77 | len(keysizes)) 78 | 79 | for i in range(len(blocksize_factor)): 80 | for keysize in keysizes: 81 | self.assertEqual( 82 | plaintext_blocks[i], 83 | plaintext_decrypted_blocks[keysize][i]) 84 | self.assertNotEqual( 85 | plaintext_blocks[i], 86 | ciphertext_blocks[keysize][i]) 87 | if enc_func is AES.CTREnc: 88 | self.assertEqual( 89 | len(ciphertext_blocks[keysize][i]), 90 | len(plaintext_blocks[i]) + AES.block_size) 91 | else: 92 | assert enc_func is AES.GCMEnc 93 | self.assertEqual( 94 | len(ciphertext_blocks[keysize][i]), 95 | len(plaintext_blocks[i]) + 2*AES.block_size) 96 | # check IND-CPA 97 | key = keys[keysize] 98 | alt_ciphertext = enc_func(key, plaintext_blocks[i]) 99 | self.assertNotEqual( 100 | ciphertext_blocks[keysize][i], 101 | alt_ciphertext) 102 | self.assertEqual( 103 | len(ciphertext_blocks[keysize][i]), 104 | len(alt_ciphertext)) 105 | self.assertNotEqual( 106 | ciphertext_blocks[keysize][i][:AES.block_size], 107 | alt_ciphertext[:AES.block_size]) 108 | self.assertNotEqual( 109 | ciphertext_blocks[keysize][i][AES.block_size:], 110 | alt_ciphertext[AES.block_size:]) 111 | 112 | if __name__ == "__main__": 113 | unittest.main() # pragma: no cover 114 | -------------------------------------------------------------------------------- /src/pyoram/tests/test_examples.py: -------------------------------------------------------------------------------- 1 | import os 2 | import glob 3 | import sys 4 | import unittest 5 | 6 | thisfile = os.path.abspath(__file__) 7 | thisdir = os.path.dirname(thisfile) 8 | topdir = os.path.dirname( 9 | os.path.dirname( 10 | os.path.dirname(thisdir))) 11 | exdir = os.path.join(topdir, 'examples') 12 | examples = glob.glob(os.path.join(exdir,"*.py")) 13 | 14 | assert os.path.exists(exdir) 15 | assert thisfile not in examples 16 | 17 | tdict = {} 18 | for fname in examples: 19 | basename = os.path.basename(fname) 20 | assert basename.endswith('.py') 21 | assert len(basename) >= 3 22 | basename = basename[:-3] 23 | tname = 'test_'+basename 24 | tdict[tname] = fname, basename 25 | 26 | assert len(tdict) == len(examples) 27 | 28 | assert 'test_encrypted_storage_s3' in tdict 29 | assert 'test_path_oram_s3' in tdict 30 | if 'PYORAM_AWS_TEST_BUCKET' not in os.environ: 31 | del tdict['test_encrypted_storage_s3'] 32 | del tdict['test_path_oram_s3'] 33 | assert 'test_encrypted_storage_sftp' in tdict 34 | assert 'test_path_oram_sftp' in tdict 35 | assert 'test_path_oram_sftp_setup' in tdict 36 | assert 'test_path_oram_sftp_test' in tdict 37 | if 'PYORAM_SSH_TEST_HOST' not in os.environ: 38 | del tdict['test_encrypted_storage_sftp'] 39 | del tdict['test_path_oram_sftp'] 40 | del tdict['test_path_oram_sftp_setup'] 41 | del tdict['test_path_oram_sftp_test'] 42 | 43 | def _execute_example(example_name): 44 | filename, basename = tdict[example_name] 45 | assert os.path.exists(filename) 46 | try: 47 | sys.path.insert(0, exdir) 48 | m = __import__(basename) 49 | m.main() 50 | finally: 51 | sys.path.remove(exdir) 52 | 53 | # this is recognized by nosetests as 54 | # a dynamic test generator 55 | def test_generator(): 56 | for example_name in sorted(tdict): 57 | yield _execute_example, example_name 58 | 59 | if __name__ == "__main__": 60 | for tfunc, tname in test_generator(): # pragma: no cover 61 | tfunc(tname) # pragma: no cover 62 | -------------------------------------------------------------------------------- /src/pyoram/tests/test_misc.py: -------------------------------------------------------------------------------- 1 | import os 2 | import unittest 3 | import tempfile 4 | 5 | import pyoram.util.misc 6 | 7 | class Test(unittest.TestCase): 8 | 9 | def test_log2floor(self): 10 | self.assertEqual(pyoram.util.misc.log2floor(1), 0) 11 | self.assertEqual(pyoram.util.misc.log2floor(2), 1) 12 | self.assertEqual(pyoram.util.misc.log2floor(3), 1) 13 | self.assertEqual(pyoram.util.misc.log2floor(4), 2) 14 | self.assertEqual(pyoram.util.misc.log2floor(5), 2) 15 | self.assertEqual(pyoram.util.misc.log2floor(6), 2) 16 | self.assertEqual(pyoram.util.misc.log2floor(7), 2) 17 | self.assertEqual(pyoram.util.misc.log2floor(8), 3) 18 | self.assertEqual(pyoram.util.misc.log2floor(9), 3) 19 | 20 | def test_log2ceil(self): 21 | self.assertEqual(pyoram.util.misc.log2ceil(1), 0) 22 | self.assertEqual(pyoram.util.misc.log2ceil(2), 1) 23 | self.assertEqual(pyoram.util.misc.log2ceil(3), 2) 24 | self.assertEqual(pyoram.util.misc.log2ceil(4), 2) 25 | self.assertEqual(pyoram.util.misc.log2ceil(5), 3) 26 | self.assertEqual(pyoram.util.misc.log2ceil(6), 3) 27 | self.assertEqual(pyoram.util.misc.log2ceil(7), 3) 28 | self.assertEqual(pyoram.util.misc.log2ceil(8), 3) 29 | self.assertEqual(pyoram.util.misc.log2ceil(9), 4) 30 | 31 | def test_intdivceil(self): 32 | 33 | with self.assertRaises(ZeroDivisionError): 34 | pyoram.util.misc.intdivceil(0, 0) 35 | with self.assertRaises(ZeroDivisionError): 36 | pyoram.util.misc.intdivceil(1, 0) 37 | 38 | self.assertEqual(pyoram.util.misc.intdivceil(1, 1), 1) 39 | self.assertEqual(pyoram.util.misc.intdivceil(2, 3), 1) 40 | self.assertEqual(2 // 3, 0) 41 | self.assertEqual(pyoram.util.misc.intdivceil( 42 | 123123123123123123123123123123123123123123123123, 43 | 123123123123123123123123123123123123123123123123), 1) 44 | self.assertEqual(pyoram.util.misc.intdivceil( 45 | 2 * 123123123123123123123123123123123123123123123123, 46 | 123123123123123123123123123123123123123123123123), 2) 47 | self.assertEqual(pyoram.util.misc.intdivceil( 48 | 2 * 123123123123123123123123123123123123123123123123 + 1, 49 | 123123123123123123123123123123123123123123123123), 3) 50 | self.assertEqual(pyoram.util.misc.intdivceil( 51 | 2 * 123123123123123123123123123123123123123123123123 - 1, 52 | 123123123123123123123123123123123123123123123123), 2) 53 | self.assertEqual( 54 | (2 * 123123123123123123123123123123123123123123123123 - 1) // \ 55 | 123123123123123123123123123123123123123123123123, 56 | 1) 57 | 58 | def test_MemorySize(self): 59 | self.assertTrue("b" in str(pyoram.util.misc.MemorySize(0.1))) 60 | self.assertTrue("B" in str(pyoram.util.misc.MemorySize(1))) 61 | self.assertTrue("B" in str(pyoram.util.misc.MemorySize(999))) 62 | self.assertTrue("KB" in str(pyoram.util.misc.MemorySize(1000))) 63 | self.assertTrue("KB" in str(pyoram.util.misc.MemorySize(999999))) 64 | self.assertTrue("MB" in str(pyoram.util.misc.MemorySize(1000000))) 65 | self.assertTrue("MB" in str(pyoram.util.misc.MemorySize(999999999))) 66 | self.assertTrue("GB" in str(pyoram.util.misc.MemorySize(1000000000))) 67 | self.assertTrue("GB" in str(pyoram.util.misc.MemorySize(9999999999))) 68 | self.assertTrue("TB" in str(pyoram.util.misc.MemorySize(1000000000000))) 69 | self.assertTrue("b" in str(pyoram.util.misc.MemorySize(1, unit="b"))) 70 | self.assertTrue("b" in str(pyoram.util.misc.MemorySize(2, unit="b"))) 71 | self.assertTrue("b" in str(pyoram.util.misc.MemorySize(7.9, unit="b"))) 72 | 73 | self.assertTrue("B" in str(pyoram.util.misc.MemorySize(8, unit="b"))) 74 | self.assertTrue("B" in str(pyoram.util.misc.MemorySize(1, unit="B"))) 75 | self.assertTrue("B" in str(pyoram.util.misc.MemorySize(999, unit="B"))) 76 | 77 | self.assertTrue("KB" in str(pyoram.util.misc.MemorySize(1000, unit="B"))) 78 | self.assertTrue("KB" in str(pyoram.util.misc.MemorySize(1, unit="KB"))) 79 | self.assertTrue("KB" in str(pyoram.util.misc.MemorySize(999, unit="KB"))) 80 | self.assertTrue("MB" in str(pyoram.util.misc.MemorySize(1000, unit="KB"))) 81 | self.assertTrue("MB" in str(pyoram.util.misc.MemorySize(1, unit="MB"))) 82 | self.assertTrue("MB" in str(pyoram.util.misc.MemorySize(999, unit="MB"))) 83 | self.assertTrue("GB" in str(pyoram.util.misc.MemorySize(1000, unit="MB"))) 84 | self.assertTrue("GB" in str(pyoram.util.misc.MemorySize(1, unit="GB"))) 85 | self.assertTrue("GB" in str(pyoram.util.misc.MemorySize(999, unit="GB"))) 86 | self.assertTrue("TB" in str(pyoram.util.misc.MemorySize(1000, unit="GB"))) 87 | self.assertTrue("TB" in str(pyoram.util.misc.MemorySize(1, unit="TB"))) 88 | 89 | self.assertEqual(pyoram.util.misc.MemorySize(1024).KiB, 1) 90 | self.assertEqual(pyoram.util.misc.MemorySize(1024**2).MiB, 1) 91 | self.assertEqual(pyoram.util.misc.MemorySize(1024**3).GiB, 1) 92 | self.assertEqual(pyoram.util.misc.MemorySize(1024**4).TiB, 1) 93 | 94 | def test_saveload_private_key(self): 95 | with tempfile.NamedTemporaryFile(delete=False) as f: 96 | filename = f.name 97 | try: 98 | key = os.urandom(32) 99 | pyoram.util.misc.save_private_key(filename, key) 100 | loaded_key = pyoram.util.misc.load_private_key(filename) 101 | self.assertEqual(key, loaded_key) 102 | finally: 103 | os.remove(filename) 104 | 105 | def test_chunkiter(self): 106 | self.assertEqual(list(pyoram.util.misc.chunkiter([1,2,3,4,5], 1)), 107 | [[1],[2],[3],[4],[5]]) 108 | self.assertEqual(list(pyoram.util.misc.chunkiter([1,2,3,4,5], 2)), 109 | [[1,2],[3,4],[5]]) 110 | self.assertEqual(list(pyoram.util.misc.chunkiter([1,2,3,4,5], 3)), 111 | [[1,2,3],[4,5]]) 112 | self.assertEqual(list(pyoram.util.misc.chunkiter([1,2,3,4,5], 4)), 113 | [[1,2,3,4],[5]]) 114 | self.assertEqual(list(pyoram.util.misc.chunkiter([1,2,3,4,5], 5)), 115 | [[1,2,3,4,5]]) 116 | self.assertEqual(list(pyoram.util.misc.chunkiter([1,2,3,4,5], 6)), 117 | [[1,2,3,4,5]]) 118 | self.assertEqual(list(pyoram.util.misc.chunkiter([], 1)), 119 | []) 120 | self.assertEqual(list(pyoram.util.misc.chunkiter([], 2)), 121 | []) 122 | 123 | if __name__ == "__main__": 124 | unittest.main() # pragma: no cover 125 | -------------------------------------------------------------------------------- /src/pyoram/tests/test_package.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import unittest 3 | 4 | import pyoram 5 | 6 | is_pypy = False 7 | try: 8 | import __pypy__ 9 | is_pypy = True 10 | except ImportError: 11 | is_pypy = False 12 | 13 | class Test(unittest.TestCase): 14 | 15 | # See what Python versions the combined 16 | # coverage report includes 17 | def test_show_coverage(self): 18 | if not is_pypy: 19 | if sys.version_info.major == 2: 20 | if sys.version_info.minor == 7: 21 | print(sys.version_info) 22 | elif sys.version_info.major == 3: 23 | if sys.version_info.minor == 4: 24 | print(sys.version_info) 25 | elif sys.version_info.minor == 5: 26 | print(sys.version_info) 27 | elif sys.version_info.minor == 6: 28 | print(sys.version_info) 29 | elif sys.version_info.minor == 7: 30 | print(sys.version_info) 31 | elif sys.version_info.minor == 8: 32 | print(sys.version_info) 33 | if is_pypy: 34 | if sys.version_info.major == 2: 35 | print(sys.version_info) 36 | elif sys.version_info.major == 3: 37 | print(sys.version_info) 38 | 39 | def test_version(self): 40 | pyoram.__version__ 41 | 42 | if __name__ == "__main__": 43 | unittest.main() # pragma: no cover 44 | -------------------------------------------------------------------------------- /src/pyoram/util/__init__.py: -------------------------------------------------------------------------------- 1 | import pyoram.util.misc 2 | import pyoram.util.virtual_heap 3 | -------------------------------------------------------------------------------- /src/pyoram/util/misc.py: -------------------------------------------------------------------------------- 1 | import base64 2 | 3 | import six 4 | 5 | def log2floor(n): 6 | """ 7 | Returns the exact value of floor(log2(n)). 8 | No floating point calculations are used. 9 | Requires positive integer type. 10 | """ 11 | assert n > 0 12 | return n.bit_length() - 1 13 | 14 | def log2ceil(n): 15 | """ 16 | Returns the exact value of ceil(log2(n)). 17 | No floating point calculations are used. 18 | Requires positive integer type. 19 | """ 20 | if n == 1: 21 | return 0 22 | return log2floor(n-1) + 1 23 | 24 | def intdivceil(x, y): 25 | """ 26 | Returns the exact value of ceil(x // y). 27 | No floating point calculations are used. 28 | Requires positive integer types. The result 29 | is undefined if at least one of the inputs 30 | is floating point. 31 | """ 32 | result = x // y 33 | if (x % y): 34 | result += 1 35 | return result 36 | 37 | def save_private_key(filename, key): 38 | with open(filename, "wb") as f: 39 | f.write(base64.b64encode(key)) 40 | 41 | def load_private_key(filename): 42 | with open(filename, "rb") as f: 43 | return base64.b64decode(f.read()) 44 | 45 | from fractions import Fraction 46 | 47 | class MemorySize(object): 48 | 49 | to_bytes = {} 50 | to_bytes['b'] = lambda x: Fraction(x,8) 51 | to_bytes['B'] = lambda x: Fraction(x,1) 52 | to_bytes['KB'] = lambda x: Fraction(1000*x,1) 53 | to_bytes['MB'] = lambda x: Fraction((1000**2)*x,1) 54 | to_bytes['GB'] = lambda x: Fraction((1000**3)*x,1) 55 | to_bytes['TB'] = lambda x: Fraction((1000**4)*x,1) 56 | to_bytes['KiB'] = lambda x: Fraction(1024*x,1) 57 | to_bytes['MiB'] = lambda x: Fraction((1024**2)*x,1) 58 | to_bytes['GiB'] = lambda x: Fraction((1024**3)*x,1) 59 | to_bytes['TiB'] = lambda x: Fraction((1024**4)*x,1) 60 | 61 | def __init__(self, size, unit='B'): 62 | assert size >= 0 63 | self.numbytes = MemorySize.to_bytes[unit](Fraction.from_float(size)) 64 | 65 | def __str__(self): 66 | if self.B < 1: 67 | return "%.3f b" % (self.b) 68 | if self.KB < 1: 69 | return "%.3f B" % (self.B) 70 | if self.MB < 1: 71 | return "%.3f KB" % (self.KB) 72 | if self.GB < 1: 73 | return "%.3f MB" % (self.MB) 74 | if self.TB < 1: 75 | return "%.3f GB" % (self.GB) 76 | return "%.3f TB" % (self.TB) 77 | 78 | @property 79 | def b(self): return self.numbytes*8 80 | @property 81 | def B(self): return self.numbytes 82 | 83 | @property 84 | def KB(self): return self.B/1000 85 | @property 86 | def MB(self): return self.KB/1000 87 | @property 88 | def GB(self): return self.MB/1000 89 | @property 90 | def TB(self): return self.GB/1000 91 | 92 | @property 93 | def KiB(self): return self.B/1024 94 | @property 95 | def MiB(self): return self.KiB/1024 96 | @property 97 | def GiB(self): return self.MiB/1024 98 | @property 99 | def TiB(self): return self.GiB/1024 100 | 101 | def chunkiter(objs, n=100): 102 | """ 103 | Chunk an iterator of unknown size. The optional 104 | keyword 'n' sets the chunk size (default 100). 105 | """ 106 | 107 | objs = iter(objs) 108 | try: 109 | while (1): 110 | chunk = [] 111 | while len(chunk) < n: 112 | chunk.append(six.next(objs)) 113 | yield chunk 114 | except StopIteration: 115 | pass 116 | if len(chunk): 117 | yield chunk 118 | -------------------------------------------------------------------------------- /src/pyoram/util/virtual_heap.py: -------------------------------------------------------------------------------- 1 | __all__ = ("VirtualHeap", 2 | "SizedVirtualHeap") 3 | 4 | import os 5 | import sys 6 | import subprocess 7 | import random 8 | import string 9 | import tempfile 10 | 11 | from six.moves import xrange 12 | 13 | from pyoram.util._virtual_heap_helper import lib as _clib 14 | from pyoram.util.misc import log2floor 15 | 16 | numerals = ''.join([c for c in string.printable \ 17 | if ((c not in string.whitespace) and \ 18 | (c != '+') and (c != '-') and \ 19 | (c != '"') and (c != "'") and \ 20 | (c != '\\') and (c != '/'))]) 21 | numeral_index = dict((c,i) for i,c in enumerate(numerals)) 22 | 23 | # The maximum heap base for which base k labels 24 | # can be produced. 25 | max_k_labeled = len(numerals) 26 | 27 | def base10_integer_to_basek_string(k, x): 28 | """Convert an integer into a base k string.""" 29 | if not (2 <= k <= max_k_labeled): 30 | raise ValueError("k must be in range [2, %d]: %s" 31 | % (max_k_labeled, k)) 32 | return ((x == 0) and numerals[0]) or \ 33 | (base10_integer_to_basek_string(k, x // k).\ 34 | lstrip(numerals[0]) + numerals[x % k]) 35 | 36 | def basek_string_to_base10_integer(k, x): 37 | """Convert a base k string into an integer.""" 38 | assert 1 < k <= max_k_labeled 39 | return sum(numeral_index[c]*(k**i) 40 | for i, c in enumerate(reversed(x))) 41 | 42 | # _clib defines a faster version of this function 43 | def calculate_bucket_level(k, b): 44 | """ 45 | Calculate the level in which a 0-based bucket 46 | lives inside of a k-ary heap. 47 | """ 48 | assert k >= 2 49 | if k == 2: 50 | return log2floor(b+1) 51 | v = (k - 1) * (b + 1) + 1 52 | h = 0 53 | while k**(h+1) < v: 54 | h += 1 55 | return h 56 | 57 | # _clib defines a faster version of this function 58 | def calculate_last_common_level(k, b1, b2): 59 | """ 60 | Calculate the highest level after which the 61 | paths from the root to these buckets diverge. 62 | """ 63 | l1 = calculate_bucket_level(k, b1) 64 | l2 = calculate_bucket_level(k, b2) 65 | while l1 > l2: 66 | b1 = (b1-1)//k 67 | l1 -= 1 68 | while l2 > l1: 69 | b2 = (b2-1)//k 70 | l2 -= 1 71 | while b1 != b2: 72 | b1 = (b1-1)//k 73 | b2 = (b2-1)//k 74 | l1 -= 1 75 | return l1 76 | 77 | def calculate_necessary_heap_height(k, n): 78 | """ 79 | Calculate the necessary k-ary heap height 80 | to store n buckets. 81 | """ 82 | assert n >= 1 83 | return calculate_bucket_level(k, n-1) 84 | 85 | def calculate_bucket_count_in_heap_with_height(k, h): 86 | """ 87 | Calculate the number of buckets in a 88 | k-ary heap of height h. 89 | """ 90 | assert h >= 0 91 | return ((k**(h+1)) - 1) // (k - 1) 92 | 93 | def calculate_bucket_count_in_heap_at_level(k, l): 94 | """ 95 | Calculate the number of buckets in a 96 | k-ary heap at level l. 97 | """ 98 | assert l >= 0 99 | return k**l 100 | 101 | def calculate_leaf_bucket_count_in_heap_with_height(k, h): 102 | """ 103 | Calculate the number of buckets in the 104 | leaf-level of a k-ary heap of height h. 105 | """ 106 | return calculate_bucket_count_in_heap_at_level(k, h) 107 | 108 | def create_node_type(k): 109 | 110 | class VirtualHeapNode(object): 111 | __slots__ = ("bucket", "level") 112 | def __init__(self, bucket): 113 | assert bucket >= 0 114 | self.bucket = bucket 115 | self.level = _clib.calculate_bucket_level(self.k, self.bucket) 116 | 117 | def __hash__(self): 118 | return self.bucket.__hash__() 119 | def __int__(self): 120 | return self.bucket 121 | def __lt__(self, other): 122 | return self.bucket < other 123 | def __le__(self, other): 124 | return self.bucket <= other 125 | def __eq__(self, other): 126 | return self.bucket == other 127 | def __ne__(self, other): 128 | return self.bucket != other 129 | def __gt__(self, other): 130 | return self.bucket > other 131 | def __ge__(self, other): 132 | return self.bucket >= other 133 | def last_common_level(self, n): 134 | return _clib.calculate_last_common_level(self.k, 135 | self.bucket, 136 | n.bucket) 137 | def child_node(self, c): 138 | assert type(c) is int 139 | assert 0 <= c < self.k 140 | return VirtualHeapNode(self.k * self.bucket + 1 + c) 141 | def parent_node(self): 142 | if self.bucket != 0: 143 | return VirtualHeapNode((self.bucket - 1)//self.k) 144 | return None 145 | def ancestor_node_at_level(self, level): 146 | if level > self.level: 147 | return None 148 | current = self 149 | while current.level != level: 150 | current = current.parent_node() 151 | return current 152 | def path_to_root(self): 153 | bucket = self.bucket 154 | yield self 155 | while bucket != 0: 156 | bucket = (bucket - 1)//self.k 157 | yield type(self)(bucket) 158 | def path_from_root(self): 159 | return list(reversed(list(self.path_to_root()))) 160 | def bucket_path_to_root(self): 161 | bucket = self.bucket 162 | yield bucket 163 | while bucket != 0: 164 | bucket = (bucket - 1)//self.k 165 | yield bucket 166 | def bucket_path_from_root(self): 167 | return list(reversed(list(self.bucket_path_to_root()))) 168 | 169 | # 170 | # Expensive Functions 171 | # 172 | def __repr__(self): 173 | try: 174 | label = self.label() 175 | except ValueError: 176 | # presumably, k is too large 177 | label = "" 178 | return ("VirtualHeapNode(k=%s, bucket=%s, level=%s, label=%r)" 179 | % (self.k, self.bucket, self.level, label)) 180 | def __str__(self): 181 | """Returns a tuple (, ).""" 182 | if self.bucket != 0: 183 | return ("(%s, %s)" 184 | % (self.level, 185 | self.bucket - 186 | calculate_bucket_count_in_heap_with_height(self.k, 187 | self.level-1))) 188 | assert self.level == 0 189 | return "(0, 0)" 190 | 191 | def label(self): 192 | assert 0 <= self.bucket 193 | if self.level == 0: 194 | return '' 195 | b_offset = self.bucket - \ 196 | calculate_bucket_count_in_heap_with_height(self.k, 197 | self.level-1) 198 | basek = base10_integer_to_basek_string(self.k, b_offset) 199 | return basek.zfill(self.level) 200 | 201 | def is_node_on_path(self, n): 202 | if n.level <= self.level: 203 | n_label = n.label() 204 | if n_label == "": 205 | return True 206 | return self.label().startswith(n_label) 207 | return False 208 | 209 | VirtualHeapNode.k = k 210 | 211 | return VirtualHeapNode 212 | 213 | class VirtualHeap(object): 214 | 215 | clib = _clib 216 | random = random.SystemRandom() 217 | 218 | def __init__(self, k, blocks_per_bucket=1): 219 | assert 1 < k 220 | assert blocks_per_bucket >= 1 221 | self._k = k 222 | self._blocks_per_bucket = blocks_per_bucket 223 | self.Node = create_node_type(k) 224 | 225 | @property 226 | def k(self): 227 | return self._k 228 | 229 | def node_label_to_bucket(self, label): 230 | if len(label) > 0: 231 | return \ 232 | (calculate_bucket_count_in_heap_with_height(self.k, 233 | len(label)-1) + 234 | basek_string_to_base10_integer(self.k, label)) 235 | return 0 236 | 237 | # 238 | # Buckets (0-based integer, equivalent to block for heap 239 | # with blocks_per_bucket=1) 240 | # 241 | 242 | @property 243 | def blocks_per_bucket(self): 244 | return self._blocks_per_bucket 245 | 246 | def bucket_count_at_level(self, l): 247 | return calculate_bucket_count_in_heap_at_level(self.k, l) 248 | def first_bucket_at_level(self, l): 249 | if l > 0: 250 | return calculate_bucket_count_in_heap_with_height(self.k, l-1) 251 | return 0 252 | def last_bucket_at_level(self, l): 253 | return calculate_bucket_count_in_heap_with_height(self.k, l) - 1 254 | def random_bucket_up_to_level(self, l): 255 | return self.random.randint(self.first_bucket_at_level(0), 256 | self.last_bucket_at_level(l)) 257 | def random_bucket_at_level(self, l): 258 | return self.random.randint(self.first_bucket_at_level(l), 259 | self.first_bucket_at_level(l+1)-1) 260 | 261 | # 262 | # Nodes (a class that helps with heap path calculations) 263 | # 264 | 265 | def root_node(self): 266 | return self.first_node_at_level(0) 267 | def node_count_at_level(self, l): 268 | return self.bucket_count_at_level(l) 269 | def first_node_at_level(self, l): 270 | return self.Node(self.first_bucket_at_level(l)) 271 | def last_node_at_level(self, l): 272 | return self.Node(self.last_bucket_at_level(l)) 273 | def random_node_up_to_level(self, l): 274 | return self.Node(self.random_bucket_up_to_level(l)) 275 | def random_node_at_level(self, l): 276 | return self.Node(self.random_bucket_at_level(l)) 277 | 278 | # 279 | # Block (0-based integer) 280 | # 281 | 282 | def bucket_to_block(self, b): 283 | assert b >= 0 284 | return b * self.blocks_per_bucket 285 | def block_to_bucket(self, s): 286 | assert s >= 0 287 | return s//self.blocks_per_bucket 288 | def first_block_in_bucket(self, b): 289 | return self.bucket_to_block(b) 290 | def last_block_in_bucket(self, b): 291 | return self.bucket_to_block(b) + self.blocks_per_bucket - 1 292 | def block_count_at_level(self, l): 293 | return self.bucket_count_at_level(l) * self.blocks_per_bucket 294 | def first_block_at_level(self, l): 295 | return self.bucket_to_block(self.first_bucket_at_level(l)) 296 | def last_block_at_level(self, l): 297 | return self.bucket_to_block(self.first_bucket_at_level(l+1)) - 1 298 | 299 | class SizedVirtualHeap(VirtualHeap): 300 | 301 | def __init__(self, k, height, blocks_per_bucket=1): 302 | super(SizedVirtualHeap, self).\ 303 | __init__(k, blocks_per_bucket=blocks_per_bucket) 304 | self._height = height 305 | 306 | # 307 | # Size properties 308 | # 309 | @property 310 | def height(self): 311 | return self._height 312 | @property 313 | def levels(self): 314 | return self.height + 1 315 | @property 316 | def first_level(self): 317 | return 0 318 | @property 319 | def last_level(self): 320 | return self.height 321 | 322 | # 323 | # Buckets (0-based integer, equivalent to block for heap 324 | # with blocks_per_bucket=1) 325 | # 326 | 327 | def bucket_count(self): 328 | return calculate_bucket_count_in_heap_with_height(self.k, 329 | self.height) 330 | def leaf_bucket_count(self): 331 | return calculate_leaf_bucket_count_in_heap_with_height(self.k, 332 | self.height) 333 | def first_leaf_bucket(self): 334 | return self.first_bucket_at_level(self.height) 335 | def last_leaf_bucket(self): 336 | return self.last_bucket_at_level(self.height) 337 | def random_bucket(self): 338 | return self.random.randint(self.first_bucket_at_level(0), 339 | self.last_leaf_bucket()) 340 | def random_leaf_bucket(self): 341 | return self.random_bucket_at_level(self.height) 342 | 343 | # 344 | # Nodes (a class that helps with heap path calculations) 345 | # 346 | 347 | def is_nil_node(self, n): 348 | return n.bucket >= self.bucket_count() 349 | def node_count(self): 350 | return self.bucket_count() 351 | def leaf_node_count(self): 352 | return self.leaf_bucket_count() 353 | def first_leaf_node(self): 354 | return self.Node(self.first_leaf_bucket()) 355 | def last_leaf_node(self): 356 | return self.Node(self.last_leaf_bucket()) 357 | def random_leaf_node(self): 358 | return self.Node(self.random_leaf_bucket()) 359 | def random_node(self): 360 | return self.Node(self.random_bucket()) 361 | 362 | # 363 | # Block (0-based integer) 364 | # 365 | 366 | def block_count(self): 367 | return self.bucket_count() * self.blocks_per_bucket 368 | def leaf_block_count(self): 369 | return self.leaf_bucket_count() * self.blocks_per_bucket 370 | def first_leaf_block(self): 371 | return self.first_block_in_bucket(self.first_leaf_bucket()) 372 | def last_leaf_block(self): 373 | return self.last_block_in_bucket(self.last_leaf_bucket()) 374 | 375 | # 376 | # Visualization 377 | # 378 | 379 | def write_as_dot(self, f, data=None, max_levels=None): 380 | "Write the tree in the dot language format to f." 381 | assert (max_levels is None) or (max_levels >= 0) 382 | def visit_node(n, levels): 383 | lbl = "{" 384 | if data is None: 385 | if self.k <= max_k_labeled: 386 | lbl = repr(n.label()).\ 387 | replace("{","\{").\ 388 | replace("}","\}").\ 389 | replace("|","\|").\ 390 | replace("<","\<").\ 391 | replace(">","\>") 392 | else: 393 | lbl = str(n) 394 | else: 395 | s = self.bucket_to_block(n.bucket) 396 | for i in xrange(self.blocks_per_bucket): 397 | lbl += "{%s}" % (data[s+i]) 398 | if i + 1 != self.blocks_per_bucket: 399 | lbl += "|" 400 | lbl += "}" 401 | f.write(" %s [penwidth=%s,label=\"%s\"];\n" 402 | % (n.bucket, 1, lbl)) 403 | levels += 1 404 | if (max_levels is None) or (levels <= max_levels): 405 | for i in xrange(self.k): 406 | cn = n.child_node(i) 407 | if not self.is_nil_node(cn): 408 | visit_node(cn, levels) 409 | f.write(" %s -> %s ;\n" % (n.bucket, cn.bucket)) 410 | 411 | f.write("// Created by SizedVirtualHeap.write_as_dot(...)\n") 412 | f.write("digraph heaptree {\n") 413 | f.write("node [shape=record]\n") 414 | 415 | if (max_levels is None) or (max_levels > 0): 416 | visit_node(self.root_node(), 1) 417 | f.write("}\n") 418 | 419 | def save_image_as_pdf(self, filename, data=None, max_levels=None): 420 | "Write the heap as PDF file." 421 | assert (max_levels is None) or (max_levels >= 0) 422 | import os 423 | if not filename.endswith('.pdf'): 424 | filename = filename+'.pdf' 425 | tmpfd, tmpname = tempfile.mkstemp(suffix='dot') 426 | with open(tmpname, 'w') as f: 427 | self.write_as_dot(f, data=data, max_levels=max_levels) 428 | os.close(tmpfd) 429 | try: 430 | subprocess.call(['dot', 431 | tmpname, 432 | '-Tpdf', 433 | '-o', 434 | ('%s'%filename)]) 435 | except OSError: 436 | sys.stderr.write( 437 | "DOT -> PDF conversion failed. See DOT file: %s\n" 438 | % (tmpname)) 439 | return False 440 | os.remove(tmpname) 441 | return True 442 | --------------------------------------------------------------------------------