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_store_data(self):
data = {'foo', 'bar'}
filename = 'scan_info_1'
pickled_data = pickle.dumps(data)
tmp_hash = sha256()
tmp_hash.update(pickled_data)
data_pickler = DataPickler('/tmp')
ret = data_pickler.store_data(filename, data)
self.assertEqual(ret, tmp_hash.hexdigest())
data_pickler.remove_file(filename)
def create_scan(
self,
scan_id: str = '',
target: Dict = None,
options: Optional[Dict] = None,
vts: Dict = None,
) -> str:
"""Creates a new scan with provided scan information.
@target: Target to scan.
@options: Miscellaneous scan options supplied via <scanner_params>
XML element.
@return: Scan's ID. None if error occurs.
"""
if not options:
options = dict()
credentials = target.pop('credentials')
scan_info = self.data_manager.dict() # type: Dict
scan_info['status'] = ScanStatus.QUEUED
scan_info['credentials'] = credentials
scan_info['start_time'] = int(time.time())
scan_info['end_time'] = 0
scan_info_to_pickle = {
'target': target,
'options': options,
'vts': vts,
}
if scan_id is None or scan_id == '':
scan_id = str(uuid.uuid4())
pickler = DataPickler(self.file_storage_dir)
scan_info_hash = None
try:
scan_info_hash = pickler.store_data(scan_id, scan_info_to_pickle)
except OspdCommandError as e:
LOGGER.error(e)
return
scan_info['scan_id'] = scan_id
scan_info['scan_info_hash'] = scan_info_hash
self.scans_table[scan_id] = scan_info
return scan_id
def test_store_data_failed(self):
data = {'foo', 'bar'}
filename = 'scan_info_1'
data_pickler = DataPickler('/root')
self.assertRaises(OspdCommandError, data_pickler.store_data, filename,
data)
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(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)
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_store_data_check_permission(self):
OWNER_ONLY_RW_PERMISSION = '0o100600' # pylint: disable=invalid-name
data = {'foo', 'bar'}
filename = 'scan_info_1'
data_pickler = DataPickler('/tmp')
data_pickler.store_data(filename, data)
file_path = (
Path(data_pickler._storage_path) # pylint: disable=protected-access
/ filename)
self.assertEqual(oct(file_path.stat().st_mode),
OWNER_ONLY_RW_PERMISSION)
data_pickler.remove_file(filename)
def remove_file_pickled_scan_info(self, scan_id: str) -> None:
pickler = DataPickler(self.file_storage_dir)
pickler.remove_file(scan_id)
def test_remove_file_failed(self, mock_logger):
filename = 'inenxistent_file'
data_pickler = DataPickler('/root')
data_pickler.remove_file(filename)
assert_called_once(mock_logger.error)
