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)