Make use of get_or_create to reduce the number of objects created

Fixes #40

Signed-off-by: Ranvir Singh <[email protected]>

Author: singh1114

vulnerabilities/data_dump.py

@@ -34,39 +34,39 @@ def debian_dump(extract_data):
     Save data scraped from Debian' security tracker.
     """
     for data in extract_data:
-        vulnerability = Vulnerability.objects.create(
+        vulnerability, _ = Vulnerability.objects.get_or_create(
             summary=data.get('description', ''),
         )
-        VulnerabilityReference.objects.create(
+        VulnerabilityReference.objects.get_or_create(
             vulnerability=vulnerability,
             reference_id=data.get('vulnerability_id', ''),
         )
 
         pkg_name = data.get('package_name', '')
-        package = Package.objects.create(
+        package, _ = Package.objects.get_or_create(
             name=pkg_name,
             version=data.get('version', ''),
         )
 
         if data['status'] == 'open':
-            ImpactedPackage.objects.create(
+            ImpactedPackage.objects.get_or_create(
                 vulnerability=vulnerability,
                 package=package
             )
         else:
-            ResolvedPackage.objects.create(
+            ResolvedPackage.objects.get_or_create(
                 vulnerability=vulnerability,
                 package=package
             )
 
             fixed_version = data.get('fixed_version')
             if fixed_version:
-                package = Package.objects.create(
+                package, _ = Package.objects.get_or_create(
                     name=pkg_name,
                     version=fixed_version,
                 )
 
-                ResolvedPackage.objects.create(
+                ResolvedPackage.objects.get_or_create(
                     vulnerability=vulnerability,
                     package=package
                 )
@@ -77,17 +77,19 @@ def ubuntu_dump(html):
     Dump data scraped from Ubuntu's security tracker.
     """
     for data in html:
