Skip to content

Commit bfda9af

Browse files
sbs2001sbs2001
committed
Remove magic methods and properties in ImportRunner
Signed-off-by: Shivam Sandbhor <shivam.sandbhor@gmail.com>
1 parent c421fe1 commit bfda9af

1 file changed

Lines changed: 81 additions & 115 deletions

File tree

‎vulnerabilities/import_runner.py‎

Lines changed: 81 additions & 115 deletions
Original file line numberDiff line numberDiff line change
@@ -61,46 +61,6 @@ class ImportRunner:
6161
def __init__(self, importer: models.Importer, batch_size: int):
6262
self.importer = importer
6363
self.batch_size = batch_size
64-
self._bulk_create_vuln_pkg_refs = []
65-
self._bulk_update_vuln_pkg_refs = []
66-
self._bulk_create_vuln_refs = []
67-
68-
def __del__(self):
69-
models.Vulnerability_Package_Relation.objects.bulk_create(self._bulk_create_vuln_pkg_refs, ignore_conflicts=True)
70-
models.Vulnerability_Package_Relation.objects.bulk_update(self._bulk_update_vuln_pkg_refs,['is_vulnerable'])
71-
models.VulnerabilityReference.objects.bulk_create(self._bulk_create_vuln_refs, ignore_conflicts=True)
72-
73-
# TODO: Currently the implementation may contain duplicate model instances in `_bulk_create_vuln_refs` and `_bulk_create_vuln_pkg_refs`
74-
# to make this work , `ignore_conflicts=True` is used. Find some way to remove duplicates of model instances(before writing them db)
75-
# and remove the `ignore_conflicts=True` flag.
76-
@property
77-
def bulk_create_vuln_pkg_refs(self):
78-
79-
if len(self._bulk_create_vuln_pkg_refs) >= self.batch_size :
80-
models.Vulnerability_Package_Relation.objects.bulk_create(self._bulk_create_vuln_pkg_refs, ignore_conflicts=True)
81-
self._bulk_create_vuln_pkg_refs = []
82-
83-
return self._bulk_create_vuln_pkg_refs
84-
85-
86-
@property
87-
def bulk_update_vuln_pkg_refs(self):
88-
89-
if len(self._bulk_update_vuln_pkg_refs) >= self.batch_size :
90-
models.Vulnerability_Package_Relation.objects.bulk_update(self._bulk_update_vuln_pkg_refs,['is_vulnerable'])
91-
self._bulk_update_vuln_pkg_refs = []
92-
93-
return self._bulk_update_vuln_pkg_refs
94-
95-
@property
96-
def bulk_create_vuln_refs(self):
97-
98-
if len(self._bulk_create_vuln_refs) >= self.batch_size :
99-
models.VulnerabilityReference.objects.bulk_create(self._bulk_create_vuln_refs, ignore_conflicts=True)
100-
self._bulk_create_vuln_refs = []
101-
102-
return self._bulk_create_vuln_refs
103-
10464

