Skip to content

Commit b303b74

Browse files
singh1114singh1114
committed
Use get or create to reduce the number of objects created
Fixes #40 Signed-off-by: Ranvir Singh <ranvir.singh1114@gmail.com>
1 parent d5ab2fe commit b303b74

3 files changed

Lines changed: 111 additions & 70 deletions

File tree

‎vulnerabilities/data_dump.py‎

Lines changed: 17 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -34,18 +34,18 @@ def debian_dump(extract_data):
3434
Save data scraped from Debian' security tracker.
3535
"""
3636
for data in extract_data:
37-
vulnerability = Vulnerability.objects.create(
37+
vulnerability, _ = Vulnerability.objects.get_or_create(
3838
summary=data.get('description', ''),
3939
)
40-
VulnerabilityReference.objects.create(
40+
VulnerabilityReference.objects.get_or_create(
4141
vulnerability=vulnerability,
4242
reference_id=data.get('vulnerability_id', ''),
4343
)
44-
package = Package.objects.create(
44+
package, _ = Package.objects.get_or_create(
4545
name=data.get('package_name', ''),
4646
version=data.get('fixed_version', ''),
4747
)
48-
ImpactedPackage.objects.create(
48+
ImpactedPackage.objects.get_or_create(
4949
vulnerability=vulnerability,
5050
package=package
5151
)
@@ -56,17 +56,17 @@ def ubuntu_dump(html):
5656
Dump data scraped from Ubuntu's security tracker.
5757
"""
5858
for data in html:
59-
vulnerability = Vulnerability.objects.create(
60-
summary='',
59+
vulnerability, _ = Vulnerability.objects.get_or_create(
60+
summary='{} vulnerability'.format(data.get('cve_id')),
6161
)
62-
VulnerabilityReference.objects.create(
62+
VulnerabilityReference.objects.get_or_create(
6363
vulnerability=vulnerability,
6464
reference_id=data.get('cve_id'),
6565
)
66-
package = Package.objects.create(
66+
package, _ = Package.objects.get_or_create(
6767
name=data.get('package_name'),
6868
)
69-
ImpactedPackage.objects.create(
69+
ImpactedPackage.objects.get_or_create(
7070
vulnerability=vulnerability,
7171
package=package
7272
)
@@ -90,39 +90,39 @@ def archlinux_dump(extract_data):
9090
if not fixed_version:
9191
fixed_version = 'None'
9292

93-
vulnerability = Vulnerability.objects.create(
93+
vulnerability, _ = Vulnerability.objects.get_or_create(
9494
summary=item['type'],
9595
)
9696

9797
for vulnerability_id in vulnerabilities:
98-
VulnerabilityReference.objects.create(
98+
VulnerabilityReference.objects.get_or_create(
9999
vulnerability=vulnerability,
100100
reference_id=vulnerability_id,
101101
url='https://security.archlinux.org/{}'.format(vulnerability_id)
102102
)
103103

104104
for package_name in packages_name:
105-
package_affected = Package.objects.create(
105+
package_affected, _ = Package.objects.get_or_create(
106106
name=package_name,
107107
version=affected_version
108108
)
109-
ImpactedPackage.objects.create(
109+
ImpactedPackage.objects.get_or_create(
110110
vulnerability=vulnerability,
111111
package=package_affected
112112
)
113-
PackageReference.objects.create(
113+
PackageReference.objects.get_or_create(
114114
package=package_affected,
115115
repository='https://security.archlinux.org/package/{}'.format(package_name)
116116
)
117-
package_fixed = Package.objects.create(
117+
package_fixed, _ = Package.objects.get_or_create(
118118
name=package_name,
119119
version=fixed_version
120120
)
121-
ResolvedPackage.objects.create(
121+
ResolvedPackage.objects.get_or_create(
122122
vulnerability=vulnerability,
123123
package=package_fixed
124124
)
125-
PackageReference.objects.create(
125+
PackageReference.objects.get_or_create(
126126
package=package_fixed,
127127
repository='https://security.archlinux.org/package/{}'.format(package_name)
128128
)

‎vulnerabilities/tests/test_api.py‎

Lines changed: 42 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -46,34 +46,31 @@ def test_debian_response(self):
4646
extract_data = debian.extract_vulnerabilities(test_data)
4747
debian_dump(extract_data)
4848
response = self.client.get('/api/packages/?name=mimetex', format='json')
49-
5049
expected = [{
5150
"name": "mimetex",
5251
"version": "1.50-1.1",
5352
"platform": "",
54-
"vulnerabilities": [{
55-
"summary": "Multiple stack-based buffer overflows in mimetex.cgi in mimeTeX",
56-
"cvss": None,
57-
"references": [{
58-
"reference_id": "CVE-2009-2458",
59-
"source": "",
60-
"url": "",
61-
}]
62-
}],
63-
"references": [],
64-
}, {
65-
"name": "mimetex",
66-
"version": "1.50-1.1",
67-
"platform": "",
68-
"vulnerabilities": [{
69-
"summary": "Multiple unspecified vulnerabilities in mimeTeX.",
70-
"cvss": None,
71-
"references": [{
72-
"reference_id": "CVE-2009-2459",
73-
"source": "",
74-
"url": "",
75-
}]
76-
}],
53+
"vulnerabilities": [
54+
{
55+
"summary": "Multiple stack-based buffer overflows in mimetex.cgi in mimeTeX",
56+
"cvss": None,
57+
"references": [{
58+
"reference_id": "CVE-2009-2458",
59+
"source": "",
60+
"url": "",
61+
}]
62+
},
63+
{
64+
"summary": "Multiple unspecified vulnerabilities in mimeTeX.",
65+
"cvss": None,
66+
"references": [{
67+
"reference_id": "CVE-2009-2459",
68+
"source": "",
69+
"url": "",
70+
}]
71+
}
72+
73+
],
7774
"references": [],
7875
}]
7976

@@ -86,13 +83,12 @@ def test_ubuntu_response(self):
8683
extract_data = ubuntu.extract_cves(test_data)
8784
ubuntu_dump(extract_data)
8885
response = self.client.get('/api/packages/?name=automake', format='json')
89-
9086
expected = [{
9187
"name": "automake",
9288
"version": "",
9389
"platform": "",
9490
"vulnerabilities": [{
95-
"summary": "",
91+
"summary": "CVE-2012-3386 vulnerability",
9692
"cvss": None,
9793
"references": [{
9894
"reference_id": "CVE-2012-3386",
@@ -121,30 +117,26 @@ def test_serializers(self):
121117
"name": "mimetex",
122118
"version": "1.50-1.1",
123119
"platform": "",
124-
"vulnerabilities": [{
125-
"summary": "Multiple stack-based buffer overflows in mimetex.cgi in mimeTeX",
126-
"cvss": None,
127-
"references": [{
128-
"reference_id": "CVE-2009-2458",
129-
"source": "",
130-
"url": "",
131-
}]
132-
}],
133-
"references": [],
134-
},
135-
{
136-
"name": "mimetex",
137-
"version": "1.50-1.1",
138-
"platform": "",
139-
"vulnerabilities": [{
140-
"summary": "Multiple unspecified vulnerabilities in mimeTeX.",
141-
"cvss": None,
142-
"references": [{
143-
"reference_id": "CVE-2009-2459",
144-
"source": "",
145-
"url": "",
146-
}]
147-
}],
120+
"vulnerabilities": [
121+
{
122+
"summary": "Multiple stack-based buffer overflows in mimetex.cgi in mimeTeX",
123+
"cvss": None,
124+
"references": [{
125+
"reference_id": "CVE-2009-2458",
126+
"source": "",
127+
"url": "",
128+
}]
129+
},
130+
{
131+
"summary": "Multiple unspecified vulnerabilities in mimeTeX.",
132+
"cvss": None,
133+
"references": [{
134+
"reference_id": "CVE-2009-2459",
135+
"source": "",
136+
"url": "",
137+
}]
138+
}
139+
],
148140
"references": [],
149141
}
150142
]

‎vulnerabilities/tests/test_data_dump.py‎

Lines changed: 52 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,6 @@
3535
from vulnerabilities.models import ResolvedPackage
3636
from vulnerabilities.models import Vulnerability
3737
from vulnerabilities.models import VulnerabilityReference
38-
from vulnerabilities.scraper import archlinux
3938
from vulnerabilities.scraper import debian
4039
from vulnerabilities.scraper import ubuntu
4140

@@ -58,7 +57,7 @@ def test_debian_data_dump(self):
5857

5958
self.assertEqual(3, Vulnerability.objects.count())
6059
self.assertEqual(3, VulnerabilityReference.objects.count())
61-
self.assertEqual(3, Package.objects.count())
60+
self.assertEqual(2, Package.objects.count())
6261

6362
self.assertTrue(Vulnerability.objects.get(
6463
summary='Multiple stack-based buffer overflows in mimetex.cgi in mimeTeX'))
@@ -76,6 +75,22 @@ def test_debian_data_dump(self):
7675
self.assertTrue(Package.objects.get(name='git-repair'))
7776
self.assertEqual(Package.objects.filter(version='1.50-1.1')[0].version, '1.50-1.1')
7877

78+
def test_debian_data_dump_twice(self):
79+
"""
80+
Scrape data from Debian' main tracker, save it
81+
in the database and verify entries.
82+
"""
83+
with open(os.path.join(TEST_DATA, 'debian.json')) as f:
84+
test_data = json.loads(f.read())
85+
86+
extract_data = debian.extract_vulnerabilities(test_data)
87+
debian_dump(extract_data)
88+
debian_dump(extract_data)
89+
90+
self.assertEqual(3, Vulnerability.objects.count())
91+
self.assertEqual(3, VulnerabilityReference.objects.count())
92+
self.assertEqual(2, Package.objects.count())
93+
7994
def test_ubuntu_data_dump(self):
8095
"""
8196
Scrape data from Ubuntu' main tracker, save it
@@ -91,6 +106,23 @@ def test_ubuntu_data_dump(self):
91106
self.assertEqual(reference.reference_id, 'CVE-2002-2439')
92107
self.assertTrue(Package.objects.filter(name='gcc-4.6')[0].name, 'gcc-4.6')
93108

109+
def test_ubuntu_data_dump_twice(self):
110+
"""
111+
Scrape data from Ubuntu twice from main tracker, save it
112+
in the database and verify single time entry.
113+
"""
114+
with open(os.path.join(TEST_DATA, 'ubuntu_main.html')) as f:
115+
test_data = f.read()
116+
117+
data = ubuntu.extract_cves(test_data)
118+
ubuntu_dump(data)
119+
count = Package.objects.all().count()
120+
ubuntu_dump(data)
121+
self.assertEqual(count, Package.objects.all().count())
122+
reference = VulnerabilityReference.objects.filter(reference_id='CVE-2002-2439')
123+
self.assertEqual(reference[0].reference_id, 'CVE-2002-2439')
124+
self.assertTrue(Package.objects.filter(name='gcc-4.6')[0].name, 'gcc-4.6')
125+
94126
def test_archlinux_data_dump(self):
95127
"""
96128
Scrape data from Archlinux' main tracker, save it
@@ -116,4 +148,21 @@ def test_archlinux_data_dump(self):
116148

117149
self.assertTrue(VulnerabilityReference.objects.get(reference_id='AVG-708'))
118150

119-
self.assertEqual(Package.objects.filter(name='wireshark-cli')[0].name, 'wireshark-cli')
151+
self.assertEqual(Package.objects.filter(name='wireshark-cli')[0].name, 'wireshark-cli')
152+
153+
def test_archlinux_data_dump_twice(self):
154+
"""
155+
Scrape data from Archlinux' main tracker twice, save it
156+
in the database and verify no multiple entries.
157+
"""
158+
with open(os.path.join(TEST_DATA, 'archlinux.json')) as f:
159+
test_data = json.loads(f.read())
160+
161+
archlinux_dump(test_data)
162+
archlinux_dump(test_data)
163+
self.assertEqual(1, Vulnerability.objects.count())
164+
self.assertEqual(14, VulnerabilityReference.objects.count())
165+
self.assertEqual(8, Package.objects.count())
166+
self.assertEqual(8, PackageReference.objects.count())
167+
self.assertEqual(4, ImpactedPackage.objects.count())
168+
self.assertEqual(4, ResolvedPackage.objects.count())

0 commit comments

Comments
 (0)