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_update_header_data(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 = 1
header_data = bytes(bytearray([0,1,2]))
fsetup = EncryptedHeapStorage.setup(
fname,
bsize,
self._heap_height,
heap_base=self._heap_base,
blocks_per_bucket=blocks_per_bucket,
header_data=header_data)
fsetup.close()
new_header_data = bytes(bytearray([1,1,1]))
with TopCachedEncryptedHeapStorage(
EncryptedHeapStorage(
fname,
key=fsetup.key,
storage_type=self._storage_type),
**self._init_kwds) as f:
self.assertEqual(f.header_data, header_data)
f.update_header_data(new_header_data)
self.assertEqual(f.header_data, new_header_data)
with TopCachedEncryptedHeapStorage(
EncryptedHeapStorage(
fname,
key=fsetup.key,
storage_type=self._storage_type),
**self._init_kwds) as f:
self.assertEqual(f.header_data, new_header_data)
with self.assertRaises(ValueError):
with TopCachedEncryptedHeapStorage(
EncryptedHeapStorage(
fname,
key=fsetup.key,
storage_type=self._storage_type),
**self._init_kwds) as f:
f.update_header_data(bytes(bytearray([1,1])))
with self.assertRaises(ValueError):
with TopCachedEncryptedHeapStorage(
EncryptedHeapStorage(
fname,
key=fsetup.key,
storage_type=self._storage_type),
**self._init_kwds) as f:
f.update_header_data(bytes(bytearray([1,1,1,1])))
with TopCachedEncryptedHeapStorage(
EncryptedHeapStorage(
fname,
key=fsetup.key,
storage_type=self._storage_type),
**self._init_kwds) as f:
self.assertEqual(f.header_data, new_header_data)
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 main():
#
# get a unique filename in the current directory
#
fid, tmpname = tempfile.mkstemp(dir=os.getcwd())
os.close(fid)
os.remove(tmpname)
print("Storage Name: %s" % (tmpname))
key_size = 32
header_data = b'a message'
heap_base = 3
heap_height = 2
block_size = 8
blocks_per_bucket=4
initialize = lambda i: \
bytes(bytearray([i] * block_size * blocks_per_bucket))
vheap = SizedVirtualHeap(
heap_base,
heap_height,
blocks_per_bucket=blocks_per_bucket)
with EncryptedHeapStorage.setup(
tmpname,
block_size,
heap_height,
key_size=key_size,
header_data=header_data,
heap_base=heap_base,
blocks_per_bucket=blocks_per_bucket,
initialize=initialize) as f:
assert tmpname == f.storage_name
assert f.header_data == header_data
print(f.read_path(vheap.random_bucket()))
key = f.key
assert os.path.exists(tmpname)
with EncryptedHeapStorage(tmpname, key=key) as f:
assert tmpname == f.storage_name
assert f.header_data == header_data
print(f.read_path(vheap.random_bucket()))
#
# cleanup
#
os.remove(tmpname)
def main():
#
# get a unique filename in the current directory
#
fid, tmpname = tempfile.mkstemp(dir=os.getcwd())
os.close(fid)
os.remove(tmpname)
print("Storage Name: %s" % (tmpname))
key_size = 32
header_data = b'a message'
heap_base = 3
heap_height = 2
block_size = 8
blocks_per_bucket = 4
initialize = lambda i: \
bytes(bytearray([i] * block_size * blocks_per_bucket))
vheap = SizedVirtualHeap(heap_base,
heap_height,
blocks_per_bucket=blocks_per_bucket)
with EncryptedHeapStorage.setup(tmpname,
block_size,
heap_height,
key_size=key_size,
header_data=header_data,
heap_base=heap_base,
blocks_per_bucket=blocks_per_bucket,
initialize=initialize) as f:
assert tmpname == f.storage_name
assert f.header_data == header_data
print(f.read_path(vheap.random_bucket()))
key = f.key
assert os.path.exists(tmpname)
with EncryptedHeapStorage(tmpname, key=key) as f:
assert tmpname == f.storage_name
assert f.header_data == header_data
print(f.read_path(vheap.random_bucket()))
#
# cleanup
#
os.remove(tmpname)
def setUpClass(cls):
assert cls._init_kwds is not None
assert cls._storage_type is not None
assert cls._heap_base is not None
assert cls._heap_height is not None
fd, cls._dummy_name = tempfile.mkstemp()
os.close(fd)
try:
os.remove(cls._dummy_name)
except OSError: # pragma: no cover
pass # pragma: no cover
cls._block_size = 50
cls._blocks_per_bucket = 3
cls._bucket_count = \
((cls._heap_base**(cls._heap_height+1)) - 1)//(cls._heap_base-1)
cls._block_count = cls._bucket_count * \
cls._blocks_per_bucket
cls._testfname = cls.__name__ + "_testfile.bin"
cls._buckets = []
f = EncryptedHeapStorage.setup(
cls._testfname,
cls._block_size,
cls._heap_height,
heap_base=cls._heap_base,
blocks_per_bucket=cls._blocks_per_bucket,
storage_type=cls._storage_type,
initialize=lambda i: bytes(bytearray([i]) * \
cls._block_size * \
cls._blocks_per_bucket),
ignore_existing=True)
f.close()
cls._key = f.key
for i in range(cls._bucket_count):
data = bytearray([i]) * \
cls._block_size * \
cls._blocks_per_bucket
cls._buckets.append(data)
def setup(cls,
storage_name,
block_size,
block_count,
bucket_capacity=4,
heap_base=2,
cached_levels=3,
concurrency_level=None,
**kwds):
if 'heap_height' in kwds:
raise ValueError("'heap_height' keyword is not accepted")
if (bucket_capacity <= 0) or \
(bucket_capacity != int(bucket_capacity)):
raise ValueError(
"Bucket capacity must be a positive integer: %s"
% (bucket_capacity))
if (block_size <= 0) or (block_size != int(block_size)):
raise ValueError(
"Block size (bytes) must be a positive integer: %s"
% (block_size))
if (block_count <= 0) or (block_count != int(block_count)):
raise ValueError(
"Block count must be a positive integer: %s"
% (block_count))
if heap_base < 2:
raise ValueError(
"heap base must be 2 or greater. Invalid value: %s"
% (heap_base))
heap_height = calculate_necessary_heap_height(heap_base,
block_count)
stash = {}
vheap = SizedVirtualHeap(
heap_base,
heap_height,
blocks_per_bucket=bucket_capacity)
position_map = cls._init_position_map(vheap, block_count)
oram_block_size = block_size + \
TreeORAMStorageManagerExplicitAddressing.\
block_info_storage_size
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)))
initialize = kwds.pop('initialize', None)
header_data = struct.pack(
cls._header_struct_string,
block_count)
kwds['header_data'] = bytes(header_data) + user_header_data
empty_bucket = bytearray(oram_block_size * bucket_capacity)
empty_bucket_view = memoryview(empty_bucket)
for i in xrange(bucket_capacity):
TreeORAMStorageManagerExplicitAddressing.tag_block_as_empty(
empty_bucket_view[(i*oram_block_size):\
((i+1)*oram_block_size)])
empty_bucket = bytes(empty_bucket)
kwds['initialize'] = lambda i: empty_bucket
f = None
try:
log.info("%s: setting up encrypted heap storage"
% (cls.__name__))
f = EncryptedHeapStorage.setup(storage_name,
oram_block_size,
heap_height,
heap_base=heap_base,
blocks_per_bucket=bucket_capacity,
**kwds)
if cached_levels != 0:
f = TopCachedEncryptedHeapStorage(
f,
cached_levels=cached_levels,
concurrency_level=concurrency_level)
elif concurrency_level is not None:
raise ValueError( # pragma: no cover
"'concurrency_level' keyword is " # pragma: no cover
"not used when no heap levels " # pragma: no cover
"are cached") # pragma: no cover
oram = TreeORAMStorageManagerExplicitAddressing(
f, stash, position_map)
if initialize is None:
zeros = bytes(bytearray(block_size))
initialize = lambda i: zeros
initial_oram_block = bytearray(oram_block_size)
for i in tqdm.tqdm(xrange(block_count),
desc=("Initializing %s Blocks" % (cls.__name__)),
total=block_count,
disable=not pyoram.config.SHOW_PROGRESS_BAR):
oram.tag_block_with_id(initial_oram_block, i)
initial_oram_block[oram.block_info_storage_size:] = \
initialize(i)[:]
bucket = oram.position_map[i]
bucket_level = vheap.Node(bucket).level
oram.position_map[i] = \
oram.storage_heap.virtual_heap.\
random_bucket_at_level(bucket_level)
oram.load_path(bucket)
oram.push_down_path()
# place a copy in the stash
oram.stash[i] = bytearray(initial_oram_block)
oram.fill_path_from_stash()
oram.evict_path()
header_data = bytearray(header_data)
stash_digest = cls.stash_digest(
oram.stash,
digestmod=hmac.HMAC(key=oram.storage_heap.key,
digestmod=hashlib.sha384))
position_map_digest = cls.position_map_digest(
oram.position_map,
digestmod=hmac.HMAC(key=oram.storage_heap.key,
digestmod=hashlib.sha384))
header_data[:len(stash_digest)] = stash_digest[:]
header_data[len(stash_digest):\
(len(stash_digest)+len(position_map_digest))] = \
position_map_digest[:]
f.update_header_data(bytes(header_data) + user_header_data)
return PathORAM(f, stash, position_map=position_map)
except:
if f is not None:
f.close() # pragma: no cover
raise
def test_setup_fails(self):
self.assertEqual(os.path.exists(self._dummy_name), False)
with self.assertRaises(IOError):
EncryptedHeapStorage.setup(os.path.join(thisdir, "baselines",
"exists.empty"),
block_size=10,
heap_height=1,
key_size=AES.key_sizes[0],
blocks_per_bucket=1,
storage_type=self._type_name)
self.assertEqual(os.path.exists(self._dummy_name), False)
with self.assertRaises(IOError):
EncryptedHeapStorage.setup(os.path.join(thisdir, "baselines",
"exists.empty"),
block_size=10,
heap_height=1,
key_size=AES.key_sizes[0],
blocks_per_bucket=1,
storage_type=self._type_name,
ignore_existing=False)
self.assertEqual(os.path.exists(self._dummy_name), False)
# bad block_size
with self.assertRaises(ValueError):
EncryptedHeapStorage.setup(self._dummy_name,
block_size=0,
heap_height=1,
key_size=AES.key_sizes[0],
blocks_per_bucket=1,
storage_type=self._type_name)
self.assertEqual(os.path.exists(self._dummy_name), False)
# bad heap_height
with self.assertRaises(ValueError):
EncryptedHeapStorage.setup(self._dummy_name,
block_size=1,
heap_height=-1,
blocks_per_bucket=1,
storage_type=self._type_name)
self.assertEqual(os.path.exists(self._dummy_name), False)
# bad blocks_per_bucket
with self.assertRaises(ValueError):
EncryptedHeapStorage.setup(self._dummy_name,
block_size=1,
heap_height=1,
key_size=AES.key_sizes[0],
blocks_per_bucket=0,
storage_type=self._type_name)
self.assertEqual(os.path.exists(self._dummy_name), False)
# bad heap_base
with self.assertRaises(ValueError):
EncryptedHeapStorage.setup(self._dummy_name,
block_size=1,
heap_height=1,
key_size=AES.key_sizes[0],
blocks_per_bucket=1,
heap_base=1,
storage_type=self._type_name)
self.assertEqual(os.path.exists(self._dummy_name), False)
# bad header_data
with self.assertRaises(TypeError):
EncryptedHeapStorage.setup(self._dummy_name,
block_size=1,
heap_height=1,
key_size=AES.key_sizes[0],
blocks_per_bucket=1,
storage_type=self._type_name,
header_data=2)
self.assertEqual(os.path.exists(self._dummy_name), False)
# uses block_count
with self.assertRaises(ValueError):
EncryptedHeapStorage.setup(self._dummy_name,
block_size=1,
heap_height=1,
key_size=AES.key_sizes[0],
blocks_per_bucket=1,
block_count=1,
storage_type=self._type_name)
self.assertEqual(os.path.exists(self._dummy_name), False)
def setup(cls,
storage_name,
block_size,
block_count,
bucket_capacity=4,
heap_base=2,
cached_levels=3,
concurrency_level=None,
**kwds):
if 'heap_height' in kwds:
raise ValueError("'heap_height' keyword is not accepted")
if (bucket_capacity <= 0) or \
(bucket_capacity != int(bucket_capacity)):
raise ValueError("Bucket capacity must be a positive integer: %s" %
(bucket_capacity))
if (block_size <= 0) or (block_size != int(block_size)):
raise ValueError(
"Block size (bytes) must be a positive integer: %s" %
(block_size))
if (block_count <= 0) or (block_count != int(block_count)):
raise ValueError("Block count must be a positive integer: %s" %
(block_count))
if heap_base < 2:
raise ValueError(
"heap base must be 2 or greater. Invalid value: %s" %
(heap_base))
heap_height = calculate_necessary_heap_height(heap_base, block_count)
stash = {}
vheap = SizedVirtualHeap(heap_base,
heap_height,
blocks_per_bucket=bucket_capacity)
position_map = cls._init_position_map(vheap, block_count)
oram_block_size = block_size + \
TreeORAMStorageManagerExplicitAddressing.\
block_info_storage_size
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)))
initialize = kwds.pop('initialize', None)
header_data = struct.pack(cls._header_struct_string, block_count)
kwds['header_data'] = bytes(header_data) + user_header_data
empty_bucket = bytearray(oram_block_size * bucket_capacity)
empty_bucket_view = memoryview(empty_bucket)
for i in xrange(bucket_capacity):
TreeORAMStorageManagerExplicitAddressing.tag_block_as_empty(
empty_bucket_view[(i*oram_block_size):\
((i+1)*oram_block_size)])
empty_bucket = bytes(empty_bucket)
kwds['initialize'] = lambda i: empty_bucket
f = None
try:
log.info("%s: setting up encrypted heap storage" % (cls.__name__))
f = EncryptedHeapStorage.setup(storage_name,
oram_block_size,
heap_height,
heap_base=heap_base,
blocks_per_bucket=bucket_capacity,
**kwds)
if cached_levels != 0:
f = TopCachedEncryptedHeapStorage(
f,
cached_levels=cached_levels,
concurrency_level=concurrency_level)
elif concurrency_level is not None:
raise ValueError( # pragma: no cover
"'concurrency_level' keyword is " # pragma: no cover
"not used when no heap levels " # pragma: no cover
"are cached") # pragma: no cover
oram = TreeORAMStorageManagerExplicitAddressing(
f, stash, position_map)
if initialize is None:
zeros = bytes(bytearray(block_size))
initialize = lambda i: zeros
initial_oram_block = bytearray(oram_block_size)
for i in tqdm.tqdm(xrange(block_count),
desc=("Initializing %s Blocks" %
(cls.__name__)),
total=block_count,
disable=not pyoram.config.SHOW_PROGRESS_BAR):
oram.tag_block_with_id(initial_oram_block, i)
initial_oram_block[oram.block_info_storage_size:] = \
initialize(i)[:]
bucket = oram.position_map[i]
bucket_level = vheap.Node(bucket).level
oram.position_map[i] = \
oram.storage_heap.virtual_heap.\
random_bucket_at_level(bucket_level)
oram.load_path(bucket)
oram.push_down_path()
# place a copy in the stash
oram.stash[i] = bytearray(initial_oram_block)
oram.fill_path_from_stash()
oram.evict_path()
header_data = bytearray(header_data)
stash_digest = cls.stash_digest(oram.stash,
digestmod=hmac.HMAC(
key=oram.storage_heap.key,
digestmod=hashlib.sha384))
position_map_digest = cls.position_map_digest(
oram.position_map,
digestmod=hmac.HMAC(key=oram.storage_heap.key,
digestmod=hashlib.sha384))
header_data[:len(stash_digest)] = stash_digest[:]
header_data[len(stash_digest):\
(len(stash_digest)+len(position_map_digest))] = \
position_map_digest[:]
f.update_header_data(bytes(header_data) + user_header_data)
return PathORAM(f, stash, position_map=position_map)
except:
if f is not None:
f.close() # pragma: no cover
raise
def main():
storage_name = "heap.bin"
print("Storage Name: %s" % (storage_name))
key_size = 32
heap_base = 2
heap_height = 2
block_size = struct.calcsize("!?LL")
blocks_per_bucket = 2
vheap = SizedVirtualHeap(
heap_base,
heap_height,
blocks_per_bucket=blocks_per_bucket)
print("Block Size: %s" % (block_size))
print("Blocks Per Bucket: %s" % (blocks_per_bucket))
position_map = {}
def initialize(i):
bucket = bytes()
for j in range(blocks_per_bucket):
if (i*j) % 3:
bucket += struct.pack(
"!?LL", False, 0, 0)
else:
x = vheap.Node(i)
while not vheap.is_nil_node(x):
x = x.child_node(random.randint(0, heap_base-1))
x = x.parent_node()
bucket += struct.pack(
"!?LL", True, initialize.id_, x.bucket)
position_map[initialize.id_] = x.bucket
initialize.id_ += 1
return bucket
initialize.id_ = 1
with EncryptedHeapStorage.setup(
storage_name,
block_size,
heap_height,
heap_base=heap_base,
key_size=key_size,
blocks_per_bucket=blocks_per_bucket,
initialize=initialize,
ignore_existing=True) as f:
assert storage_name == f.storage_name
stash = {}
oram = TreeORAMStorageManagerPointerAddressing(f, stash)
b = vheap.random_bucket()
oram.load_path(b)
print("")
print(repr(vheap.Node(oram.path_stop_bucket)))
print(oram.path_block_ids)
print(oram.path_block_eviction_levels)
oram.push_down_path()
print("")
print(repr(vheap.Node(oram.path_stop_bucket)))
print(oram.path_block_ids)
print(oram.path_block_eviction_levels)
print(oram.path_block_reordering)
oram.evict_path()
oram.load_path(b)
print("")
print(repr(vheap.Node(oram.path_stop_bucket)))
print(oram.path_block_ids)
print(oram.path_block_eviction_levels)
oram.push_down_path()
print("")
print(repr(vheap.Node(oram.path_stop_bucket)))
print(oram.path_block_ids)
print(oram.path_block_eviction_levels)
print(oram.path_block_reordering)
assert all(x is None for x in oram.path_block_reordering)
os.remove(storage_name)
def test_setup_fails(self):
self.assertEqual(os.path.exists(self._dummy_name), False)
with self.assertRaises(IOError):
EncryptedHeapStorage.setup(
os.path.join(thisdir,
"baselines",
"exists.empty"),
block_size=10,
heap_height=1,
key_size=AES.key_sizes[0],
blocks_per_bucket=1,
storage_type=self._type_name)
self.assertEqual(os.path.exists(self._dummy_name), False)
with self.assertRaises(IOError):
EncryptedHeapStorage.setup(
os.path.join(thisdir,
"baselines",
"exists.empty"),
block_size=10,
heap_height=1,
key_size=AES.key_sizes[0],
blocks_per_bucket=1,
storage_type=self._type_name,
ignore_existing=False)
self.assertEqual(os.path.exists(self._dummy_name), False)
# bad block_size
with self.assertRaises(ValueError):
EncryptedHeapStorage.setup(
self._dummy_name,
block_size=0,
heap_height=1,
key_size=AES.key_sizes[0],
blocks_per_bucket=1,
storage_type=self._type_name)
self.assertEqual(os.path.exists(self._dummy_name), False)
# bad heap_height
with self.assertRaises(ValueError):
EncryptedHeapStorage.setup(
self._dummy_name,
block_size=1,
heap_height=-1,
blocks_per_bucket=1,
storage_type=self._type_name)
self.assertEqual(os.path.exists(self._dummy_name), False)
# bad blocks_per_bucket
with self.assertRaises(ValueError):
EncryptedHeapStorage.setup(
self._dummy_name,
block_size=1,
heap_height=1,
key_size=AES.key_sizes[0],
blocks_per_bucket=0,
storage_type=self._type_name)
self.assertEqual(os.path.exists(self._dummy_name), False)
# bad heap_base
with self.assertRaises(ValueError):
EncryptedHeapStorage.setup(
self._dummy_name,
block_size=1,
heap_height=1,
key_size=AES.key_sizes[0],
blocks_per_bucket=1,
heap_base=1,
storage_type=self._type_name)
self.assertEqual(os.path.exists(self._dummy_name), False)
# bad header_data
with self.assertRaises(TypeError):
EncryptedHeapStorage.setup(
self._dummy_name,
block_size=1,
heap_height=1,
key_size=AES.key_sizes[0],
blocks_per_bucket=1,
storage_type=self._type_name,
header_data=2)
self.assertEqual(os.path.exists(self._dummy_name), False)
# uses block_count
with self.assertRaises(ValueError):
EncryptedHeapStorage.setup(
self._dummy_name,
block_size=1,
heap_height=1,
key_size=AES.key_sizes[0],
blocks_per_bucket=1,
block_count=1,
storage_type=self._type_name)
self.assertEqual(os.path.exists(self._dummy_name), False)
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)