Skip to content

Commit f5bd3bc

Browse files
sbs2001sbs2001
committed
♻️ Refactor serailizer methods into nested serializer
Signed-off-by: Shivam Sandbhor <shivam.sandbhor@gmail.com>
1 parent 4ae838c commit f5bd3bc

2 files changed

Lines changed: 73 additions & 116 deletions

File tree

‎vulnerabilities/api.py‎

Lines changed: 41 additions & 80 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
#
2-
# Copyright (c) 2017 nexB Inc. and others. All rights reserved.
2+
# Copyright (c) nexB Inc. and others. All rights reserved.
33
# http://nexb.com and https://github.com/nexB/vulnerablecode/
44
# The VulnerableCode software is licensed under the Apache License version 2.0.
55
# Data generated with VulnerableCode require an acknowledgment.
@@ -18,7 +18,7 @@
1818
# OR CONDITIONS OF ANY KIND, either express or implied. No content created from
1919
# VulnerableCode should be considered or used as legal advice. Consult an Attorney
2020
# for any legal advice.
21-
# VulnerableCode is a free software code scanning tool from nexB Inc. and others.
21+
# VulnerableCode is a free software tool from nexB Inc. and others.
2222
# Visit https://github.com/nexB/vulnerablecode/ for support and download.
2323

2424
from urllib.parse import unquote
@@ -44,49 +44,43 @@ class Meta:
4444
]
4545

4646

47+
class HyperLinkedPackageSerializer(serializers.HyperlinkedModelSerializer):
48+
purl = serializers.CharField(source="package_url")
49+
50+
class Meta:
51+
model = Package
52+
fields = ["url", "purl"]
53+
54+
55+
class HyperLinkedVulnerabilitySerializer(serializers.HyperlinkedModelSerializer):
56+
vulnerability_id = serializers.CharField(source="cve_id")
57+
58+
class Meta:
59+
model = Vulnerability
60+
fields = ["url", "vulnerability_id"]
61+
62+
4763
class VulnerabilitySerializer(serializers.ModelSerializer):
4864
references = VulnerabilityReferenceSerializer(many=True, source="vulnerabilityreference_set")
49-
resolved_packages = serializers.SerializerMethodField()
50-
unresolved_packages = serializers.SerializerMethodField()
65+
resolved_packages = HyperLinkedPackageSerializer(
66+
many=True, source="resolved_to", read_only=True
67+
)
68+
unresolved_packages = HyperLinkedPackageSerializer(
69+
many=True, source="vulnerable_to", read_only=True
70+
)
5171

5272
class Meta:
5373
model = Vulnerability
5474
fields = "__all__"
5575

56-
def get_resolved_packages(self, vulnerability):
57-
request = self.context.get("request")
58-
59-
# Instead of iterating and querying the db for eve id FOR EACH item
60-
# in `resolved_to`, consider prefetching all cve_id of `resolved_to` in one query.
61-
return [
62-
{
63-
"purl": Package.objects.get(id=rel.package_id).package_url,
64-
"url": request.build_absolute_uri(
65-
reverse("package-detail", kwargs={"pk": rel.package_id})
66-
),
67-
}
68-
for rel in vulnerability.resolved_to
69-
]
70-
71-
def get_unresolved_packages(self, vulnerability):
72-
request = self.context.get("request")
73-
74-
# Instead of iterating and querying the db for eve id FOR EACH item
75-
# in `vulnerable_to`, consider prefetching all cve_id of vulnerable_to` in one query.
76-
return [
77-
{
78-
"purl": Package.objects.get(id=rel.package_id).package_url,
79-
"url": request.build_absolute_uri(
80-
reverse("package-detail", kwargs={"pk": rel.package_id})
81-
),
82-
}
83-
for rel in vulnerability.vulnerable_to
84-
]
85-
8676

8777
class PackageSerializer(serializers.ModelSerializer):
88-
unresolved_vulnerabilities = serializers.SerializerMethodField()
89-
resolved_vulnerabilities = serializers.SerializerMethodField()
78+
unresolved_vulnerabilities = HyperLinkedVulnerabilitySerializer(
79+
many=True, source="vulnerable_to", read_only=True
80+
)
81+
resolved_vulnerabilities = HyperLinkedVulnerabilitySerializer(
82+
many=True, source="resolved_to", read_only=True
83+
)
9084
purl = serializers.CharField(source="package_url")
9185

9286
class Meta:
@@ -103,36 +97,6 @@ class Meta:
10397
"unresolved_vulnerabilities",
10498
]
10599

