Use get or create to reduce the number of objects created

singh1114 · singh1114 · commit b303b742b739 · 2019-09-25T00:14:37.000+05:30
Fixes #40 Signed-off-by: Ranvir Singh <ranvir.singh1114@gmail.com>
diff --git a/vulnerabilities/data_dump.py b/vulnerabilities/data_dump.py
@@ -34,18 +34,18 @@ 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', ''),
         )
-        package = Package.objects.create(
+        package, _ = Package.objects.get_or_create(
             name=data.get('package_name', ''),
             version=data.get('fixed_version', ''),
         )
-        ImpactedPackage.objects.create(
+        ImpactedPackage.objects.get_or_create(
             vulnerability=vulnerability,
             package=package
         )
@@ -56,17 +56,17 @@ def ubuntu_dump(html):
     Dump data scraped from Ubuntu's security tracker.
     """
     for data in html:
-        vulnerability = Vulnerability.objects.create(
-            summary='',
+        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
         )
@@ -90,39 +90,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)
             )
diff --git a/vulnerabilities/tests/test_api.py b/vulnerabilities/tests/test_api.py
@@ -46,34 +46,31 @@ def test_debian_response(self):
         extract_data = debian.extract_vulnerabilities(test_data)
         debian_dump(extract_data)
         response = self.client.get('/api/packages/?name=mimetex', format='json')
-
         expected = [{
             "name": "mimetex",
             "version": "1.50-1.1",
             "platform": "",
-            "vulnerabilities": [{
-                "summary": "Multiple stack-based buffer overflows in mimetex.cgi in mimeTeX",
-                "cvss": None,
-                "references": [{
-                    "reference_id": "CVE-2009-2458",
-                    "source": "",
-                    "url": "",
-                }]
-            }],
-            "references": [],
-        }, {
-            "name": "mimetex",
-            "version": "1.50-1.1",
-            "platform": "",
-            "vulnerabilities": [{
-                "summary": "Multiple unspecified vulnerabilities in mimeTeX.",
-                "cvss": None,
-                "references": [{
-                    "reference_id": "CVE-2009-2459",
-                    "source": "",
-                    "url": "",
-                }]
-            }],
+            "vulnerabilities": [
+                {
+                    "summary": "Multiple stack-based buffer overflows in mimetex.cgi in mimeTeX",
+                    "cvss": None,
+                    "references": [{
+                        "reference_id": "CVE-2009-2458",
+                        "source": "",
+                        "url": "",
+                    }]
+                },
+                {
+                    "summary": "Multiple unspecified vulnerabilities in mimeTeX.",
+                    "cvss": None,
+                    "references": [{
+                        "reference_id": "CVE-2009-2459",
+                        "source": "",
+                        "url": "",
+                    }]
+                }
+
+            ],
             "references": [],
         }]
 
@@ -86,13 +83,12 @@ def test_ubuntu_response(self):
         extract_data = ubuntu.extract_cves(test_data)
         ubuntu_dump(extract_data)
         response = self.client.get('/api/packages/?name=automake', format='json')
-
         expected = [{
             "name": "automake",
             "version": "",
             "platform": "",
             "vulnerabilities": [{
-                "summary": "",
+                "summary": "CVE-2012-3386 vulnerability",
                 "cvss": None,
                 "references": [{
                     "reference_id": "CVE-2012-3386",
@@ -121,30 +117,26 @@ def test_serializers(self):
                 "name": "mimetex",
                 "version": "1.50-1.1",
                 "platform": "",
-                "vulnerabilities": [{
-                    "summary": "Multiple stack-based buffer overflows in mimetex.cgi in mimeTeX",
-                    "cvss": None,
-                    "references": [{
-                        "reference_id": "CVE-2009-2458",
-                        "source": "",
-                        "url": "",
-                    }]
-                }],
-                "references": [],
-            },
-            {
-                "name": "mimetex",
-                "version": "1.50-1.1",
-                "platform": "",
-                "vulnerabilities": [{
-                    "summary": "Multiple unspecified vulnerabilities in mimeTeX.",
-                    "cvss": None,
-                    "references": [{
-                        "reference_id": "CVE-2009-2459",
-                        "source": "",
-                        "url": "",
-                    }]
-                }],
+                "vulnerabilities": [
+                    {
+                        "summary": "Multiple stack-based buffer overflows in mimetex.cgi in mimeTeX",
+                        "cvss": None,
+                        "references": [{
+                            "reference_id": "CVE-2009-2458",
+                            "source": "",
+                            "url": "",
+                        }]
+                    },
+                    {
+                        "summary": "Multiple unspecified vulnerabilities in mimeTeX.",
+                        "cvss": None,
+                        "references": [{
+                            "reference_id": "CVE-2009-2459",
+                            "source": "",
+                            "url": "",
+                        }]
+                    }
+                ],
                 "references": [],
             }
         ]
diff --git a/vulnerabilities/tests/test_data_dump.py b/vulnerabilities/tests/test_data_dump.py
@@ -35,7 +35,6 @@
 from vulnerabilities.models import ResolvedPackage
 from vulnerabilities.models import Vulnerability
 from vulnerabilities.models import VulnerabilityReference
-from vulnerabilities.scraper import archlinux
 from vulnerabilities.scraper import debian
 from vulnerabilities.scraper import ubuntu
 
@@ -58,7 +57,7 @@ def test_debian_data_dump(self):
 
         self.assertEqual(3, Vulnerability.objects.count())
         self.assertEqual(3, VulnerabilityReference.objects.count())
-        self.assertEqual(3, Package.objects.count())
+        self.assertEqual(2, Package.objects.count())
 
         self.assertTrue(Vulnerability.objects.get(
                         summary='Multiple stack-based buffer overflows in mimetex.cgi in mimeTeX'))
@@ -76,6 +75,22 @@ def test_debian_data_dump(self):
         self.assertTrue(Package.objects.get(name='git-repair'))
         self.assertEqual(Package.objects.filter(version='1.50-1.1')[0].version, '1.50-1.1')
 
+    def test_debian_data_dump_twice(self):
+        """
+        Scrape data from Debian' main tracker, save it
+        in the database and verify entries.
+        """
+        with open(os.path.join(TEST_DATA, 'debian.json')) as f:
+            test_data = json.loads(f.read())
+
+        extract_data = debian.extract_vulnerabilities(test_data)
+        debian_dump(extract_data)
+        debian_dump(extract_data)
+
+        self.assertEqual(3, Vulnerability.objects.count())
+        self.assertEqual(3, VulnerabilityReference.objects.count())
+        self.assertEqual(2, Package.objects.count())
+
     def test_ubuntu_data_dump(self):
         """
         Scrape data from Ubuntu' main tracker, save it
@@ -91,6 +106,23 @@ def test_ubuntu_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.
+        """
+        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)
+        count = Package.objects.all().count()
+        ubuntu_dump(data)
+        self.assertEqual(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')
+
     def test_archlinux_data_dump(self):
         """
         Scrape data from Archlinux' main tracker, save it
@@ -116,4 +148,21 @@ def test_archlinux_data_dump(self):
 
         self.assertTrue(VulnerabilityReference.objects.get(reference_id='AVG-708'))
 
-        self.assertEqual(Package.objects.filter(name='wireshark-cli')[0].name, 'wireshark-cli')
+        self.assertEqual(Package.objects.filter(name='wireshark-cli')[0].name, 'wireshark-cli')
+
+    def test_archlinux_data_dump_twice(self):
+        """
+        Scrape data from Archlinux' main tracker twice, save it
+        in the database and verify no multiple entries.
+        """
+        with open(os.path.join(TEST_DATA, 'archlinux.json')) as f:
+            test_data = json.loads(f.read())
+
+        archlinux_dump(test_data)
+        archlinux_dump(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())
