def start_processing_per_tenant(vulnerability_index: str, asset_index: str): LOGGER.info( F'Calculation for {vulnerability_index} and {asset_index} started') try: assets_count = AssetDocument.search(index=asset_index).filter( ~Q('match', tags=AssetStatus.DELETED)).count() vuln_search = VulnerabilityDocument.search( index=vulnerability_index).filter( ~Q('match', tags=VulnerabilityStatus.FIXED) & ~Q('match', asset__tags=AssetStatus.DELETED)) prepare(vulnerability_index) workers_count = get_workers_count() vuln_count = vuln_search.count() slices_count = 1 if vuln_count > 500: slices_count = vuln_count // workers_count slices_count = slices_count if slices_count <= workers_count else workers_count (group( _processing.si(idx, slices_count, assets_count, vulnerability_index) for idx in range(slices_count)) | _end_processing.si(vulnerability_index, asset_index))() except Exception as ex: LOGGER.error(F'Unknown processing exception {ex}')
def prepare(vulnerability_index): s = VulnerabilityDocument.search(index=vulnerability_index).filter( ~Q('match', tags=VulnerabilityStatus.FIXED) & ~Q('match', asset__tags=AssetStatus.DELETED)) s.aggs.bucket('cves', 'terms', field='cve.id', size=10000000) s = s.execute() for result in s.aggregations.cves.buckets: key = '{}-{}'.format(vulnerability_index, result['key']) cache.set(key, result['doc_count'], LOCK_EXPIRE)
def get_or_create(ip_address, config=None): index = AssetDocument.get_index(config) result = AssetDocument.search(index=index).filter( Q('term', ip_address=ip_address) & ~Q('match', tags=AssetStatus.DELETED)).execute() if result: return result[0] return AssetDocument(id=ip_address, ip_address=ip_address, tags=[AssetStatus.DISCOVERED]).save(index=index, refresh=True)
def _processing(idx, slices_count, assets_count, vulnerability_index): docs = [] try: vuln_search = VulnerabilityDocument.search( index=vulnerability_index).filter( ~Q('match', tags=VulnerabilityStatus.FIXED) & ~Q('match', asset__tags=AssetStatus.DELETED)) LOGGER.debug( F'Calculation for {vulnerability_index} and {idx}, {slices_count} started' ) if slices_count > 1: vuln_search = vuln_search.extra(slice={ "id": idx, "max": slices_count }).params(scroll="60m") # List competence used due to better performance vulns = [vuln for vuln in vuln_search.scan()] LOGGER.debug(F'all vulns for slice {idx} downloaded') for vuln in vulns: score, vector = calculate_environmental_score_v3(vuln) vuln.environmental_score_vector_v3 = vector vuln.environmental_score_v3 = score vuln_count = get_cve_count(vulnerability_index, vuln.cve.id) score, vector = calculate_environmental_score_v2( vuln, vuln_count, assets_count) vuln.environmental_score_vector_v2 = vector vuln.environmental_score_v2 = score docs.append(vuln.to_dict(include_meta=True)) if len(docs) > 10000: async_bulk(docs, vulnerability_index) docs = [] async_bulk(docs, vulnerability_index) except Exception as ex: LOGGER.error(F'Unknown processing exception {ex}') finally: thread_pool_executor.wait_for_all() LOGGER.debug( F'Calculation for {vulnerability_index} and {idx}, {slices_count} done' )
def _update_existing_assets(assets: dict, index): if assets: updated = [] assets_search = AssetDocument.search( index=index).filter(~Q('match', tags=AssetStatus.DISCOVERED)) current_assets = [a for a in assets_search.scan()] for current_asset in current_assets: asset_id = current_asset.id if asset_id in assets: if current_asset.has_changed(assets[asset_id]): current_asset.update(assets[asset_id], index=index, weak=True) updated.append( current_asset.to_dict(include_meta=True)) del assets[asset_id] elif asset_id not in assets and AssetStatus.DELETED not in current_asset.tags: current_asset.tags.append(AssetStatus.DELETED) updated.append( current_asset.save( index=index, weak=True).to_dict(include_meta=True)) if len(updated) > 500: async_bulk(updated) updated = [] async_bulk(updated) return assets
def get_or_create(ip_address, config=None): index = AssetDocument.get_index(config) result = AssetDocument.search(index=index).filter( Q('term', ip_address=ip_address) & ~Q('match', tags=AssetStatus.DELETED)).execute() if result: return result[0] if hasattr(config, 'scanner'): source = config.scanner.split('.')[-1].capitalize() else: source = 'vmc' return AssetDocument(id=ip_address, ip_address=ip_address, tenant=config.tenant.name if config and config.tenant else None, source=source, tags=[AssetStatus.DISCOVERED]).save(index=index, refresh=True)
def _update_discovered_assets(assets: dict, index): if assets: updated = [] assets = {a.ip_address: a for a in assets.values()} assets_search = AssetDocument.search(index=index).filter(Q('match', tags=AssetStatus.DISCOVERED)) discovered_assets = [a for a in assets_search.scan()] for discovered_asset in discovered_assets: if discovered_asset.ip_address in assets: discovered_asset.update(assets[discovered_asset.ip_address], index=index, weak=True) updated.append(discovered_asset.to_dict(include_meta=True)) del assets[discovered_asset.ip_address] if len(updated) > 500: async_bulk(updated) updated = [] async_bulk(updated) return assets
def test_delete_asset(self): asset_1 = self.create_asset(asset_id=1, ip_address='10.0.0.1', hostname='hostname_1') asset_2 = self.create_asset(asset_id=2, ip_address='10.0.0.2', hostname='hostname_2') self.assertEqual(2, Search().index(AssetDocument.Index.name).count()) AssetDocument.create_or_update({asset_1.id: asset_1}, AssetConfigMock()) thread_pool_executor.wait_for_all() result = AssetDocument.search().filter( Q('match', tags=AssetStatus.DELETED)).execute() self.assertEqual(1, len(result.hits)) self.assertEqual(result.hits[0].ip_address, asset_2.ip_address) self.assertEqual(result.hits[0].id, asset_2.id)
def get_assets_with_tag(tag: str, config=None): index = AssetDocument.get_index(config) result = AssetDocument.search(index=index).filter(Q("match", tags=tag)) return result