Skip to content

Commit

Permalink
Fix tests
Browse files Browse the repository at this point in the history
Signed-off-by: Tushar Goel <tushar.goel.dav@gmail.com>
  • Loading branch information
TG1999 committed Oct 15, 2024
1 parent 3388081 commit d4403d9
Show file tree
Hide file tree
Showing 14 changed files with 123 additions and 274 deletions.
4 changes: 0 additions & 4 deletions vulnerabilities/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -619,10 +619,6 @@ def get_queryset(self):
super()
.get_queryset()
.prefetch_related(
# Prefetch(
# "packages",
# queryset=self.get_packages_qs(),
# ),
"weaknesses",
Prefetch(
"fixed_by_packages",
Expand Down
2 changes: 0 additions & 2 deletions vulnerabilities/api_extension.py
Original file line number Diff line number Diff line change
Expand Up @@ -238,8 +238,6 @@ class Meta:
"qualifiers",
"subpath",
"purl",
# this hurts
"packagerelatedvulnerability__fix",
]

def filter_purl(self, queryset, name, value):
Expand Down
2 changes: 2 additions & 0 deletions vulnerabilities/import_runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,9 @@
from vulnerabilities.improver import Inference
from vulnerabilities.improvers.default import DefaultImporter
from vulnerabilities.models import Advisory
from vulnerabilities.models import AffectedByPackageRelatedVulnerability
from vulnerabilities.models import Alias
from vulnerabilities.models import FixingPackageRelatedVulnerability
from vulnerabilities.models import Package
from vulnerabilities.models import Vulnerability
from vulnerabilities.models import VulnerabilityChangeLog
Expand Down
6 changes: 3 additions & 3 deletions vulnerabilities/improve_runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,9 +17,10 @@
from vulnerabilities.importers import IMPORTERS_REGISTRY
from vulnerabilities.improver import Inference
from vulnerabilities.models import Advisory
from vulnerabilities.models import AffectedByPackageRelatedVulnerability
from vulnerabilities.models import Alias
from vulnerabilities.models import FixingPackageRelatedVulnerability
from vulnerabilities.models import Package
from vulnerabilities.models import PackageRelatedVulnerability
from vulnerabilities.models import Vulnerability
from vulnerabilities.models import VulnerabilityChangeLog
from vulnerabilities.models import VulnerabilityReference
Expand Down Expand Up @@ -148,12 +149,11 @@ def process_inferences(
fixed_package, created = Package.objects.get_or_create_from_purl(
purl=inference.fixed_purl
)
PackageRelatedVulnerability(
FixingPackageRelatedVulnerability(
vulnerability=vulnerability,
package=fixed_package,
created_by=improver_name,
confidence=inference.confidence,
fix=True,
).update_or_create(
advisory=advisory,
)
Expand Down
14 changes: 9 additions & 5 deletions vulnerabilities/management/commands/export.py
Original file line number Diff line number Diff line change
Expand Up @@ -116,7 +116,7 @@ def export_data(self, base_path: Path):
}
package_vulnerabilities.append(package_data)

for vuln in pkg_version.vulnerabilities.all():
for vuln in pkg_version.vulnerabilities:
vcid = vuln.vulnerability_id
# do not write twice the same file
if vcid in seen_vcid:
Expand Down Expand Up @@ -158,10 +158,14 @@ def packages_by_type_ns_name():
qs = (
Package.objects.order_by("type", "namespace", "name", "version")
.prefetch_related(
"vulnerabilities",
"vulnerabilities__references",
"vulnerabilities__weaknesses",
"vulnerabilities__references__vulnerabilityseverity_set",
"affected_by_vulnerabilities",
"affected_by_vulnerabilities__references",
"affected_by_vulnerabilities__weaknesses",
"affected_by_vulnerabilities__references__vulnerabilityseverity_set",
"fixing_vulnerabilities",
"fixing_vulnerabilities__references",
"fixing_vulnerabilities__weaknesses",
"fixing_vulnerabilities__references__vulnerabilityseverity_set",
)
.paginated()
)
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,16 @@
# Generated by Django 4.2.15 on 2024-10-15 10:37

from django.db import migrations


class Migration(migrations.Migration):

dependencies = [
("vulnerabilities", "0072_remove_package_vulnerabilities_and_more"),
]

operations = [
migrations.DeleteModel(
name="PackageRelatedVulnerability",
),
]
131 changes: 14 additions & 117 deletions vulnerabilities/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -87,7 +87,7 @@ def affecting_vulnerabilities(self):
"""
Return a queryset of Vulnerability that affect a package.
"""
return self.filter(packagerelatedvulnerability__fix=False)
return self.filter(affecting_packages__isnull=False)

def with_cpes(self):
"""
Expand Down Expand Up @@ -187,12 +187,6 @@ class Vulnerability(models.Model):
to="VulnerabilityReference", through="VulnerabilityRelatedReference"
)

# packages = models.ManyToManyField(
# to="Package",
# through="PackageRelatedVulnerability",
# related_name="all_vulnerabilities",
# )

affecting_packages = models.ManyToManyField(
to="Package",
through="AffectedByPackageRelatedVulnerability",
Expand Down Expand Up @@ -452,7 +446,7 @@ def get_fixed_by_package_versions(self, purl: PackageURL, fix=True):
}

if fix:
filter_dict["packagerelatedvulnerability__fix"] = True
filter_dict["fixing_vulnerabilities__isnull"] = False

return Package.objects.filter(**filter_dict).distinct()

Expand All @@ -474,15 +468,15 @@ def affected(self):
"""
Return only packages affected by a vulnerability.
"""
return self.filter(packagerelatedvulnerability__fix=False)
return self.filter(affected_by_vulnerabilities__isnull=False)

vulnerable = affected

def fixing(self):
"""
Return only packages fixing a vulnerability .
"""
return self.filter(packagerelatedvulnerability__fix=True)
return self.filter(fixing_vulnerabilities__isnull=False)

def with_vulnerability_counts(self):
return self.annotate(
Expand Down Expand Up @@ -585,6 +579,12 @@ def _vulnerable(self, vulnerable=True):
"""
return self.with_is_vulnerable().filter(is_vulnerable=vulnerable)

def vulnerable(self):
"""
Return only packages that are vulnerable.
"""
return self.filter(affected_by_vulnerabilities__isnull=False)


def get_purl_query_lookups(purl):
"""
Expand All @@ -606,12 +606,6 @@ class Package(PackageURLMixin):
# https://github.com/package-url/packageurl-python/pull/67
# gets merged

# vulnerabilities = models.ManyToManyField(
# to="Vulnerability",
# through="PackageRelatedVulnerability",
# related_name="all_packages",
# )

affected_by_vulnerabilities = models.ManyToManyField(
to="Vulnerability",
through="AffectedByPackageRelatedVulnerability",
Expand Down Expand Up @@ -751,6 +745,10 @@ def next_non_vulnerable_version(self):
next_non_vulnerable, _ = self.get_non_vulnerable_versions()
return next_non_vulnerable.version if next_non_vulnerable else None

@property
def vulnerabilities(self):
return self.affected_by_vulnerabilities.all() | self.fixing_vulnerabilities.all()

@property
def latest_non_vulnerable_version(self):
"""
Expand Down Expand Up @@ -880,107 +878,6 @@ def affecting_vulns(self):
)


class PackageRelatedVulnerability(models.Model):
"""
Track the relationship between a Package and Vulnerability.
"""

# TODO: Fix related_name
package = models.ForeignKey(
Package,
on_delete=models.CASCADE,
)

vulnerability = models.ForeignKey(
Vulnerability,
on_delete=models.CASCADE,
)

created_by = models.CharField(
max_length=100,
blank=True,
help_text="Fully qualified name of the improver prefixed with the"
"module name responsible for creating this relation. Eg:"
"vulnerabilities.importers.nginx.NginxBasicImprover",
)

from vulnerabilities.improver import MAX_CONFIDENCE

confidence = models.PositiveIntegerField(
default=MAX_CONFIDENCE,
validators=[MinValueValidator(0), MaxValueValidator(MAX_CONFIDENCE)],
help_text="Confidence score for this relation",
)

fix = models.BooleanField(
default=False,
db_index=True,
help_text="Does this relation fix the specified vulnerability ?",
)

class Meta:
unique_together = ["package", "vulnerability"]
verbose_name_plural = "PackageRelatedVulnerabilities"
indexes = [models.Index(fields=["fix"])]
ordering = ["package", "vulnerability"]

def __str__(self):
return f"{self.package.package_url} {self.vulnerability.vulnerability_id}"

def update_or_create(self, advisory):
"""
Update if supplied record has more confidence than existing record
Create if doesn't exist
"""
try:
existing = PackageRelatedVulnerability.objects.get(
vulnerability=self.vulnerability, package=self.package
)
if self.confidence > existing.confidence:
existing.created_by = self.created_by
existing.confidence = self.confidence
existing.fix = self.fix
existing.save()
# TODO: later we want these to be part of a log field in the DB
logger.info(
f"Confidence improved for {self.package} R {self.vulnerability}, "
f"new confidence: {self.confidence}"
)
self.add_package_vulnerability_changelog(advisory=advisory)

except self.DoesNotExist:
PackageRelatedVulnerability.objects.create(
vulnerability=self.vulnerability,
created_by=self.created_by,
package=self.package,
confidence=self.confidence,
fix=self.fix,
)

logger.info(
f"New relationship {self.package} R {self.vulnerability}, "
f"fix: {self.fix}, confidence: {self.confidence}"
)

self.add_package_vulnerability_changelog(advisory=advisory)

@transaction.atomic
def add_package_vulnerability_changelog(self, advisory):
from vulnerabilities.utils import get_importer_name

importer_name = get_importer_name(advisory)
if self.fix:
change_logger = PackageChangeLog.log_fixing
else:
change_logger = PackageChangeLog.log_affected_by
change_logger(
package=self.package,
importer=importer_name,
source_url=advisory.url or None,
related_vulnerability=str(self.vulnerability),
)


class PackageRelatedVulnerabilityBase(models.Model):
"""
Abstract base class for package-vulnerability relations.
Expand Down
Loading

0 comments on commit d4403d9

Please sign in to comment.