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)