10565
def run(self, cutoff_date: datetime.datetime = None) -> None:
10666
"""
@@ -117,91 +77,97 @@ def run(self, cutoff_date: datetime.datetime = None) -> None:
11777
data_source = self.importer.make_data_source(self.batch_size, cutoff_date=cutoff_date)
11878
with data_source:
11979
_process_added_advisories(data_source)
120-
self._process_updated_advisories(data_source)
80+
_process_updated_advisories(data_source)
12181
self.importer.last_run = datetime.datetime.now(tz=datetime.timezone.utc)
12282
self.importer.data_source_cfg = dataclasses.asdict(data_source.config)
12383
self.importer.save()
12484

12585
logger.debug(f'Successfully finished import for {self.importer.name}.')
12686

12787

128-
def _process_updated_advisories(self,data_source: DataSource) -> None:
129-
"""
130-
TODO: Break this method into smaller functions
131-
"""
132-
133-
for batch in data_source.updated_advisories():
134-
for advisory in batch:
135-
vuln, vuln_created = _get_or_create_vulnerability(advisory)
136-
137-
if vuln_created:
138-
# This means vulnerability didn't previously exist in the db, so bulk create
139-
# is used without any hesitation
140-
for id_ in set(advisory.reference_ids):
141-
self.bulk_create_vuln_refs.append(models.VulnerabilityReference(vulnerability=vuln, reference_id=id_))
142-
143-
for url in set(advisory.reference_urls):
144-
self.bulk_create_vuln_refs.append(models.VulnerabilityReference(vulnerability=vuln, url=url))
145-
146-
else:
147-
vuln_refs_qs = models.VulnerabilityReference.objects.filter(vulnerability=vuln)
148-
# This is to avoid making additional SELECT queries, and do further filtering in python
149-
# Check https://stackoverflow.com/a/5989530
150-
vuln_ids = {ref.reference_id for ref in vuln_refs_qs}
151-
vuln_urls = {ref.url for ref in vuln_refs_qs}
152-
153-
for id_ in advisory.reference_ids:
154-
if id_ not in vuln_ids :
155-
# Delete the item because it will allow the duplicates pass to through if they are not
156-
# present in vuln_ids
157-
vuln_ids.add(id_)
158-
self.bulk_create_vuln_refs.append(models.VulnerabilityReference(vulnerability=vuln, reference_id=id_))
159-
160-
for url in advisory.reference_urls:
161-
# Delete the item because it will allow the duplicates pass to through if they are not
162-
# present in vuln_urls
163-
if url not in vuln_urls:
164-
vuln_urls.add(url)
165-
self.bulk_create_vuln_refs.append(models.VulnerabilityReference(vulnerability=vuln, url=url))
88+
def _process_updated_advisories(data_source: DataSource) -> None:
89+
"""
90+
TODO: Break this method into smaller functions
91+
"""
92+
for batch in data_source.updated_advisories():
93+
for advisory in batch:
16694

95+
bulk_create_vuln_refs = []
96+
bulk_update_vuln_pkg_refs = []
97+
bulk_create_vuln_pkg_refs = []
16798

168-
99+
vuln, vuln_created = _get_or_create_vulnerability(advisory)
169100

170-
for ipurl in advisory.impacted_package_urls :
171-
pkg, pkg_created = _get_or_create_package(ipurl)
172-
vuln_pkgs_ref = models.Vulnerability_Package_Relation(package=pkg,vulnerability=vuln,is_vulnerable=True)
173-
174-
if pkg_created or vuln_created :
175-
self.bulk_create_vuln_pkg_refs.append(vuln_pkgs_ref)
176-
177-
else:
178-
qs = models.Vulnerability_Package_Relation.objects.filter(package=pkg,vulnerability=vuln)
179-
if qs:
180-
if not qs[0].is_vulnerable:
181-
qs[0].is_vulnerable = True
182-
self.bulk_update_vuln_pkg_refs.append(qs[0])
183-
184-
else:
185-
self.bulk_create_vuln_pkg_refs.append(vuln_pkgs_ref)
186-
187-
for rpurl in advisory.resolved_package_urls :
188-
pkg, pkg_created = _get_or_create_package(rpurl)
189-
vuln_pkgs_ref = models.Vulnerability_Package_Relation(package=pkg,vulnerability=vuln,is_vulnerable=False)
190-
191-
if pkg_created or vuln_created :
192-
# For a `Vulnerability_Package_Relation` tp exist it needs both, the package and
193-
# vulnerability to already exist in the db.
194-
self.bulk_create_vuln_pkg_refs.append(vuln_pkgs_ref)
195-
196-
else:
197-
qs = models.Vulnerability_Package_Relation.objects.filter(package=pkg,vulnerability=vuln)
198-
if qs:
199-
if qs[0].is_vulnerable :
200-
qs[0].is_vulnerable = False
201-
self.bulk_update_vuln_pkg_refs.append(qs[0])
202-
203-
else:
204-
self.bulk_create_vuln_pkg_refs.append(vuln_pkgs_ref)
101+
if vuln_created:
102+
# This means vulnerability didn't previously exist in the db, so bulk create
103+
# is used without any hesitation
104+
for id_ in set(advisory.reference_ids):
105+
bulk_create_vuln_refs.append(models.VulnerabilityReference(vulnerability=vuln, reference_id=id_))
106+
107+
for url in set(advisory.reference_urls):
108+
bulk_create_vuln_refs.append(models.VulnerabilityReference(vulnerability=vuln, url=url))
109+
110+
else:
111+
vuln_refs_qs = models.VulnerabilityReference.objects.filter(vulnerability=vuln)
112+
# This is to avoid making additional SELECT queries, and do further filtering in python
113+
# Check https://stackoverflow.com/a/5989530
114+
vuln_ids = {ref.reference_id for ref in vuln_refs_qs}
115+
vuln_urls = {ref.url for ref in vuln_refs_qs}
116+
117+
for id_ in advisory.reference_ids:
118+
if id_ not in vuln_ids :
119+
# Delete the item because it will allow the duplicates pass to through if they are not
120+
# present in vuln_ids
121+
vuln_ids.add(id_)
122+
bulk_create_vuln_refs.append(models.VulnerabilityReference(vulnerability=vuln, reference_id=id_))
123+
124+
for url in advisory.reference_urls:
125+
# Delete the item because it will allow the duplicates pass to through if they are not
126+
# present in vuln_urls
127+
if url not in vuln_urls:
128+
vuln_urls.add(url)
129+
bulk_create_vuln_refs.append(models.VulnerabilityReference(vulnerability=vuln, url=url))
130+
131+
models.VulnerabilityReference.objects.bulk_create(bulk_create_vuln_refs)
132+
133+
for ipurl in advisory.impacted_package_urls :
134+
pkg, pkg_created = _get_or_create_package(ipurl)
135+
vuln_pkgs_ref = models.Vulnerability_Package_Relation(package=pkg,vulnerability=vuln,is_vulnerable=True)
136+
137+
if pkg_created or vuln_created :
138+
bulk_create_vuln_pkg_refs.append(vuln_pkgs_ref)
139+
140+
else:
141+
qs = models.Vulnerability_Package_Relation.objects.filter(package=pkg,vulnerability=vuln)
142+
if qs:
143+
if not qs[0].is_vulnerable:
144+
qs[0].is_vulnerable = True
145+
bulk_update_vuln_pkg_refs.append(qs[0])
146+
147+
else:
148+
bulk_create_vuln_pkg_refs.append(vuln_pkgs_ref)
149+
150+
for rpurl in advisory.resolved_package_urls :
151+
pkg, pkg_created = _get_or_create_package(rpurl)
152+
vuln_pkgs_ref = models.Vulnerability_Package_Relation(package=pkg,vulnerability=vuln,is_vulnerable=False)
153+
154+
if pkg_created or vuln_created :
155+
# For a `Vulnerability_Package_Relation` tp exist it needs both, the package and
156+
# vulnerability to already exist in the db.
157+
bulk_create_vuln_pkg_refs.append(vuln_pkgs_ref)
158+
159+
else:
160+
qs = models.Vulnerability_Package_Relation.objects.filter(package=pkg,vulnerability=vuln)
161+
if qs:
162+
if qs[0].is_vulnerable :
163+
qs[0].is_vulnerable = False
164+
bulk_update_vuln_pkg_refs.append(qs[0])
165+
166+
else:
167+
bulk_create_vuln_pkg_refs.append(vuln_pkgs_ref)
168+
169+
models.Vulnerability_Package_Relation.objects.bulk_update(bulk_update_vuln_pkg_refs, ['is_vulnerable'])
170+
models.Vulnerability_Package_Relation.objects.bulk_create(bulk_create_vuln_pkg_refs)
205171

206172
def _process_added_advisories(data_source: DataSource) -> None:
207173
for batch in data_source.added_advisories():

0 commit comments

Comments
 (0)