Add tests for etags

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

Author: sbs2001

vulnerabilities/data_source.py

@@ -393,10 +393,15 @@ def _include_file(
 
 
 class OvalDataSource(DataSource):
-
+    """
+    All data sources which collect data from OVAL files must inherit from this (OvalDataSource) class.
+    Subclasses must implement the methods `_fetch` and `set_api`.  
+    """
     @staticmethod
-    def create_purl(pkg_name: str, pkg_version: str, pkg_data: Mapping):
+    def create_purl(pkg_name: str, pkg_version: str, pkg_data: Mapping) -> PackageURL:
         """
+        Helper method for creating different purls for subclasses without them reimplementing 
+        get_data_from_xml_doc  method
         Note: pkg_data must include 'type' of package
         """
         return PackageURL(name=pkg_name, version=pkg_version, **pkg_data)
@@ -405,7 +410,7 @@ def create_purl(pkg_name: str, pkg_version: str, pkg_data: Mapping):
     def _collect_pkgs(parsed_oval_data: Mapping) -> Set:
         """
         Helper method, used for loading the API. It expects data from
-        OvalParser.get_data() .
+        OvalParser.get_data().
         """
         all_pkgs = set()
         for definition_data in parsed_oval_data:
@@ -419,47 +424,66 @@ def _fetch() -> Tuple[Mapping, Iterable[ET.ElementTree]]:
         """
         This method  contains logic to fetch OVAL files and yield them into
         a tuple of file's metadata and it's ET.ElementTree.
-        Subclasses must implement this method.
+        Subclasses must implement this method. 
+
+        Note: Mapping MUST INCLUDE "type" key. Example values of Mapping
+              {"type":"deb","qualifiers":{"distro":"buster"} }
+
         """
         raise NotImplementedError
 
     def added_advisories(self) -> List[Advisory]:
+        """
+        Note: metadata MUST INCLUDE "type" key, implement _fetch accordingly.
+        """
         advisories = []
         for metadata, oval_file in self._fetch():
             advisories.extend(self.get_data_from_xml_doc(oval_file, metadata))
-        return advisories
+        return self.batch_advisories(advisories)
 
     def set_api(self, all_pkgs: Iterable[str]):
         """
         This method loads the self.pkg_manager_api with the specified packages. It fetches
-        and caches the data about these packages exposes them through
-        self.pkg_manager_api.get(<package_name>)
+        and caches all the versions of these packages and exposes them through
+        self.pkg_manager_api.get(<package_name>). Example
+
+        >>> self.set_api(['electron'])
+        Assume 'electron' has only versions 1.0.0 and 1.2.0
+        >>> assert  self.pkg_manager_api.get('electron') == {'1.0.0','1.2.0'}
+
         """
         raise NotImplementedError
 
     def get_data_from_xml_doc(self, xml_doc: ET.ElementTree, pkg_metadata={}) -> List[Advisory]:
         """
-        The orchestration method of the OvalDataSource. Breaks an OVAL xml
-        ElementTree into a list of Advisory.
+        The orchestration method of the OvalDataSource. This method breaks an OVAL xml
+        ElementTree into a list of `Advisory`.
+
+        Note: pkg_metadata MUST INCLUDE "type" key. Example value of pkg_metadata, 
+              {"type":"deb","qualifiers":{"distro":"buster"} }
         """
         all_adv = []
         oval_doc = OvalParser(self.translations, xml_doc)
         raw_data = oval_doc.get_data()
         all_pkgs = self._collect_pkgs(raw_data)
         self.set_api(all_pkgs)
         for definition_data in raw_data:  # definition_data -> Advisory
+
+            # These fields are definition level, i.e common for all
+            # elements connected/linked to an OvalDefinition
             vuln_id = definition_data['vuln_id']
             description = definition_data['description']
             affected_purls = set()
             safe_purls = set()
             urls = definition_data['reference_urls']
+
             for test_data in definition_data['test_data']:
                 for package in test_data['package_list']:
                     pkg_name = package
                     aff_ver_range = test_data['version_ranges']
                     all_versions = self.pkg_manager_api.get(package)
-                    # This filter is to filter out long versions.
-                    # 50 is limit because that's what db permits atm
+                    # This filter is for filtering out long versions.
+                    # 50 is limit because that's what db permits atm.
                     all_versions = set(
                         filter(
                             lambda x: len(x) < 50,
@@ -489,4 +513,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 self.batch_advisories(all_adv)
+        return all_adv

vulnerabilities/tests/test_data_source.py

@@ -27,16 +27,29 @@
 from unittest import TestCase
 from unittest.mock import MagicMock
 from unittest.mock import patch
+import xml.etree.ElementTree as ET
 
 import pygit2
 import pytest
+from packageurl import PackageURL
 
-from vulnerabilities.data_source import GitDataSource, _include_file
+from vulnerabilities.data_source import GitDataSource, _include_file, OvalDataSource
 from vulnerabilities.data_source import InvalidConfigurationError
+from vulnerabilities.oval_parser import OvalParser
 
 BASE_DIR = os.path.dirname(os.path.abspath(__file__))
 TEST_DATA = os.path.join(BASE_DIR, 'test_data/')
 
+def load_oval_data():
+    etrees_of_oval = {}
+    for f in os.listdir(TEST_DATA):
+        if f.endswith('oval_data.xml'):
+            path = os.path.join(TEST_DATA, f)
+            provider = f.split("_")[0]
+            etrees_of_oval[provider] = ET.parse(path)
+    return etrees_of_oval
+
+
 
 def mk_ds(**kwargs):
     # just for convenience, since this is a manadory parameter we always pass a value
@@ -331,3 +344,46 @@ def test_file_changes_include_fixed_advisories(self):
         assert len(added_files) == 0
         assert len(updated_files) == 1
         assert os.path.join(self.repodir, 'crates/hyper/RUSTSEC-2020-0008.toml') in updated_files
+
+class TestOvalDataSource(TestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        cls.oval_data_src = OvalDataSource(1)
+
+    def test_create_purl(self):
+        purl1 =  PackageURL(name="ffmpeg",type="test",version="1.2.0")
+
+        assert purl1 == self.oval_data_src.create_purl(pkg_name="ffmpeg",
+        pkg_version="1.2.0", pkg_data={"type":"test"})
+
+        purl2 = PackageURL(name="notepad",type="example",version="7.9.6",namespace="ns",
+        qualifiers={"distro":"sample"},subpath="root")
+        assert purl2 == self.oval_data_src.create_purl(pkg_name="notepad",
+        pkg_version="7.9.6",pkg_data={
+        "namespace":"ns","qualifiers":{"distro":"sample"},
+        "subpath":"root","type":"example"
+        }
+        )
+    
+    def test__collect_pkgs(self):
+
+        xmls = load_oval_data()
+
+        expected_suse_pkgs = {'cacti-spine', 'apache2-mod_perl', 'cacti', 'apache2-mod_perl-devel'}
+        expected_ubuntu_pkgs = {'potrace', 'tor'}
+
+        translations = {"less than": "<"}
+
+        found_suse_pkgs = self.oval_data_src._collect_pkgs(
+            OvalParser(translations,xmls['suse']).get_data())
+
+        found_ubuntu_pkgs = self.oval_data_src._collect_pkgs(
+            OvalParser(translations,xmls['ubuntu']).get_data())
+
+        assert found_suse_pkgs == expected_suse_pkgs
+        assert found_ubuntu_pkgs == expected_ubuntu_pkgs
+
+
+
+

vulnerabilities/tests/test_ubuntu.py

@@ -1,6 +1,7 @@
 import os
 import unittest
 from unittest.mock import patch
+from unittest.mock import MagicMock
 import xml.etree.ElementTree as ET
 from collections import OrderedDict
 import asyncio
@@ -15,6 +16,9 @@
 BASE_DIR = os.path.dirname(os.path.abspath(__file__))
 TEST_DATA = os.path.join(BASE_DIR, "test_data/")
 
+class MockResponse:
+
+    headers = {"ETag":"0x1234"}
 
 class TestUbuntuOvalParser(unittest.TestCase):
     @classmethod
@@ -165,7 +169,7 @@ def test_get_data(self):
 
         assert expected_data == self.parsed_oval.get_data()
 
-#This is horrible, there must be a better way
+#This is horrible, there might be a better way
 async def mock(a,b):
     pass
 
@@ -176,7 +180,10 @@ class TestUbuntuDataSource(unittest.TestCase):
 
     @classmethod
     def setUpClass(cls):
-        pass
+        data_source_cfg = {
+        'releases': 'eg-ubuntu',"etags":{}}
+        cls.ubuntu_data_src = UbuntuDataSource(
+            batch_size=1, config=data_source_cfg)
 
     @patch(
         'vulnerabilities.importers.ubuntu.VersionAPI.get',
@@ -186,11 +193,6 @@ def setUpClass(cls):
             '2.14-2'})
     @patch('vulnerabilities.importers.ubuntu.VersionAPI.load_api',new=mock)
     def test_get_data_from_xml_doc(self, mock_write):
-
-        data_source_cfg = {
-            'releases': 'eg-ubuntu',"etags":{}}
-        ubuntu_data_src = UbuntuDataSource(
-            batch_size=1, config=data_source_cfg)
         expected_data = {
             Advisory(
                 summary=('Tor before 0.2.8.9 and 0.2.9.x before 0.2.9.4-alpha had '
@@ -268,8 +270,19 @@ 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"))
-        # Dirty quick patch to deal with batch_advisories
+        # Dirty quick patch to mock 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"})}
+            data = {i for i in self.ubuntu_data_src.get_data_from_xml_doc(xml_doc,{"type":"deb"})}
         assert expected_data == data
+
+    def test_create_etag(self):
+        
+        assert self.ubuntu_data_src.config.etags == {}
+        with patch('vulnerabilities.importers.ubuntu.requests.head', return_value=MockResponse()):
+            assert True == self.ubuntu_data_src.create_etag("https://example.org")
+            assert self.ubuntu_data_src.config.etags == {"https://example.org":"0x1234"}
+            assert False == self.ubuntu_data_src.create_etag("https://example.org")
+        
+
+