106-
def get_unresolved_vulnerabilities(self, package):
107-
request = self.context["request"]
108-
109-
# Instead of iterating and querying the db for eve id FOR EACH item
110-
# in `vulnerable_to`, consider prefetching all cve_id of vulnerable_to` in one query.
111-
return [
112-
{
113-
"vulnerability_id": Vulnerability.objects.get(id=i.vulnerability_id).cve_id,
114-
"url": request.build_absolute_uri(
115-
reverse("vulnerability-detail", kwargs={"pk": i.vulnerability_id})
116-
),
117-
}
118-
for i in package.vulnerable_to
119-
]
120-
121-
def get_resolved_vulnerabilities(self, package):
122-
request = self.context["request"]
123-
124-
# Instead of iterating and querying the db for eve id FOR EACH item
125-
# in `resolved_to`, consider prefetching all cve_id of `resolved_to` in one query.
126-
return [
127-
{
128-
"vulnerability_id": Vulnerability.objects.get(id=i.vulnerability_id).cve_id,
129-
"url": request.build_absolute_uri(
130-
reverse("vulnerability-detail", kwargs={"pk": i.vulnerability_id})
131-
),
132-
}
133-
for i in package.resolved_to
134-
]
135-
136100

137101
class PackageFilterSet(filters.FilterSet):
138102
purl = filters.CharFilter(method="filter_purl")
@@ -162,20 +126,17 @@ class PackageViewSet(viewsets.ReadOnlyModelViewSet):
162126
filterset_class = PackageFilterSet
163127

164128

165-
class VulnerabilityViewSet(viewsets.ReadOnlyModelViewSet):
166-
serializer_class = VulnerabilitySerializer
167-
paginate_by = 50
129+
class VulnerabilityFilterSet(filters.FilterSet):
130+
vulnerability_id = filters.CharFilter(field_name="cve_id")
168131

169-
def get_queryset(self):
170-
if "vulnerability_id" in self.request.query_params:
171-
return Vulnerability.objects.filter(
172-
cve_id__contains=self.request.query_params["vulnerability_id"]
173-
)
132+
class Meta:
133+
model = Vulnerability
134+
fields = ["vulnerability_id"]
174135

175-
return Vulnerability.objects.all()
176136

177-
def get_serializer_context(self):
178-
context = super(VulnerabilityViewSet, self).get_serializer_context()
179-
# Passing this context allows construction of absolute urls.
180-
context.update({"request": self.request})
181-
return context
137+
class VulnerabilityViewSet(viewsets.ReadOnlyModelViewSet):
138+
queryset = Vulnerability.objects.all()
139+
serializer_class = VulnerabilitySerializer
140+
paginate_by = 50
141+
filter_backends = (filters.DjangoFilterBackend,)
142+
filterset_class = VulnerabilityFilterSet

‎vulnerabilities/models.py‎

Lines changed: 32 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
#
2-
# Copyright (c) 2017 nexB Inc. and others. All rights reserved.
2+
# Copyright (c) nexB Inc. and others. All rights reserved.
33
# http://nexb.com and https://github.com/nexB/vulnerablecode/
44
# The VulnerableCode software is licensed under the Apache License version 2.0.
55
# Data generated with VulnerableCode require an acknowledgment.
@@ -18,7 +18,7 @@
1818
# OR CONDITIONS OF ANY KIND, either express or implied. No content created from
1919
# VulnerableCode should be considered or used as legal advice. Consult an Attorney
2020
# for any legal advice.
21-
# VulnerableCode is a free software code scanning tool from nexB Inc. and others.
21+
# VulnerableCode is a free software tool from nexB Inc. and others.
2222
# Visit https://github.com/nexB/vulnerablecode/ for support and download.
2323

