Cleanup OvalDataSource and UbuntuDataSource, add etags to ubuntu migration script

Signed-off-by: Shivam Sandbhor <shivam.sandbhor@gmail.com>

Author: sbs2001

vulnerabilities/data_source.py

@@ -489,4 +489,4 @@ def get_data_from_xml_doc(self, xml_doc: ET.ElementTree, pkg_metadata={}) -> Lis
                     resolved_package_urls=safe_purls,
                     cve_id=vuln_id,
                     reference_urls=urls))
-        return all_adv
+        return self.batch_advisories(all_adv)

vulnerabilities/importers/ubuntu.py

@@ -31,13 +31,12 @@
 import xml.etree.ElementTree as ET
 
 
-from aiohttp import ClientSession, ClientTimeout
+from aiohttp import ClientSession
 from aiohttp.client_exceptions import ClientResponseError
 import requests
-from packageurl import PackageURL
 
 
-from vulnerabilities.data_source import OvalDataSource, DataSourceConfiguration, Advisory
+from vulnerabilities.data_source import OvalDataSource, DataSourceConfiguration
 
 
 @dataclasses.dataclass
@@ -69,9 +68,14 @@ def _fetch(self):
             resp = requests.get(file_url)
             extracted = bz2.decompress(resp.content)
             yield (
-                    {'type': 'deb', 'namespace': 'ubuntu'},
-                    ET.ElementTree(ET.fromstring(extracted.decode('utf-8')))
-                  )
+                {'type': 'deb', 'namespace': 'ubuntu'},
+                ET.ElementTree(ET.fromstring(extracted.decode('utf-8')))
+            )
+        # In case every file is latest, _fetch won't yield anything(due to checking for new etags),
+        # this would return None to added_advisories
+        # which will cause error, hence this
+        # function return an empty list
+        return []
 
     def set_api(self, packages):
         asyncio.run(self.pkg_manager_api.load_api(packages))

vulnerabilities/migrations/0009_ubuntu_importer.py

@@ -31,7 +31,8 @@ def add_ubuntu_importer(apps, _):
         license='',
         last_run=None,
         data_source='UbuntuDataSource',
-        data_source_cfg={'releases':['bionic','trusty','focal','eoan','xenial']},
+        data_source_cfg={'releases':['bionic','trusty','focal','eoan','xenial'],
+        'etags':{}},
     )
 
 

vulnerabilities/tests/test_ubuntu.py

@@ -169,6 +169,9 @@ def test_get_data(self):
 async def mock(a,b):
     pass
 
+def return_adv(_,a):
+    return a
+
 class TestUbuntuDataSource(unittest.TestCase):
 
     @classmethod
@@ -185,7 +188,7 @@ def setUpClass(cls):
     def test_get_data_from_xml_doc(self, mock_write):
 
         data_source_cfg = {
-            'releases': 'eg-ubuntu'}
+            'releases': 'eg-ubuntu',"etags":{}}
         ubuntu_data_src = UbuntuDataSource(
             batch_size=1, config=data_source_cfg)
         expected_data = {
@@ -265,6 +268,8 @@ def test_get_data_from_xml_doc(self, mock_write):
                 cve_id='CVE-2016-8703')}
 
         xml_doc = ET.parse(os.path.join(TEST_DATA, "ubuntu_oval_data.xml"))
-        data = set(ubuntu_data_src.get_data_from_xml_doc(xml_doc,{"type":"deb"}))
-
+        # Dirty quick patch to deal with batch_advisories
+        with patch('vulnerabilities.importers.ubuntu.UbuntuDataSource.batch_advisories',
+         new=return_adv):
+            data = {i for i in ubuntu_data_src.get_data_from_xml_doc(xml_doc,{"type":"deb"})}
         assert expected_data == data