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)))