def get_encrypted_record(self,
enc_key: bytes = None,
nonce: bytes = None) -> dict:
"""
Return the encrypted form of the record.
:param enc_key: Key used for encryption.
:param nonce: DEBUG ONLY!
:return: Encrypted record as dict
"""
if enc_key is None and self._encryption_key is None:
raise ValueError("No encryption key defined.")
if enc_key is not None:
self._encryption_key = enc_key
length = len(self.record).to_bytes(
(len(self.record).bit_length() + 7) // 8, byteorder='big')
buf = struct.pack('%sd' % len(self.record), *self.record)
data = buf
longhash = self.get_long_hash()
key = self._encryption_key
# log.debug(f"Encryption - Using key: {self._encryption_key}")
if nonce is not None:
cipher = AES.new(key, AES.MODE_GCM, nonce=nonce)
else:
cipher = AES.new(key, AES.MODE_GCM)
cipher.update(length)
cipher.update(longhash)
ciphertext, mac = cipher.encrypt_and_digest(data)
json_k = ['nonce', 'length', 'hash', 'ciphertext', 'mac']
json_v = [
to_base64(x)
for x in (cipher.nonce, length, longhash, ciphertext, mac)
]
result = dict(zip(json_k, json_v))
return result
def get_upload_format(self) -> Tuple[str, str, str]:
"""
Return format required for upload to storage server
:return: [Base64(long_hash), json.dumps(ciphertext), owner]
"""
return (to_base64(self.get_long_hash()),
json.dumps(self.get_encrypted_record()), self.get_owner())
def test_client_integrity_bpe(self):
# Full Integrity Test including flask
self.c = client.Client(self.user)
self.c.set_password(self.password)
self.c.metric = "offset-0.1"
# Redirect requests to flask
target = self.target
# Server Records:
matches: List[Record] = self.matches
str_backend = StorageServer(test_dir)
with self.str_app.app_context():
for m in matches:
str_backend.store_record(
to_base64(m.get_long_hash()),
json.dumps(m.get_encrypted_record(self.enc_keys[0])),
'OwnerA')
# PSI Matches
psi_matches = []
for m in matches:
psi_matches.append(m.get_psi_index())
s = Session(True)
with patch("requests.get", s.get), \
patch("requests.post", s.post), \
patch.object(self.c, "_receive_ots",
Mock(return_value=self.enc_keys_int[:3])):
res = self.c.full_retrieve(target)
# Set hash key for comparison
for r in res:
r.set_hash_key(self.hash_key)
# Compare without order
for m in matches:
self.assertIn(m, res)
for r in res:
self.assertIn(r, matches)
def compute_matches_bloom(self,
candidate_iterator: SimilarityMetricIterator) \
-> List[Record]:
"""
Compute list of records stored on server side using a bloom filter.
:param candidate_iterator: Iterator over candidates
:return: List of Records found on server side.
"""
if self._psi_mode:
raise RuntimeError("Matches cannot be computed with bloom filter "
"because PSI-Mode is enabled.")
log.info(f"3.1 Retrieve bloom filter.")
b = self._get_bloom_filter()
self.eval['bloom_filter_retrieve_time'] = time.monotonic()
log.info(f"3.2 Compute matches with Bloom Filter.")
if config.PARALLEL:
# parallel
num_procs = multiprocessing.cpu_count()
its = candidate_iterator.split(num_procs)
num_procs = len(its)
log.debug(
f"Using {num_procs} parallel processes for bloom matching.")
processes = []
m = multiprocessing.Manager()
output = m.list()
def matching(client_set: RecordIterator,
i: int): # pragma no cover
"""Compute matching with given part of client set."""
log.debug(
f"Proc {i} iterating over {len(client_set)} elements.")
m = [
i for i in client_set if to_base64(i.get_long_hash()) in b
]
output.extend(m)
for i in range(num_procs):
client_set = RecordIterator(its[i], self._hash_key)
p = multiprocessing.Process(target=matching,
args=(client_set, i))
processes.append(p)
p.start()
atexit.register(p.terminate)
for i, p in enumerate(processes):
p.join()
atexit.unregister(p.terminate)
matches = list(output)
log.debug("All processes terminated.")
else:
client_set = RecordIterator(candidate_iterator,
self.get_hash_key())
matches = [
i for i in client_set if to_base64(i.get_long_hash()) in b
]
self.eval['bloom_matching_time'] = time.monotonic()
return matches
def test_data_provider_int(self):
# Full integrity Test including flask
self.dp = data_provider.DataProvider(self.provider)
self.dp.set_password(self.password)
str_backend = StorageServer(test_dir)
with self.str_app.app_context():
for r in self.sr:
# check that bloom filter is empty
b = str_backend.bloom
self.assertNotIn(to_base64(r.get_long_hash()), b)
# Check that DB empty
res = str_backend.batch_get_records(
[to_base64(r.get_long_hash()) for r in self.sr], "client")
# Decrypt
result = [
Record.from_ciphertext(json.loads(r), self.enc_keys[0])
for h, r in res
]
self.assertEqual([], result)
s = Session(True)
with patch("requests.get", s.get), \
patch("requests.post", s.post), \
patch.object(self.dp, "_receive_ots",
Mock(return_value=self.enc_keys_int[:len(self.sr)])):
self.dp.store_records(self.sr)
str_backend = StorageServer(test_dir)
with self.str_app.app_context():
for r in self.sr:
# check that records are in bloom filter
b = str_backend.bloom
self.assertIn(to_base64(r.get_long_hash()), b)
# Check records in db
res = str_backend.batch_get_records(
[to_base64(r.get_long_hash()) for r in self.sr], "client")
# Decrypt
result = [
Record.from_ciphertext(json.loads(r), self.enc_keys[0])
for h, r in res
]
for m in self.sr:
self.assertIn(m, result)
for r in result:
self.assertIn(r, self.sr)
def matching(client_set: RecordIterator,
i: int): # pragma no cover
"""Compute matching with given part of client set."""
log.debug(
f"Proc {i} iterating over {len(client_set)} elements.")
m = [
i for i in client_set if to_base64(i.get_long_hash()) in b
]
output.extend(m)
def test_get_hash_key(self):
url = f"{KEYSERVER}/mock/hash_key"
j = {
'success': True,
'hash_key': to_base64(int(1).to_bytes(16, 'big'))
}
responses.add(responses.GET, url, json=j, status=200)
# Success
key = self.m.get_hash_key()
self.assertEqual(int(1).to_bytes(16, 'big'), key)
def test_get_all_record_psi_hashes(self):
records = []
correct = []
for i in range(10):
r = Record([1, 2, 3, 4, 5])
r.set_hash_key(b'fake_key')
m = Mock()
m.hash = helpers.to_base64(r.get_long_hash())
records.append(m)
correct.append(r.get_psi_index())
with patch("lib.storage_server_backend.StoredRecord") as c:
c.query.all.return_value = records
s = server.StorageServer()
res = s.get_all_record_psi_hashes()
self.assertEqual(correct, res)
def test_store_record_success(self, m):
url = (f"https://{config.STORAGESERVER_HOSTNAME}:"
f"{config.STORAGE_API_PORT}/"
f"{UserType.OWNER}/store_record")
j = {'success': True, 'msg': None}
m.return_value.json.return_value = j
self.d._store_record_on_server(b"hash", {'cipher': "record"}, "userA")
m.assert_called_once()
expected = json.dumps({
'hash': to_base64(b"hash"),
'record': {
'cipher': "record"
},
'owner': 'userA'
}).encode()
m.called_with(url, json=expected)
def test_store_records(self):
"""Kind of integrity test, we only mock the server responses."""
# Define server response
# 1 - Get token
url = (f"https://{config.KEYSERVER_HOSTNAME}"
f":{config.KEY_API_PORT}/provider/gen_token")
j = {
'success':
True,
'token':
'XIu2a9SDGURRTzQnJdDg19Ii_CS7wy810s3_Lrx-TY7Wvh2Hf0U4xLH'
'NwnY_byYJ71II3kfUXpSZHOqAxA3zrw'
}
responses.add(responses.GET, url, json=j, status=200)
# 2 - Hash Key
url = f"{self.d.KEYSERVER}/hash_key"
hash_key = to_base64(int(1).to_bytes(16, 'big'))
j = {'success': True, 'hash_key': hash_key}
responses.add(responses.GET, url, json=j, status=200)
# 3 - Encryption Keys
j = {
'success': True,
'port': 50000,
'host': "127.0.0.1",
'totalOTs': 10,
'tls': config.OT_TLS
}
url = f"https://localhost:" \
f"{config.KEY_API_PORT}/provider/key_retrieval?totalOTs=3"
responses.add(responses.GET, url, json=j, status=200)
# Remember to mock the OT
r1 = Record([1.0, 2.1, 3.3, 4.4, 5.0])
r2 = Record([1.0532, 2.15423, 3.3453, 4.4, 5.0])
r3 = Record([1.52340, 2.1523, 3.35423, 4.4, 5.0])
records = [r1, r2, r3]
# Log in user
self.d.set_password("password")
with patch.object(self.d, "_receive_ots", return_value=[10, 9, 8]):
# Mock OT
with patch.object(self.d,
"_batch_store_records_on_server",
return_value=True):
self.d.store_records(records)
def get_hash_key(user_type: str, username: str) -> dict:
"""Return the hash key in Base64 encoding.
:param username: Username
:param user_type: Type of user
:return on success looks like this:
{
'success': True,
'hash_key': 'AAAAAAAAAAAAAAAAAAAAAQ=='
}
"""
app.logger.debug('Hash Key requested.')
_add_to_hash_key_db(user_type, username)
key = to_base64(get_keyserver_backend().get_hash_key())
return {'success': True,
'hash_key': key
}
def batch_get_records(self, candidates: List[Record]) -> List[Record]:
"""
Retrieve the records for all hashes in the list
:param candidates: Record objects for all candidates (hash_set)
:return: List of retrieved records (decyrpted)
"""
log.info("4.1 Retrieve encryption keys.")
start = time.monotonic()
ot_indices = []
# No duplicates
for r in candidates:
if r.get_ot_index() not in ot_indices:
ot_indices.append(r.get_ot_index())
enc_keys = self._get_enc_keys(ot_indices)
# Create mapping
enc_keys = dict(zip(ot_indices, enc_keys))
self.eval['key_retrieve_time'] = time.monotonic()
log.info(
f"4.1 - Retrieve keys took: {print_time(time.monotonic() - start)}"
)
log.info("4.2 Retrieve encrypted records.")
start = time.monotonic()
hash_list = [to_base64(r.get_long_hash()) for r in candidates]
records = self._batch_get_encrpyted_records(hash_list)
if not records:
self.eval['record_retrieve_time'] = time.monotonic()
self.eval['decryption_time'] = time.monotonic()
return []
res_list = []
self.eval['record_retrieve_time'] = time.monotonic()
log.info(
f"4.2 - Retrieve records: {print_time(time.monotonic() - start)}")
log.info("4.3 Decrypting.")
start = time.monotonic()
for h, c in records:
c = json.loads(c)
key = enc_keys[hash_to_index(from_base64(h), config.OT_INDEX_LEN)]
log.debug(f"Using key {key} for record {h}.")
res_list.append(Record.from_ciphertext(c, key))
self.eval['decryption_time'] = time.monotonic()
log.info(
f"4.3 - Decryption took: {print_time(time.monotonic() - start)}")
return res_list
def _store_record_on_server(self, hash_val: bytes, ciphertext: dict,
owner: str) -> None:
"""
Store the given record on the storage server.
:param hash_val: [Bytes] Long hash of record as returned by records
:param ciphertext: [Dict] encrypted record as returned by records
object
:param owner: [str] owner of record as string
:return:
"""
j = {
'hash': to_base64(hash_val),
'ciphertext': ciphertext,
'owner': owner
}
r = self.post(f"{self.STORAGESERVER}/store_record", json=j)
suc = r.json()['success']
if suc:
log.info("Successfully stored record.")
else:
msg = r.json()['msg']
raise RuntimeError(f"Failed to store record: {msg}")
def test__add_to_transaction_db(self, db, RecordRetrieval):
r1 = Record([1, 2, 3, 4, 5])
r2 = Record([1.1, 2, 3, 4, 5])
r3 = Record([1, 2.2, 3, 4, 5])
r4 = Record([1.1, 2.2, 3, 4, 5])
r5 = Record([1.2, 2.22, 3, 4, 5])
recs = [r1, r2, r3, r4, r5]
for r in recs:
r.set_hash_key(b'hash-key')
hashes = [helpers.to_base64(r.get_long_hash()) for r in recs]
records = [Mock() for _ in range(5)]
records[0].hash = helpers.to_base64(r1.get_long_hash())
records[1].hash = helpers.to_base64(r1.get_long_hash()) # Same
records[2].hash = helpers.to_base64(r3.get_long_hash())
records[3].hash = helpers.to_base64(r4.get_long_hash())
records[4].hash = helpers.to_base64(r5.get_long_hash())
server.StorageServer._add_to_transaction_db(records, "client", hashes)
self.assertEqual(1, RecordRetrieval.call_count) # 2 owners
expected = {
"client": "client",
"enc_keys_by_hash": 5,
"enc_keys_by_records": 4
}
self.assertEqual(expected, RecordRetrieval.call_args[1])
def test_base64(self):
b = b'Test'
self.assertEqual(b,
helpers.from_base64(helpers.to_base64(b)))