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