2424
import importlib
@@ -39,11 +39,10 @@ class Vulnerability(models.Model):
3939
A software vulnerability with minimal information. Identifiers other than CVE ID are stored as
4040
VulnerabilityReference.
4141
"""
42-
cve_id = models.CharField(
43-
max_length=50, help_text='CVE ID', unique=True, null=True)
44-
summary = models.TextField(
45-
help_text='Summary of the vulnerability', blank=True)
46-
cvss = models.FloatField(max_length=100, help_text='CVSS Score', null=True)
42+
43+
cve_id = models.CharField(max_length=50, help_text="CVE ID", unique=True, null=True)
44+
summary = models.TextField(help_text="Summary of the vulnerability", blank=True)
45+
cvss = models.FloatField(max_length=100, help_text="CVSS Score", null=True)
4746

4847
@property
4948
def vulnerable_to(self):
@@ -57,36 +56,37 @@ def __str__(self):
5756
return self.cve_id or self.summary
5857

5958
class Meta:
60-
verbose_name_plural = 'Vulnerabilities'
59+
verbose_name_plural = "Vulnerabilities"
6160

6261

6362
class VulnerabilityReference(models.Model):
6463
"""
6564
A reference to a vulnerability such as a security advisory from a Linux distribution or language
6665
package manager.
6766
"""
68-
vulnerability = models.ForeignKey(
69-
Vulnerability, on_delete=models.CASCADE)
70-
source = models.CharField(
71-
max_length=50, help_text='Source(s) name eg:NVD', blank=True)
67+
68+
vulnerability = models.ForeignKey(Vulnerability, on_delete=models.CASCADE)
69+
source = models.CharField(max_length=50, help_text="Source(s) name eg:NVD", blank=True)
7270
reference_id = models.CharField(
73-
max_length=50, help_text='Reference ID, eg:DSA-4465-1', blank=True)
74-
url = models.URLField(
75-
max_length=1024, help_text='URL of Vulnerability data', blank=True)
71+
max_length=50, help_text="Reference ID, eg:DSA-4465-1", blank=True
72+
)
73+
url = models.URLField(max_length=1024, help_text="URL of Vulnerability data", blank=True)
7674

7775
class Meta:
78-
unique_together = ('vulnerability', 'source', 'reference_id', 'url')
76+
unique_together = ("vulnerability", "source", "reference_id", "url")
7977

8078
def __str__(self):
81-
return f'{self.source} {self.reference_id} {self.url}'
79+
return f"{self.source} {self.reference_id} {self.url}"
8280

8381

8482
class Package(PackageURLMixin):
8583
"""
8684
A software package with links to relevant vulnerabilities.
8785
"""
86+
8887
vulnerabilities = models.ManyToManyField(
89-
to='Vulnerability', through='PackageRelatedVulnerability')
88+
to="Vulnerability", through="PackageRelatedVulnerability"
89+
)
9090

9191
@property
9292
def vulnerable_to(self):
@@ -97,17 +97,17 @@ def resolved_to(self):
9797
return self.packagerelatedvulnerability_set.filter(is_vulnerable=False)
9898

9999
class Meta:
100-
unique_together = ('name', 'namespace', 'type',
101-
'version', 'qualifiers', 'subpath')
100+
unique_together = ("name", "namespace", "type", "version", "qualifiers", "subpath")
101+
102102
# Remove the `qualifers` and `set_package_url` overrides after
103103
# https://github.com/package-url/packageurl-python/pull/35 gets merged
104104
qualifiers = pgfields.JSONField(
105105
default=dict,
106106
help_text=_(
107-
'Extra qualifying data for a package such as the name of an OS, '
108-
'architecture, distro, etc.'
107+
"Extra qualifying data for a package such as the name of an OS, "
108+
"architecture, distro, etc."
109109
),
110-
null=True
110+
null=True,
111111
)
112112

113113
def set_package_url(self, package_url):
@@ -123,8 +123,7 @@ def set_package_url(self, package_url):
123123
model_field = self._meta.get_field(field_name)
124124

125125
if value and len(value) > model_field.max_length:
126-
raise ValidationError(
127-
_('Value too long for field "{}".'.format(field_name)))
126+
raise ValidationError(_('Value too long for field "{}".'.format(field_name)))
128127

129128
setattr(self, field_name, value or None)
130129

@@ -142,7 +141,7 @@ class Meta:
142141
# Technically 'is_vulnerable' doesn't belong here. The idea is to
143142
# later filter out for a pairs of ('package', 'vulnerability') which have both
144143
# values of 'is_vulnerable' and ping the data providers to resolve such entries.
145-
unique_together = ('package', 'vulnerability', 'is_vulnerable')
144+
unique_together = ("package", "vulnerability", "is_vulnerable")
146145

147146

148147
class ImportProblem(models.Model):
@@ -155,26 +154,23 @@ class Importer(models.Model):
155154
Metadata and pointer to the implementation for a source of vulnerability data (aka security
156155
advisories)
157156
"""
158-
name = models.CharField(max_length=100, unique=True,
159-
help_text='Name of the importer')
157+
158+
name = models.CharField(max_length=100, unique=True, help_text="Name of the importer")
160159

161160
license = models.CharField(
162-
max_length=100,
163-
blank=True,
164-
help_text='License of the vulnerability data',
161+
max_length=100, blank=True, help_text="License of the vulnerability data",
165162
)
166163

167-
last_run = models.DateTimeField(
168-
null=True, help_text='UTC Timestamp of the last run')
164+
last_run = models.DateTimeField(null=True, help_text="UTC Timestamp of the last run")
169165

170166
data_source = models.CharField(
171167
max_length=100,
172-
help_text='Name of the data source implementation importable from vulnerabilities.importers'
168+
help_text="Name of the data source implementation importable from vulnerabilities.importers", # nopep8
173169
)
174170
data_source_cfg = pgfields.JSONField(
175171
null=False,
176172
default=dict,
177-
help_text='Implementation-specific configuration for the data source',
173+
help_text="Implementation-specific configuration for the data source",
178174
)
179175

180176
def make_data_source(self, batch_size: int, cutoff_date: datetime = None) -> DataSource:
@@ -184,7 +180,7 @@ def make_data_source(self, batch_size: int, cutoff_date: datetime = None) -> Dat
184180
batch_size - max. number of records to return on each iteration
185181
cutoff_date - optional timestamp of the oldest data to include in the import
186182
"""
187-
importers_module = importlib.import_module('vulnerabilities.importers')
183+
importers_module = importlib.import_module("vulnerabilities.importers")
188184
klass = getattr(importers_module, self.data_source)
189185

190186
ds = klass(

0 commit comments

Comments
 (0)