def unpickle_scan_info(self, scan_id: str) -> None:
"""Unpickle a stored scan_inf corresponding to the scan_id
and store it in the scan_table"""
scan_info = self.scans_table.get(scan_id)
scan_info_hash = scan_info.pop('scan_info_hash')
pickler = DataPickler(self.file_storage_dir)
unpickled_scan_info = pickler.load_data(scan_id, scan_info_hash)
if not unpickled_scan_info:
pickler.remove_file(scan_id)
raise OspdCommandError(
'Not possible to unpickle stored scan info for %s' % scan_id,
'start_scan',
)
scan_info['results'] = list()
scan_info['temp_results'] = list()
scan_info['progress'] = ScanProgress.INIT.value
scan_info['target_progress'] = dict()
scan_info['count_alive'] = 0
scan_info['count_dead'] = 0
scan_info['count_total'] = None
scan_info['excluded_simplified'] = None
scan_info['target'] = unpickled_scan_info.pop('target')
scan_info['vts'] = unpickled_scan_info.pop('vts')
scan_info['options'] = unpickled_scan_info.pop('options')
scan_info['start_time'] = int(time.time())
scan_info['end_time'] = 0
self.scans_table[scan_id] = scan_info
pickler.remove_file(scan_id)
def test_load_data_corrupted(self):
data_pickler = DataPickler('/tmp')
data = {'foo', 'bar'}
filename = 'scan_info_1'
pickled_data = pickle.dumps(data)
tmp_hash = sha256()
tmp_hash.update(pickled_data)
pickled_data_hash = tmp_hash.hexdigest()
ret = data_pickler.store_data(filename, data)
self.assertEqual(ret, pickled_data_hash)
# courrupt data
file_to_corrupt = (
Path(data_pickler._storage_path) # pylint: disable=protected-access
/ filename)
with file_to_corrupt.open('ab') as f:
f.write(b'bar2')
original_data = data_pickler.load_data(filename, pickled_data_hash)
self.assertIsNone(original_data)
data_pickler.remove_file(filename)
def test_load_data_no_file(self, mock_logger):
filename = 'scan_info_1'
data_pickler = DataPickler('/tmp')
data_loaded = data_pickler.load_data(filename, "1234")
assert_called_once(mock_logger.error)
self.assertIsNone(data_loaded)
data_pickler.remove_file(filename)
def test_load_data(self):
data_pickler = DataPickler('/tmp')
data = {'foo', 'bar'}
filename = 'scan_info_1'
pickled_data = pickle.dumps(data)
tmp_hash = sha256()
tmp_hash.update(pickled_data)
pickled_data_hash = tmp_hash.hexdigest()
ret = data_pickler.store_data(filename, data)
self.assertEqual(ret, pickled_data_hash)
original_data = data_pickler.load_data(filename, pickled_data_hash)
self.assertIsNotNone(original_data)
self.assertIn('foo', original_data)