def compute_storage_size(cls, block_size, block_count, aes_mode='ctr', storage_type='file', ignore_header=False, **kwds): assert (block_size > 0) and (block_size == int(block_size)) assert (block_count > 0) and (block_count == int(block_count)) assert aes_mode in ('ctr', 'gcm') if not isinstance(storage_type, BlockStorageInterface): storage_type = BlockStorageTypeFactory(storage_type) if aes_mode == 'ctr': extra_block_data = AES.block_size else: assert aes_mode == 'gcm' extra_block_data = 2 * AES.block_size if ignore_header: return (extra_block_data * block_count) + \ storage_type.compute_storage_size( block_size, block_count, ignore_header=True, **kwds) else: return cls._index_offset + \ 2 * AES.block_size + \ (extra_block_data * block_count) + \ storage_type.compute_storage_size( block_size, block_count, ignore_header=False, **kwds)
def compute_storage_size(cls, block_size, block_count, aes_mode='ctr', storage_type='file', ignore_header=False, **kwds): assert (block_size > 0) and (block_size == int(block_size)) assert (block_count > 0) and (block_count == int(block_count)) assert aes_mode in ('ctr', 'gcm') if not isinstance(storage_type, BlockStorageInterface): storage_type = BlockStorageTypeFactory(storage_type) if aes_mode == 'ctr': extra_block_data = AES.block_size else: assert aes_mode == 'gcm' extra_block_data = 2 * AES.block_size if ignore_header: return (extra_block_data * block_count) + \ storage_type.compute_storage_size( block_size, block_count, ignore_header=True, **kwds) else: return cls._index_offset + \ 2 * AES.block_size + \ (extra_block_data * block_count) + \ storage_type.compute_storage_size( block_size, block_count, ignore_header=False, **kwds)
def test_setup(self): fname = ".".join(self.id().split(".")[1:]) fname += ".bin" fname = os.path.join(thisdir, fname) if os.path.exists(fname): os.remove(fname) # pragma: no cover bsize = 10 blocks_per_bucket = 3 fsetup = EncryptedHeapStorage.setup( fname, bsize, self._heap_height, heap_base=self._heap_base, storage_type=self._storage_type, blocks_per_bucket=blocks_per_bucket) fsetup.close() self.assertEqual(type(fsetup.raw_storage), BlockStorageTypeFactory(self._storage_type)) with open(fname, 'rb') as f: flen = len(f.read()) self.assertEqual( flen, TopCachedEncryptedHeapStorage.compute_storage_size( bsize, self._heap_height, heap_base=self._heap_base, blocks_per_bucket=blocks_per_bucket)) self.assertEqual( flen > TopCachedEncryptedHeapStorage.compute_storage_size( bsize, self._heap_height, heap_base=self._heap_base, blocks_per_bucket=blocks_per_bucket, ignore_header=True), True) with TopCachedEncryptedHeapStorage( EncryptedHeapStorage( fname, key=fsetup.key, storage_type=self._storage_type), **self._init_kwds) as f: self.assertEqual(f.header_data, bytes()) self.assertEqual(fsetup.header_data, bytes()) self.assertEqual(f.key, fsetup.key) self.assertEqual(f.blocks_per_bucket, blocks_per_bucket) self.assertEqual(fsetup.blocks_per_bucket, blocks_per_bucket) self.assertEqual(f.bucket_count, (self._heap_base**(self._heap_height+1) - 1)//(self._heap_base-1)) self.assertEqual(fsetup.bucket_count, (self._heap_base**(self._heap_height+1) - 1)//(self._heap_base-1)) self.assertEqual(f.bucket_size, bsize * blocks_per_bucket) self.assertEqual(fsetup.bucket_size, bsize * blocks_per_bucket) self.assertEqual(f.storage_name, fname) self.assertEqual(fsetup.storage_name, fname) os.remove(fname)
def test_setup_withdata(self): fname = ".".join(self.id().split(".")[1:]) fname += ".bin" fname = os.path.join(thisdir, fname) if os.path.exists(fname): os.remove(fname) # pragma: no cover bsize = 10 heap_height = 2 blocks_per_bucket = 1 header_data = bytes(bytearray([0, 1, 2])) fsetup = EncryptedHeapStorage.setup( fname, bsize, heap_height, key_size=AES.key_sizes[0], blocks_per_bucket=blocks_per_bucket, header_data=header_data) fsetup.close() self.assertEqual(type(fsetup.raw_storage), BlockStorageTypeFactory(self._type_name)) with open(fname, 'rb') as f: flen = len(f.read()) self.assertEqual( flen, EncryptedHeapStorage.compute_storage_size( bsize, heap_height, header_data=header_data)) self.assertTrue(len(header_data) > 0) self.assertEqual( EncryptedHeapStorage.compute_storage_size( bsize, heap_height, storage_type=self._type_name) < EncryptedHeapStorage.compute_storage_size( bsize, heap_height, storage_type=self._type_name, header_data=header_data), True) self.assertEqual( flen > EncryptedHeapStorage.compute_storage_size( bsize, heap_height, storage_type=self._type_name, header_data=header_data, ignore_header=True), True) with EncryptedHeapStorage(fname, key=fsetup.key, storage_type=self._type_name) as f: self.assertEqual(f.header_data, header_data) self.assertEqual(fsetup.header_data, header_data) self.assertEqual(f.key, fsetup.key) self.assertEqual(f.blocks_per_bucket, blocks_per_bucket) self.assertEqual(fsetup.blocks_per_bucket, blocks_per_bucket) self.assertEqual(f.bucket_count, 2**(heap_height + 1) - 1) self.assertEqual(fsetup.bucket_count, 2**(heap_height + 1) - 1) self.assertEqual(f.bucket_size, bsize * blocks_per_bucket) self.assertEqual(fsetup.bucket_size, bsize * blocks_per_bucket) self.assertEqual(f.storage_name, fname) self.assertEqual(fsetup.storage_name, fname) os.remove(fname)
def __init__(self, storage, **kwds): self._key = kwds.pop('key', None) if self._key is None: raise ValueError("An encryption key is required using " "the 'key' keyword.") if isinstance(storage, BlockStorageInterface): storage_owned = False self._storage = storage if len(kwds): raise ValueError("Keywords not used when initializing " "with a storage device: %s" % (str(kwds))) else: storage_owned = True storage_type = kwds.pop('storage_type', 'file') self._storage = \ BlockStorageTypeFactory(storage_type)(storage, **kwds) try: header_data = AES.GCMDec(self._key, self._storage.header_data) (self._ismodegcm, ) = struct.unpack( self._index_struct_string, header_data[:self._index_offset]) self._verify_digest = header_data[:hashlib.sha384().digest_size] verify = hmac.HMAC(key=self.key, msg=struct.pack(self._verify_struct_string, self._storage.block_size, self._storage.block_count, len(self._storage.header_data)), digestmod=hashlib.sha384) if verify.digest() != self._verify_digest: raise ValueError("HMAC of plaintext index data does not match") if self._ismodegcm: self._encrypt_block_func = AES.GCMEnc self._decrypt_block_func = AES.GCMDec else: self._encrypt_block_func = AES.CTREnc self._decrypt_block_func = AES.CTRDec except: if storage_owned: self._storage.close() raise
def __init__(self, storage, **kwds): self._key = kwds.pop('key', None) if self._key is None: raise ValueError( "An encryption key is required using " "the 'key' keyword.") if isinstance(storage, BlockStorageInterface): storage_owned = False self._storage = storage if len(kwds): raise ValueError( "Keywords not used when initializing " "with a storage device: %s" % (str(kwds))) else: storage_owned = True storage_type = kwds.pop('storage_type', 'file') self._storage = \ BlockStorageTypeFactory(storage_type)(storage, **kwds) try: header_data = AES.GCMDec(self._key, self._storage.header_data) (self._ismodegcm,) = struct.unpack( self._index_struct_string, header_data[:self._index_offset]) self._verify_digest = header_data[:hashlib.sha384().digest_size] verify = hmac.HMAC( key=self.key, msg=struct.pack(self._verify_struct_string, self._storage.block_size, self._storage.block_count, len(self._storage.header_data)), digestmod=hashlib.sha384) if verify.digest() != self._verify_digest: raise ValueError( "HMAC of plaintext index data does not match") if self._ismodegcm: self._encrypt_block_func = AES.GCMEnc self._decrypt_block_func = AES.GCMDec else: self._encrypt_block_func = AES.CTREnc self._decrypt_block_func = AES.CTRDec except: if storage_owned: self._storage.close() raise
def __init__(self, storage, **kwds): if isinstance(storage, BlockStorageInterface): self._storage = storage if len(kwds): raise ValueError( "Keywords not used when initializing " "with a storage device: %s" % (str(kwds))) else: storage_type = kwds.pop('storage_type', 'file') self._storage = BlockStorageTypeFactory(storage_type)\ (storage, **kwds) heap_base, heap_height, blocks_per_bucket = \ struct.unpack( self._header_struct_string, self._storage.header_data[:self._header_offset]) self._vheap = SizedVirtualHeap( heap_base, heap_height, blocks_per_bucket=blocks_per_bucket)
def test_setup(self): fname = ".".join(self.id().split(".")[1:]) fname += ".bin" fname = os.path.join(thisdir, fname) if os.path.exists(fname): os.remove(fname) # pragma: no cover bsize = 10 heap_height = 2 blocks_per_bucket = 3 fsetup = HeapStorage.setup(fname, bsize, heap_height, blocks_per_bucket=blocks_per_bucket) fsetup.close() self.assertEqual(type(fsetup.bucket_storage), BlockStorageTypeFactory(self._type_name)) with open(fname, 'rb') as f: flen = len(f.read()) self.assertEqual( flen, HeapStorage.compute_storage_size( bsize, heap_height, blocks_per_bucket=blocks_per_bucket)) self.assertEqual( flen > HeapStorage.compute_storage_size( bsize, heap_height, blocks_per_bucket=blocks_per_bucket, ignore_header=True), True) with HeapStorage(fname, storage_type=self._type_name) as f: self.assertEqual(f.header_data, bytes()) self.assertEqual(fsetup.header_data, bytes()) self.assertEqual(f.blocks_per_bucket, blocks_per_bucket) self.assertEqual(fsetup.blocks_per_bucket, blocks_per_bucket) self.assertEqual(f.bucket_count, 2**(heap_height + 1) - 1) self.assertEqual(fsetup.bucket_count, 2**(heap_height + 1) - 1) self.assertEqual(f.bucket_size, bsize * blocks_per_bucket) self.assertEqual(fsetup.bucket_size, bsize * blocks_per_bucket) self.assertEqual(f.storage_name, fname) self.assertEqual(fsetup.storage_name, fname) os.remove(fname)
class EncryptedBlockStorage(EncryptedBlockStorageInterface): _index_struct_string = "!" + ("x" * hashlib.sha384().digest_size) + "?" _index_offset = struct.calcsize(_index_struct_string) _verify_struct_string = "!LLL" _verify_size = struct.calcsize(_verify_struct_string) def __init__(self, storage, **kwds): self._key = kwds.pop('key', None) if self._key is None: raise ValueError("An encryption key is required using " "the 'key' keyword.") if isinstance(storage, BlockStorageInterface): storage_owned = False self._storage = storage if len(kwds): raise ValueError("Keywords not used when initializing " "with a storage device: %s" % (str(kwds))) else: storage_owned = True storage_type = kwds.pop('storage_type', 'file') self._storage = \ BlockStorageTypeFactory(storage_type)(storage, **kwds) try: header_data = AES.GCMDec(self._key, self._storage.header_data) (self._ismodegcm, ) = struct.unpack( self._index_struct_string, header_data[:self._index_offset]) self._verify_digest = header_data[:hashlib.sha384().digest_size] verify = hmac.HMAC(key=self.key, msg=struct.pack(self._verify_struct_string, self._storage.block_size, self._storage.block_count, len(self._storage.header_data)), digestmod=hashlib.sha384) if verify.digest() != self._verify_digest: raise ValueError("HMAC of plaintext index data does not match") if self._ismodegcm: self._encrypt_block_func = AES.GCMEnc self._decrypt_block_func = AES.GCMDec else: self._encrypt_block_func = AES.CTREnc self._decrypt_block_func = AES.CTRDec except: if storage_owned: self._storage.close() raise # # Define EncryptedBlockStorageInterface Methods # @property def key(self): return self._key @property def raw_storage(self): return self._storage # # Define BlockStorageInterface Methods # def clone_device(self): return EncryptedBlockStorage(self._storage.clone_device(), key=self.key) @classmethod def compute_storage_size(cls, block_size, block_count, aes_mode='ctr', storage_type='file', ignore_header=False, **kwds): assert (block_size > 0) and (block_size == int(block_size)) assert (block_count > 0) and (block_count == int(block_count)) assert aes_mode in ('ctr', 'gcm') if not isinstance(storage_type, BlockStorageInterface): storage_type = BlockStorageTypeFactory(storage_type) if aes_mode == 'ctr': extra_block_data = AES.block_size else: assert aes_mode == 'gcm' extra_block_data = 2 * AES.block_size if ignore_header: return (extra_block_data * block_count) + \ storage_type.compute_storage_size( block_size, block_count, ignore_header=True, **kwds) else: return cls._index_offset + \ 2 * AES.block_size + \ (extra_block_data * block_count) + \ storage_type.compute_storage_size( block_size, block_count, ignore_header=False, **kwds) @classmethod def setup(cls, storage_name, block_size, block_count, aes_mode='ctr', key_size=None, key=None, storage_type='file', initialize=None, **kwds): if (key is not None) and (key_size is not None): raise ValueError("Only one of 'key' or 'keysize' keywords can " "be specified at a time") if key is None: if key_size is None: key_size = 32 if key_size not in AES.key_sizes: raise ValueError("Invalid key size: %s" % (key_size)) key = AES.KeyGen(key_size) else: if len(key) not in AES.key_sizes: raise ValueError("Invalid key size: %s" % (len(key))) if (block_size <= 0) or (block_size != int(block_size)): raise ValueError( "Block size (bytes) must be a positive integer: %s" % (block_size)) ismodegcm = None encrypt_block_func = None encrypted_block_size = block_size if aes_mode == 'ctr': ismodegcm = False encrypt_block_func = AES.CTREnc encrypted_block_size += AES.block_size elif aes_mode == 'gcm': ismodegcm = True encrypt_block_func = AES.GCMEnc encrypted_block_size += (2 * AES.block_size) else: raise ValueError( "AES encryption mode must be one of 'ctr' or 'gcm'. " "Invalid value: %s" % (aes_mode)) assert ismodegcm is not None assert encrypt_block_func is not None if not isinstance(storage_type, BlockStorageInterface): storage_type = BlockStorageTypeFactory(storage_type) if initialize is None: zeros = bytes(bytearray(block_size)) initialize = lambda i: zeros def encrypted_initialize(i): return encrypt_block_func(key, initialize(i)) kwds['initialize'] = encrypted_initialize user_header_data = kwds.get('header_data', bytes()) if type(user_header_data) is not bytes: raise TypeError("'header_data' must be of type bytes. " "Invalid type: %s" % (type(user_header_data))) # we generate the first time simply to # compute the length tmp = hmac.HMAC(key=key, msg=struct.pack(cls._verify_struct_string, encrypted_block_size, block_count, 0), digestmod=hashlib.sha384).digest() header_data = bytearray( struct.pack(cls._index_struct_string, ismodegcm)) header_data[:hashlib.sha384().digest_size] = tmp header_data = header_data + user_header_data header_data = AES.GCMEnc(key, bytes(header_data)) # now that we know the length of the header data # being sent to the underlying storage we can # compute the real hmac verify_digest = hmac.HMAC(key=key, msg=struct.pack(cls._verify_struct_string, encrypted_block_size, block_count, len(header_data)), digestmod=hashlib.sha384).digest() header_data = bytearray( struct.pack(cls._index_struct_string, ismodegcm)) header_data[:hashlib.sha384().digest_size] = verify_digest header_data = header_data + user_header_data kwds['header_data'] = AES.GCMEnc(key, bytes(header_data)) return EncryptedBlockStorage(storage_type.setup( storage_name, encrypted_block_size, block_count, **kwds), key=key) @property def header_data(self): return AES.GCMDec(self._key, self._storage.header_data)\ [self._index_offset:] @property def block_count(self): return self._storage.block_count @property def block_size(self): if self._ismodegcm: return self._storage.block_size - 2 * AES.block_size else: return self._storage.block_size - AES.block_size @property def storage_name(self): return self._storage.storage_name def update_header_data(self, new_header_data): self._storage.update_header_data( AES.GCMEnc( self.key, AES.GCMDec(self._key, self._storage.header_data)\ [:self._index_offset] + \ new_header_data)) def close(self): self._storage.close() def read_block(self, i): return self._decrypt_block_func(self._key, self._storage.read_block(i)) def read_blocks(self, indices, *args, **kwds): return [ self._decrypt_block_func(self._key, b) for b in self._storage.read_blocks(indices, *args, **kwds) ] def yield_blocks(self, indices, *args, **kwds): for b in self._storage.yield_blocks(indices, *args, **kwds): yield self._decrypt_block_func(self._key, b) def write_block(self, i, block, *args, **kwds): self._storage.write_block(i, self._encrypt_block_func(self._key, block), *args, **kwds) def write_blocks(self, indices, blocks, *args, **kwds): enc_blocks = [] for i, b in zip(indices, blocks): enc_blocks.append(self._encrypt_block_func(self._key, b)) self._storage.write_blocks(indices, enc_blocks, *args, **kwds) @property def bytes_sent(self): return self._storage.bytes_sent @property def bytes_received(self): return self._storage.bytes_received
# # We must cast from bytearray to bytes # when reading from a bytearray so that this # class works with the encryption layer. # def read_blocks(self, indices): return [bytes(block) for block in super(BlockStorageRAM, self).read_blocks(indices)] def yield_blocks(self, indices): for block in super(BlockStorageRAM, self).yield_blocks(indices): yield bytes(block) def read_block(self, i): return bytes(super(BlockStorageRAM, self).read_block(i)) #def write_blocks(...) #def write_block(...) @property def bytes_sent(self): return self._bytes_sent @property def bytes_received(self): return self._bytes_received BlockStorageTypeFactory.register_device("ram", BlockStorageRAM)
#def update_header_data(...) def close(self): self._prep_for_close() if self._f is not None: if self._mmap_owned: try: self._f.close() except OSError: # pragma: no cover pass # pragma: no cover self._f = None #def read_blocks(...) #def yield_blocks(...) #def read_block(...) #def write_blocks(...) #def write_block(...) #@property #def bytes_sent(...) #@property #def bytes_received(...) BlockStorageTypeFactory.register_device("mmap", BlockStorageMMap)
assert 0 <= i < self.block_count self._bytes_received += self.block_size return self._download(i) def write_blocks(self, indices, blocks, callback=None): self._check_async() # be sure not to exhaust this if it is an iterator # or generator indices = list(indices) assert all(0 <= i <= self.block_count for i in indices) self._bytes_sent += self.block_size * len(indices) indices = (self._basename % i for i in indices) self._schedule_async_write(zip(indices, blocks), callback=callback) def write_block(self, i, block): self._check_async() assert 0 <= i < self.block_count self._bytes_sent += self.block_size self._schedule_async_write((((self._basename % i), block), )) @property def bytes_sent(self): return self._bytes_sent @property def bytes_received(self): return self._bytes_received BlockStorageTypeFactory.register_device("s3", BlockStorageS3)
def test_register_invalid_type(self): with self.assertRaises(TypeError): BlockStorageTypeFactory.register_device('new_str_type', str)
def setup(cls, storage_name, block_size, block_count, aes_mode='ctr', key_size=None, key=None, storage_type='file', initialize=None, **kwds): if (key is not None) and (key_size is not None): raise ValueError("Only one of 'key' or 'keysize' keywords can " "be specified at a time") if key is None: if key_size is None: key_size = 32 if key_size not in AES.key_sizes: raise ValueError("Invalid key size: %s" % (key_size)) key = AES.KeyGen(key_size) else: if len(key) not in AES.key_sizes: raise ValueError("Invalid key size: %s" % (len(key))) if (block_size <= 0) or (block_size != int(block_size)): raise ValueError( "Block size (bytes) must be a positive integer: %s" % (block_size)) ismodegcm = None encrypt_block_func = None encrypted_block_size = block_size if aes_mode == 'ctr': ismodegcm = False encrypt_block_func = AES.CTREnc encrypted_block_size += AES.block_size elif aes_mode == 'gcm': ismodegcm = True encrypt_block_func = AES.GCMEnc encrypted_block_size += (2 * AES.block_size) else: raise ValueError( "AES encryption mode must be one of 'ctr' or 'gcm'. " "Invalid value: %s" % (aes_mode)) assert ismodegcm is not None assert encrypt_block_func is not None if not isinstance(storage_type, BlockStorageInterface): storage_type = BlockStorageTypeFactory(storage_type) if initialize is None: zeros = bytes(bytearray(block_size)) initialize = lambda i: zeros def encrypted_initialize(i): return encrypt_block_func(key, initialize(i)) kwds['initialize'] = encrypted_initialize user_header_data = kwds.get('header_data', bytes()) if type(user_header_data) is not bytes: raise TypeError("'header_data' must be of type bytes. " "Invalid type: %s" % (type(user_header_data))) # we generate the first time simply to # compute the length tmp = hmac.HMAC(key=key, msg=struct.pack(cls._verify_struct_string, encrypted_block_size, block_count, 0), digestmod=hashlib.sha384).digest() header_data = bytearray( struct.pack(cls._index_struct_string, ismodegcm)) header_data[:hashlib.sha384().digest_size] = tmp header_data = header_data + user_header_data header_data = AES.GCMEnc(key, bytes(header_data)) # now that we know the length of the header data # being sent to the underlying storage we can # compute the real hmac verify_digest = hmac.HMAC(key=key, msg=struct.pack(cls._verify_struct_string, encrypted_block_size, block_count, len(header_data)), digestmod=hashlib.sha384).digest() header_data = bytearray( struct.pack(cls._index_struct_string, ismodegcm)) header_data[:hashlib.sha384().digest_size] = verify_digest header_data = header_data + user_header_data kwds['header_data'] = AES.GCMEnc(key, bytes(header_data)) return EncryptedBlockStorage(storage_type.setup( storage_name, encrypted_block_size, block_count, **kwds), key=key)
def test_setup_withdata(self): fname = ".".join(self.id().split(".")[1:]) fname += ".bin" fname = os.path.join(thisdir, fname) if os.path.exists(fname): os.remove(fname) # pragma: no cover bsize = 10 bcount = 11 header_data = bytes(bytearray([0,1,2])) fsetup = PathORAM.setup( fname, block_size=bsize, block_count=bcount, bucket_capacity=self._bucket_capacity, heap_base=self._heap_base, key=self._test_key, key_size=self._test_key_size, aes_mode=self._aes_mode, storage_type=self._type_name, header_data=header_data, **self._kwds) fsetup.close() self.assertEqual(type(fsetup.raw_storage), BlockStorageTypeFactory(self._type_name)) with open(fname, 'rb') as f: flen = len(f.read()) self.assertEqual( flen, PathORAM.compute_storage_size( bsize, bcount, bucket_capacity=self._bucket_capacity, heap_base=self._heap_base, aes_mode=self._aes_mode, storage_type=self._type_name, header_data=header_data)) self.assertTrue(len(header_data) > 0) self.assertEqual( PathORAM.compute_storage_size( bsize, bcount, bucket_capacity=self._bucket_capacity, heap_base=self._heap_base, aes_mode=self._aes_mode, storage_type=self._type_name) < PathORAM.compute_storage_size( bsize, bcount, bucket_capacity=self._bucket_capacity, heap_base=self._heap_base, aes_mode=self._aes_mode, storage_type=self._type_name, header_data=header_data), True) self.assertEqual( flen > PathORAM.compute_storage_size( bsize, bcount, bucket_capacity=self._bucket_capacity, heap_base=self._heap_base, aes_mode=self._aes_mode, storage_type=self._type_name, header_data=header_data, ignore_header=True), True) with PathORAM(fname, fsetup.stash, fsetup.position_map, key=fsetup.key, storage_type=self._type_name, **self._kwds) as f: self.assertEqual(f.header_data, header_data) self.assertEqual(fsetup.header_data, header_data) self.assertEqual(f.key, fsetup.key) self.assertEqual(f.block_size, bsize) self.assertEqual(fsetup.block_size, bsize) self.assertEqual(f.block_count, bcount) self.assertEqual(fsetup.block_count, bcount) self.assertEqual(f.storage_name, fname) self.assertEqual(fsetup.storage_name, fname) os.remove(fname)
#def update_header_data(...) def close(self): self._prep_for_close() if self._f is not None: if self._mmap_owned: try: self._f.close() except OSError: # pragma: no cover pass # pragma: no cover self._f = None #def read_blocks(...) #def yield_blocks(...) #def read_block(...) #def write_blocks(...) #def write_block(...) #@property #def bytes_sent(...) #@property #def bytes_received(...) BlockStorageTypeFactory.register_device("mmap", BlockStorageMMap)
def test_setup(self): fname = ".".join(self.id().split(".")[1:]) fname += ".bin" fname = os.path.join(thisdir, fname) if os.path.exists(fname): os.remove(fname) # pragma: no cover bsize = 10 bcount = 11 fsetup = PathORAM.setup( fname, bsize, bcount, bucket_capacity=self._bucket_capacity, heap_base=self._heap_base, key=self._test_key, key_size=self._test_key_size, aes_mode=self._aes_mode, storage_type=self._type_name, **self._kwds) fsetup.close() self.assertEqual(type(fsetup.raw_storage), BlockStorageTypeFactory(self._type_name)) # test that these can be called with default keyword values fsetup.stash_digest(fsetup.stash) fsetup.position_map_digest(fsetup.position_map) with open(fname, 'rb') as f: flen = len(f.read()) self.assertEqual( flen, PathORAM.compute_storage_size( bsize, bcount, bucket_capacity=self._bucket_capacity, heap_base=self._heap_base, aes_mode=self._aes_mode, storage_type=self._type_name)) self.assertEqual( flen > PathORAM.compute_storage_size( bsize, bcount, bucket_capacity=self._bucket_capacity, heap_base=self._heap_base, aes_mode=self._aes_mode, storage_type=self._type_name, ignore_header=True), True) with PathORAM(fname, fsetup.stash, fsetup.position_map, key=fsetup.key, storage_type=self._type_name, **self._kwds) as f: self.assertEqual(f.header_data, bytes()) self.assertEqual(fsetup.header_data, bytes()) self.assertEqual(f.key, fsetup.key) self.assertEqual(f.block_size, bsize) self.assertEqual(fsetup.block_size, bsize) self.assertEqual(f.block_count, bcount) self.assertEqual(fsetup.block_count, bcount) self.assertEqual(f.storage_name, fname) self.assertEqual(fsetup.storage_name, fname) os.remove(fname)
def test_file(self): self.assertIs(BlockStorageTypeFactory('file'), BlockStorageFile)
def test_mmap(self): self.assertIs(BlockStorageTypeFactory('mmap'), BlockStorageMMap)
def test_register_invalid_name(self): with self.assertRaises(ValueError): BlockStorageTypeFactory.register_device( 's3', BlockStorageFile)
self._f.seek(self._header_offset + i * self.block_size) return self._f.read(self.block_size) def write_blocks(self, indices, blocks, callback=None): self._check_async() chunks = [] for i, block in zip(indices, blocks): assert 0 <= i < self.block_count assert len(block) == self.block_size, \ ("%s != %s" % (len(block), self.block_size)) self._bytes_sent += self.block_size chunks.append((i, block)) self._schedule_async_write(chunks, callback=callback) def write_block(self, i, block): self._check_async() assert 0 <= i < self.block_count assert len(block) == self.block_size self._bytes_sent += self.block_size self._schedule_async_write(((i, block),)) @property def bytes_sent(self): return self._bytes_sent @property def bytes_received(self): return self._bytes_received BlockStorageTypeFactory.register_device("file", BlockStorageFile)
for i in indices: assert 0 <= i < self.block_count self._bytes_received += self.block_size args.append((self._header_offset + i * self.block_size, self.block_size)) return self._f.readv(args) def yield_blocks(self, indices, chunksize=100): for chunk in chunkiter(indices, n=chunksize): assert all(0 <= i <= self.block_count for i in chunk) self._bytes_received += self.block_size * len(chunk) args = [(self._header_offset + i * self.block_size, self.block_size) for i in chunk] for block in self._f.readv(args): yield block #def read_block(...) #def write_blocks(...) #def write_block(...) #@property #def bytes_sent(...) #@property #def bytes_received(...) BlockStorageTypeFactory.register_device("sftp", BlockStorageSFTP)
def setup(cls, storage_name, block_size, block_count, aes_mode='ctr', key_size=None, key=None, storage_type='file', initialize=None, **kwds): if (key is not None) and (key_size is not None): raise ValueError( "Only one of 'key' or 'keysize' keywords can " "be specified at a time") if key is None: if key_size is None: key_size = 32 if key_size not in AES.key_sizes: raise ValueError( "Invalid key size: %s" % (key_size)) key = AES.KeyGen(key_size) else: if len(key) not in AES.key_sizes: raise ValueError( "Invalid key size: %s" % (len(key))) if (block_size <= 0) or (block_size != int(block_size)): raise ValueError( "Block size (bytes) must be a positive integer: %s" % (block_size)) ismodegcm = None encrypt_block_func = None encrypted_block_size = block_size if aes_mode == 'ctr': ismodegcm = False encrypt_block_func = AES.CTREnc encrypted_block_size += AES.block_size elif aes_mode == 'gcm': ismodegcm = True encrypt_block_func = AES.GCMEnc encrypted_block_size += (2 * AES.block_size) else: raise ValueError( "AES encryption mode must be one of 'ctr' or 'gcm'. " "Invalid value: %s" % (aes_mode)) assert ismodegcm is not None assert encrypt_block_func is not None if not isinstance(storage_type, BlockStorageInterface): storage_type = BlockStorageTypeFactory(storage_type) if initialize is None: zeros = bytes(bytearray(block_size)) initialize = lambda i: zeros def encrypted_initialize(i): return encrypt_block_func(key, initialize(i)) kwds['initialize'] = encrypted_initialize user_header_data = kwds.get('header_data', bytes()) if type(user_header_data) is not bytes: raise TypeError( "'header_data' must be of type bytes. " "Invalid type: %s" % (type(user_header_data))) # we generate the first time simply to # compute the length tmp = hmac.HMAC( key=key, msg=struct.pack(cls._verify_struct_string, encrypted_block_size, block_count, 0), digestmod=hashlib.sha384).digest() header_data = bytearray(struct.pack(cls._index_struct_string, ismodegcm)) header_data[:hashlib.sha384().digest_size] = tmp header_data = header_data + user_header_data header_data = AES.GCMEnc(key, bytes(header_data)) # now that we know the length of the header data # being sent to the underlying storage we can # compute the real hmac verify_digest = hmac.HMAC( key=key, msg=struct.pack(cls._verify_struct_string, encrypted_block_size, block_count, len(header_data)), digestmod=hashlib.sha384).digest() header_data = bytearray(struct.pack(cls._index_struct_string, ismodegcm)) header_data[:hashlib.sha384().digest_size] = verify_digest header_data = header_data + user_header_data kwds['header_data'] = AES.GCMEnc(key, bytes(header_data)) return EncryptedBlockStorage( storage_type.setup(storage_name, encrypted_block_size, block_count, **kwds), key=key)
def test_s3(self): self.assertIs(BlockStorageTypeFactory('s3'), BlockStorageS3)
class EncryptedBlockStorage(EncryptedBlockStorageInterface): _index_struct_string = "!"+("x"*hashlib.sha384().digest_size)+"?" _index_offset = struct.calcsize(_index_struct_string) _verify_struct_string = "!LLL" _verify_size = struct.calcsize(_verify_struct_string) def __init__(self, storage, **kwds): self._key = kwds.pop('key', None) if self._key is None: raise ValueError( "An encryption key is required using " "the 'key' keyword.") if isinstance(storage, BlockStorageInterface): storage_owned = False self._storage = storage if len(kwds): raise ValueError( "Keywords not used when initializing " "with a storage device: %s" % (str(kwds))) else: storage_owned = True storage_type = kwds.pop('storage_type', 'file') self._storage = \ BlockStorageTypeFactory(storage_type)(storage, **kwds) try: header_data = AES.GCMDec(self._key, self._storage.header_data) (self._ismodegcm,) = struct.unpack( self._index_struct_string, header_data[:self._index_offset]) self._verify_digest = header_data[:hashlib.sha384().digest_size] verify = hmac.HMAC( key=self.key, msg=struct.pack(self._verify_struct_string, self._storage.block_size, self._storage.block_count, len(self._storage.header_data)), digestmod=hashlib.sha384) if verify.digest() != self._verify_digest: raise ValueError( "HMAC of plaintext index data does not match") if self._ismodegcm: self._encrypt_block_func = AES.GCMEnc self._decrypt_block_func = AES.GCMDec else: self._encrypt_block_func = AES.CTREnc self._decrypt_block_func = AES.CTRDec except: if storage_owned: self._storage.close() raise # # Define EncryptedBlockStorageInterface Methods # @property def key(self): return self._key @property def raw_storage(self): return self._storage # # Define BlockStorageInterface Methods # def clone_device(self): return EncryptedBlockStorage(self._storage.clone_device(), key=self.key) @classmethod def compute_storage_size(cls, block_size, block_count, aes_mode='ctr', storage_type='file', ignore_header=False, **kwds): assert (block_size > 0) and (block_size == int(block_size)) assert (block_count > 0) and (block_count == int(block_count)) assert aes_mode in ('ctr', 'gcm') if not isinstance(storage_type, BlockStorageInterface): storage_type = BlockStorageTypeFactory(storage_type) if aes_mode == 'ctr': extra_block_data = AES.block_size else: assert aes_mode == 'gcm' extra_block_data = 2 * AES.block_size if ignore_header: return (extra_block_data * block_count) + \ storage_type.compute_storage_size( block_size, block_count, ignore_header=True, **kwds) else: return cls._index_offset + \ 2 * AES.block_size + \ (extra_block_data * block_count) + \ storage_type.compute_storage_size( block_size, block_count, ignore_header=False, **kwds) @classmethod def setup(cls, storage_name, block_size, block_count, aes_mode='ctr', key_size=None, key=None, storage_type='file', initialize=None, **kwds): if (key is not None) and (key_size is not None): raise ValueError( "Only one of 'key' or 'keysize' keywords can " "be specified at a time") if key is None: if key_size is None: key_size = 32 if key_size not in AES.key_sizes: raise ValueError( "Invalid key size: %s" % (key_size)) key = AES.KeyGen(key_size) else: if len(key) not in AES.key_sizes: raise ValueError( "Invalid key size: %s" % (len(key))) if (block_size <= 0) or (block_size != int(block_size)): raise ValueError( "Block size (bytes) must be a positive integer: %s" % (block_size)) ismodegcm = None encrypt_block_func = None encrypted_block_size = block_size if aes_mode == 'ctr': ismodegcm = False encrypt_block_func = AES.CTREnc encrypted_block_size += AES.block_size elif aes_mode == 'gcm': ismodegcm = True encrypt_block_func = AES.GCMEnc encrypted_block_size += (2 * AES.block_size) else: raise ValueError( "AES encryption mode must be one of 'ctr' or 'gcm'. " "Invalid value: %s" % (aes_mode)) assert ismodegcm is not None assert encrypt_block_func is not None if not isinstance(storage_type, BlockStorageInterface): storage_type = BlockStorageTypeFactory(storage_type) if initialize is None: zeros = bytes(bytearray(block_size)) initialize = lambda i: zeros def encrypted_initialize(i): return encrypt_block_func(key, initialize(i)) kwds['initialize'] = encrypted_initialize user_header_data = kwds.get('header_data', bytes()) if type(user_header_data) is not bytes: raise TypeError( "'header_data' must be of type bytes. " "Invalid type: %s" % (type(user_header_data))) # we generate the first time simply to # compute the length tmp = hmac.HMAC( key=key, msg=struct.pack(cls._verify_struct_string, encrypted_block_size, block_count, 0), digestmod=hashlib.sha384).digest() header_data = bytearray(struct.pack(cls._index_struct_string, ismodegcm)) header_data[:hashlib.sha384().digest_size] = tmp header_data = header_data + user_header_data header_data = AES.GCMEnc(key, bytes(header_data)) # now that we know the length of the header data # being sent to the underlying storage we can # compute the real hmac verify_digest = hmac.HMAC( key=key, msg=struct.pack(cls._verify_struct_string, encrypted_block_size, block_count, len(header_data)), digestmod=hashlib.sha384).digest() header_data = bytearray(struct.pack(cls._index_struct_string, ismodegcm)) header_data[:hashlib.sha384().digest_size] = verify_digest header_data = header_data + user_header_data kwds['header_data'] = AES.GCMEnc(key, bytes(header_data)) return EncryptedBlockStorage( storage_type.setup(storage_name, encrypted_block_size, block_count, **kwds), key=key) @property def header_data(self): return AES.GCMDec(self._key, self._storage.header_data)\ [self._index_offset:] @property def block_count(self): return self._storage.block_count @property def block_size(self): if self._ismodegcm: return self._storage.block_size - 2 * AES.block_size else: return self._storage.block_size - AES.block_size @property def storage_name(self): return self._storage.storage_name def update_header_data(self, new_header_data): self._storage.update_header_data( AES.GCMEnc( self.key, AES.GCMDec(self._key, self._storage.header_data)\ [:self._index_offset] + \ new_header_data)) def close(self): self._storage.close() def read_block(self, i): return self._decrypt_block_func( self._key, self._storage.read_block(i)) def read_blocks(self, indices, *args, **kwds): return [self._decrypt_block_func(self._key, b) for b in self._storage.read_blocks(indices, *args, **kwds)] def yield_blocks(self, indices, *args, **kwds): for b in self._storage.yield_blocks(indices, *args, **kwds): yield self._decrypt_block_func(self._key, b) def write_block(self, i, block, *args, **kwds): self._storage.write_block( i, self._encrypt_block_func(self._key, block), *args, **kwds) def write_blocks(self, indices, blocks, *args, **kwds): enc_blocks = [] for i, b in zip(indices, blocks): enc_blocks.append( self._encrypt_block_func(self._key, b)) self._storage.write_blocks(indices, enc_blocks, *args, **kwds) @property def bytes_sent(self): return self._storage.bytes_sent @property def bytes_received(self): return self._storage.bytes_received
def test_invalid(self): with self.assertRaises(ValueError): BlockStorageTypeFactory(None)
class HeapStorage(HeapStorageInterface): _header_struct_string = "!LLL" _header_offset = struct.calcsize(_header_struct_string) def _new_storage(self, storage, **kwds): storage_type = kwds.pop('storage_type', 'file') def __init__(self, storage, **kwds): if isinstance(storage, BlockStorageInterface): self._storage = storage if len(kwds): raise ValueError( "Keywords not used when initializing " "with a storage device: %s" % (str(kwds))) else: storage_type = kwds.pop('storage_type', 'file') self._storage = BlockStorageTypeFactory(storage_type)\ (storage, **kwds) heap_base, heap_height, blocks_per_bucket = \ struct.unpack( self._header_struct_string, self._storage.header_data[:self._header_offset]) self._vheap = SizedVirtualHeap( heap_base, heap_height, blocks_per_bucket=blocks_per_bucket) # # Define HeapStorageInterface Methods # def clone_device(self): return HeapStorage(self._storage.clone_device()) @classmethod def compute_storage_size(cls, block_size, heap_height, blocks_per_bucket=1, heap_base=2, ignore_header=False, storage_type='file', **kwds): assert (block_size > 0) and (block_size == int(block_size)) assert heap_height >= 0 assert blocks_per_bucket >= 1 assert heap_base >= 2 assert 'block_count' not in kwds vheap = SizedVirtualHeap( heap_base, heap_height, blocks_per_bucket=blocks_per_bucket) if ignore_header: return BlockStorageTypeFactory(storage_type).\ compute_storage_size( vheap.blocks_per_bucket * block_size, vheap.bucket_count(), ignore_header=True, **kwds) else: return cls._header_offset + \ BlockStorageTypeFactory(storage_type).\ compute_storage_size( vheap.blocks_per_bucket * block_size, vheap.bucket_count(), ignore_header=False, **kwds) @classmethod def setup(cls, storage_name, block_size, heap_height, blocks_per_bucket=1, heap_base=2, storage_type='file', **kwds): if 'block_count' in kwds: raise ValueError("'block_count' keyword is not accepted") if heap_height < 0: raise ValueError( "heap height must be 0 or greater. Invalid value: %s" % (heap_height)) if blocks_per_bucket < 1: raise ValueError( "blocks_per_bucket must be 1 or greater. " "Invalid value: %s" % (blocks_per_bucket)) if heap_base < 2: raise ValueError( "heap base must be 2 or greater. Invalid value: %s" % (heap_base)) vheap = SizedVirtualHeap( heap_base, heap_height, blocks_per_bucket=blocks_per_bucket) user_header_data = kwds.pop('header_data', bytes()) if type(user_header_data) is not bytes: raise TypeError( "'header_data' must be of type bytes. " "Invalid type: %s" % (type(user_header_data))) kwds['header_data'] = \ struct.pack(cls._header_struct_string, heap_base, heap_height, blocks_per_bucket) + \ user_header_data return HeapStorage( BlockStorageTypeFactory(storage_type).setup( storage_name, vheap.blocks_per_bucket * block_size, vheap.bucket_count(), **kwds)) @property def header_data(self): return self._storage.header_data[self._header_offset:] @property def bucket_count(self): return self._storage.block_count @property def bucket_size(self): return self._storage.block_size @property def blocks_per_bucket(self): return self._vheap.blocks_per_bucket @property def storage_name(self): return self._storage.storage_name @property def virtual_heap(self): return self._vheap @property def bucket_storage(self): return self._storage def update_header_data(self, new_header_data): self._storage.update_header_data( self._storage.header_data[:self._header_offset] + \ new_header_data) def close(self): self._storage.close() def read_path(self, b, level_start=0): assert 0 <= b < self._vheap.bucket_count() bucket_list = self._vheap.Node(b).bucket_path_from_root() assert 0 <= level_start < len(bucket_list) return self._storage.read_blocks(bucket_list[level_start:]) def write_path(self, b, buckets, level_start=0): assert 0 <= b < self._vheap.bucket_count() bucket_list = self._vheap.Node(b).bucket_path_from_root() assert 0 <= level_start < len(bucket_list) self._storage.write_blocks(bucket_list[level_start:], buckets) @property def bytes_sent(self): return self._storage.bytes_sent @property def bytes_received(self): return self._storage.bytes_received
def test_register_invalid_name(self): with self.assertRaises(ValueError): BlockStorageTypeFactory.register_device('s3', BlockStorageFile)
def test_ram(self): self.assertIs(BlockStorageTypeFactory('ram'), BlockStorageRAM)
def test_register_invalid_type(self): with self.assertRaises(TypeError): BlockStorageTypeFactory.register_device( 'new_str_type', str)
def test_sftp(self): self.assertIs(BlockStorageTypeFactory('sftp'), BlockStorageSFTP)
# class works with the encryption layer. # def read_blocks(self, indices): return [ bytes(block) for block in super(BlockStorageRAM, self).read_blocks(indices) ] def yield_blocks(self, indices): for block in super(BlockStorageRAM, self).yield_blocks(indices): yield bytes(block) def read_block(self, i): return bytes(super(BlockStorageRAM, self).read_block(i)) #def write_blocks(...) #def write_block(...) @property def bytes_sent(self): return self._bytes_sent @property def bytes_received(self): return self._bytes_received BlockStorageTypeFactory.register_device("ram", BlockStorageRAM)