-        vulnerability = Vulnerability.objects.create(
-            summary='',
+        # If we don't do this, every VulnerabilityReference will have single
+        # Vulnerability.
+        vulnerability, _ = Vulnerability.objects.get_or_create(
+            summary='{} vulnerability'.format(data.get('cve_id')),
         )
-        VulnerabilityReference.objects.create(
+        VulnerabilityReference.objects.get_or_create(
             vulnerability=vulnerability,
             reference_id=data.get('cve_id'),
         )
-        package = Package.objects.create(
+        package, _ = Package.objects.get_or_create(
             name=data.get('package_name'),
         )
-        ImpactedPackage.objects.create(
+        ImpactedPackage.objects.get_or_create(
             vulnerability=vulnerability,
             package=package
         )
@@ -111,39 +113,39 @@ def archlinux_dump(extract_data):
         if not fixed_version:
             fixed_version = 'None'
 
-        vulnerability = Vulnerability.objects.create(
+        vulnerability, _ = Vulnerability.objects.get_or_create(
             summary=item['type'],
         )
 
         for vulnerability_id in vulnerabilities:
-            VulnerabilityReference.objects.create(
+            VulnerabilityReference.objects.get_or_create(
                 vulnerability=vulnerability,
                 reference_id=vulnerability_id,
                 url='https://security.archlinux.org/{}'.format(vulnerability_id)
             )
 
         for package_name in packages_name:
-            package_affected = Package.objects.create(
+            package_affected, _ = Package.objects.get_or_create(
                 name=package_name,
                 version=affected_version
             )
-            ImpactedPackage.objects.create(
+            ImpactedPackage.objects.get_or_create(
                 vulnerability=vulnerability,
                 package=package_affected
             )
-            PackageReference.objects.create(
+            PackageReference.objects.get_or_create(
                 package=package_affected,
                 repository='https://security.archlinux.org/package/{}'.format(package_name)
             )
-            package_fixed = Package.objects.create(
+            package_fixed, _ = Package.objects.get_or_create(
                 name=package_name,
                 version=fixed_version
             )
-            ResolvedPackage.objects.create(
+            ResolvedPackage.objects.get_or_create(
                 vulnerability=vulnerability,
                 package=package_fixed
             )
-            PackageReference.objects.create(
+            PackageReference.objects.get_or_create(
                 package=package_fixed,
                 repository='https://security.archlinux.org/package/{}'.format(package_name)
             )

vulnerabilities/tests/test_api.py

@@ -47,7 +47,7 @@ def test_debian_response(self):
         debian_dump(extract_data)
         response = self.client.get('/api/packages/?name=mimetex', format='json').data
 
-        self.assertEqual(4, response['count'])
+        self.assertEqual(2, response['count'])
 
         first_result = response['results'][0]
         self.assertEqual('mimetex', first_result['name'])
@@ -69,7 +69,7 @@ def test_ubuntu_response(self):
             "version": "",
             "platform": "",
             "vulnerabilities": [{
-                "summary": "",
+                "summary": "CVE-2012-3386 vulnerability",
                 "cvss": None,
                 "references": [{
                     "reference_id": "CVE-2012-3386",
@@ -93,7 +93,7 @@ def test_serializers(self):
         pk = Package.objects.filter(name="mimetex")
         response = PackageSerializer(pk, many=True).data
 
-        self.assertEqual(4, len(response))
+        self.assertEqual(2, len(response))
 
         first_result = response[0]
         self.assertEqual('mimetex', first_result['name'])

vulnerabilities/tests/test_data_dump.py

@@ -51,8 +51,8 @@ def setUpTestData(self):
         with open(os.path.join(TEST_DATA, 'debian.json')) as f:
             test_data = json.load(f)
 
-        extract_data = debian.extract_vulnerabilities(test_data)
-        debian_dump(extract_data)
+        self.extract_data = debian.extract_vulnerabilities(test_data)
+        debian_dump(self.extract_data)
 
     def test_Vulnerability(self):
         """
@@ -83,9 +83,7 @@ def test_Package(self):
         """
         Check that all packages from the test data are stored in the database
         """
-        # There are five rows in Package because currently the models allow duplicates
-        # (see issue #28).
-        self.assertEqual(5, Package.objects.count())
+        self.assertEqual(3, Package.objects.count())
 
         self.assertTrue(Package.objects.filter(name='mimetex'))
         self.assertTrue(Package.objects.get(name='librsync'))
@@ -114,15 +112,33 @@ def test_ResolvedPackage(self):
         self.assertIn('1.50-1.1', versions)
         self.assertIn('1.74-1', versions)
 
+    def test_debian_data_dump_twice(self):
+        """
+        Scrape data from Debian' main tracker, save it
+        in the database and verify entries.
+        """
+        debian_dump(self.extract_data)
+
+        self.assertEqual(3, Vulnerability.objects.count())
+        self.assertEqual(3, VulnerabilityReference.objects.count())
+        self.assertEqual(3, Package.objects.count())
+
+        # Dumping the data twice doesn't create new objects.
+        debian_dump(self.extract_data)
+
+        self.assertEqual(3, Vulnerability.objects.count())
+        self.assertEqual(3, VulnerabilityReference.objects.count())
+        self.assertEqual(3, Package.objects.count())
+
 
 class TestUbuntuDataDump(TestCase):
     @classmethod
     def setUpTestData(self):
         with open(os.path.join(TEST_DATA, 'ubuntu_main.html')) as f:
             test_data = f.read()
 
-        data = ubuntu.extract_cves(test_data)
-        ubuntu_dump(data)
+        self.data = ubuntu.extract_cves(test_data)
+        ubuntu_dump(self.data)
 
     def test_data_dump(self):
         """
@@ -132,15 +148,32 @@ def test_data_dump(self):
         self.assertEqual(reference.reference_id, 'CVE-2002-2439')
         self.assertTrue(Package.objects.filter(name='gcc-4.6')[0].name, 'gcc-4.6')
 
+    def test_ubuntu_data_dump_twice(self):
+        """
+        Scrape data from Ubuntu twice from main tracker, save it
+        in the database and verify single time entry.
+        """
+        ubuntu_dump(self.data)
+        count = Package.objects.all().count()
+        reference = VulnerabilityReference.objects.filter(
+            reference_id='CVE-2002-2439')
+        self.assertEqual(reference[0].reference_id, 'CVE-2002-2439')
+        self.assertTrue(Package.objects.filter(name='gcc-4.6')[0].name,
+                        'gcc-4.6')
+
+        # Dumping the data twice doesn't create new objects.
+        ubuntu_dump(self.data)
+        self.assertEqual(count, Package.objects.all().count())
+
 
 class TestArchLinuxDataDump(TestCase):
 
     @classmethod
     def setUpTestData(self):
         with open(os.path.join(TEST_DATA, 'archlinux.json')) as f:
-            test_data = json.load(f)
+            self.test_data = json.load(f)
 
-        archlinux_dump(test_data)
+        archlinux_dump(self.test_data)
 
     def test_Vulnerability(self):
         """
@@ -190,3 +223,25 @@ def test_ResolvedPackage(self):
 
         self.assertEqual(4, len(resolved_pkgs))
         self.assertEqual('2.6.1-1', resolved_pkg.package.version)
+
+    def test_archlinux_data_dump_twice(self):
+        """
+        Scrape data from Archlinux' main tracker twice, save it
+        in the database and verify no multiple entries.
+        """
+        archlinux_dump(self.test_data)
+        self.assertEqual(1, Vulnerability.objects.count())
+        self.assertEqual(14, VulnerabilityReference.objects.count())
+        self.assertEqual(8, Package.objects.count())
+        self.assertEqual(8, PackageReference.objects.count())
+        self.assertEqual(4, ImpactedPackage.objects.count())
+        self.assertEqual(4, ResolvedPackage.objects.count())
+
+        # Dumping the data twice doesn't create new objects.
+        archlinux_dump(self.test_data)
+        self.assertEqual(1, Vulnerability.objects.count())
+        self.assertEqual(14, VulnerabilityReference.objects.count())
+        self.assertEqual(8, Package.objects.count())
+        self.assertEqual(8, PackageReference.objects.count())
+        self.assertEqual(4, ImpactedPackage.objects.count())
+        self.assertEqual(4, ResolvedPackage.objects.count())