Add OvalDataSource, refactor UbuntuDataSource to inherit from OvalDataSource

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

Author: sbs2001

vulnerabilities/data_source.py

@@ -35,10 +35,13 @@
 from typing import Sequence
 from typing import Set
 from typing import Tuple
+import xml.etree.ElementTree as ET
 
 import pygit2
 from packageurl import PackageURL
 
+from vulnerabilities.oval_parser import OvalParser
+
 
 @dataclasses.dataclass
 class Advisory:
@@ -258,7 +261,8 @@ def file_changes(
 
             return {str(p) for p in path.glob(glob) if p.is_file()}, set()
 
-        return self._collect_file_changes(subdir=subdir, recursive=recursive, file_ext=file_ext)
+        return self._collect_file_changes(
+            subdir=subdir, recursive=recursive, file_ext=file_ext)
 
     def _collect_file_changes(
             self,
@@ -270,7 +274,8 @@ def _collect_file_changes(
         previous_commit = None
         added_files, updated_files = set(), set()
 
-        for commit in self._repo.walk(self._repo.head.target, pygit2.GIT_SORT_TIME):
+        for commit in self._repo.walk(
+                self._repo.head.target, pygit2.GIT_SORT_TIME):
             commit_time = commit.commit_time + commit.commit_time_offset  # convert to UTC
 
             if commit_time < self.cutoff_timestamp:
@@ -281,13 +286,16 @@ def _collect_file_changes(
                 continue
 
             for d in commit.tree.diff_to_tree(previous_commit.tree).deltas:
-                if not _include_file(d.new_file.path, subdir, recursive, file_ext) or d.is_binary:
+                if not _include_file(
+                        d.new_file.path, subdir, recursive, file_ext) or d.is_binary:
                     continue
 
-                abspath = os.path.join(self.config.working_directory, d.new_file.path)
+                abspath = os.path.join(
+                    self.config.working_directory, d.new_file.path)
                 # TODO
                 # Just filtering on the two status values for "added" and "modified" is too
-                # simplistic. This does not cover file renames, copies & deletions.
+                # simplistic. This does not cover file renames, copies &
+                # deletions.
                 if d.status == pygit2.GIT_DELTA_ADDED:
                     added_files.add(abspath)
                 elif d.status == pygit2.GIT_DELTA_MODIFIED:
@@ -382,3 +390,103 @@ def _include_file(
         match = match and path.endswith(f'.{file_ext}')
 
     return match
+
+
+class OvalDataSource(DataSource):
+
+    @staticmethod
+    def create_purl(pkg_name: str, pkg_version: str, pkg_data: Mapping):
+        """
+        Note: pkg_data must include 'type' of package
+        """
+        return PackageURL(name=pkg_name, version=pkg_version, **pkg_data)
+
+    @staticmethod
+    def _collect_pkgs(parsed_oval_data: Mapping) -> Set:
+        """
+        Helper method, used for loading the API. It expects data from
+        OvalParser.get_data() .
+        """
+        all_pkgs = set()
+        for definition_data in parsed_oval_data:
+            for test_data in definition_data['test_data']:
+                for package in test_data['package_list']:
+                    all_pkgs.add(package)
+
+        return all_pkgs
+
+    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.
+        """
+        raise NotImplementedError
+
+    def added_advisories(self) -> List[Advisory]:
+        advisories = []
+        for metadata, oval_file in self._fetch():
+            advisories.extend(self.get_data_from_xml_doc(oval_file, metadata))
+        return 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>)
+        """
+        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.
+        """
+        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
+            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
+                    all_versions = set(
+                        filter(
+                            lambda x: len(x) < 50,
+                            all_versions))
+                    if not all_versions:
+                        continue
+                    affected_versions = set(
+                        filter(
+                            lambda x: x in aff_ver_range,
+                            all_versions))
+                    safe_versions = all_versions - affected_versions
+
+                    for version in affected_versions:
+                        pkg_url = self.create_purl(
+                            pkg_name=pkg_name, pkg_version=version, pkg_data=pkg_metadata)
+                        affected_purls.add(pkg_url)
+
+                    for version in safe_versions:
+                        pkg_url = self.create_purl(
+                            pkg_name=pkg_name, pkg_version=version, pkg_data=pkg_metadata)
+                        safe_purls.add(pkg_url)
+
+            all_adv.append(
+                Advisory(
+                    summary=description,
+                    impacted_package_urls=affected_purls,
+                    resolved_package_urls=safe_purls,
+                    cve_id=vuln_id,
+                    reference_urls=urls))
+        return all_adv

vulnerabilities/importers/__init__.py

@@ -20,6 +20,7 @@
 #  VulnerableCode is a free software code scanning tool from nexB Inc. and others.
 #  Visit https://github.com/nexB/vulnerablecode/ for support and download.
 
+
 from vulnerabilities.importers.alpine_linux import AlpineDataSource
 from vulnerabilities.importers.archlinux import ArchlinuxDataSource
 from vulnerabilities.importers.debian import DebianDataSource

vulnerabilities/importers/ubuntu.py

@@ -31,22 +31,21 @@
 import xml.etree.ElementTree as ET
 
 
-from aiohttp import ClientSession
+from aiohttp import ClientSession, ClientTimeout
 from aiohttp.client_exceptions import ClientResponseError
 import requests
 from packageurl import PackageURL
 
 
-from vulnerabilities.data_source import DataSource, DataSourceConfiguration, Advisory
-from vulnerabilities.importers import oval_parser
+from vulnerabilities.data_source import OvalDataSource, DataSourceConfiguration, Advisory
 
 
 @dataclasses.dataclass
 class UbuntuConfiguration(DataSourceConfiguration):
     releases: list
 
 
-class UbuntuDataSource(DataSource):
+class UbuntuDataSource(OvalDataSource):
 
     CONFIG_CLASS = UbuntuConfiguration
 
@@ -56,86 +55,21 @@ def __init__(self, *args, **kwargs):
         # set by default in the OvalParser, but we don't yet know
         # whether all OVAL providers use the same format
         self.translations = {'less than': '<'}
-        self._versions = VersionAPI()
+        self.pkg_manager_api = VersionAPI()
 
     def _fetch(self):
         base_url = 'https://people.canonical.com/~ubuntu-security/oval/'
         file_name = 'com.ubuntu.{}.cve.oval.xml.bz2'
         releases = self.config.releases
         for release in releases:
+            print("getting ", release)
             resp = requests.get(base_url + file_name.format(release))
             extracted = bz2.decompress(resp.content)
-            yield ET.ElementTree(ET.fromstring(extracted.decode('utf-8')))
-
-    def added_advisories(self) -> List[Advisory]:
-        advisories = []
-        for oval_file in self._fetch():
-            advisories.extend(self.get_data_from_xml_doc(oval_file))
-        return advisories
-
-    @staticmethod
-    def _collect_pkgs(parsed_oval_data) -> Set:
-        all_pkgs = set()
-        for definition_data in parsed_oval_data:
-            for test_data in definition_data['test_data']:
-                for package in test_data['package_list']:
-                    all_pkgs.add(package)
-
-        return all_pkgs
-
-    def get_data_from_xml_doc(self, xml_doc) -> List[Advisory]:
-        all_adv = []
-        oval_doc = oval_parser.OvalParser(self.translations, xml_doc)
-        raw_data = oval_doc.get_data()
-        all_pkgs = self._collect_pkgs(raw_data)
-
-        asyncio.run(self._versions.load_api(all_pkgs))
-
-        for definition_data in raw_data:  # definition_data -> Advisory
-            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._versions.get(package)
-                    # This filter is to filter out long versions.
-                    # 50 is limit because that's what db permits atm
-                    all_versions = set(
-                        filter(
-                            lambda x: len(x) < 50,
-                            all_versions))
-                    if not all_versions:
-                        continue
-                    affected_versions = set(
-                        filter(
-                            lambda x: x in aff_ver_range,
-                            all_versions))
-                    safe_versions = all_versions - affected_versions
-
-                    for version in affected_versions:
-                        # should we add a qualifier like 'distro:ubuntu'?
-                        pkg_url = PackageURL(
-                            name=pkg_name, type='deb', version=version)
-                        affected_purls.add(pkg_url)
-
-                    for version in safe_versions:
-                        # should we add a qualifier like 'distro:ubuntu'?
-                        pkg_url = PackageURL(
-                            name=pkg_name, type='deb', version=version)
-                        safe_purls.add(pkg_url)
-
-            all_adv.append(
-                Advisory(
-                    summary=description,
-                    impacted_package_urls=affected_purls,
-                    resolved_package_urls=safe_purls,
-                    cve_id=vuln_id,
-                    reference_urls=urls))
-        return all_adv
+            print("done ")
+            yield ({'type': 'deb'}, ET.ElementTree(ET.fromstring(extracted.decode('utf-8'))))
+
+    def set_api(self, packages):
+        asyncio.run(self.pkg_manager_api.load_api(packages))
 
 
 class VersionAPI:
@@ -146,7 +80,9 @@ def get(self, package_name: str) -> Set[str]:
         return self.cache[package_name]
 
     async def load_api(self, pkg_set):
-        async with ClientSession() as session:
+        # This is debatable
+        timeout = ClientTimeout(total=None)
+        async with ClientSession(raise_for_status=True, timeout=timeout) as session:
             await asyncio.gather(*[self.set_api(pkg, session)
                                    for pkg in pkg_set if pkg not in self.cache])
 
@@ -156,9 +92,8 @@ async def set_api(self, pkg, session):
                'source_name={}&exact_match=true'.format(pkg))
         try:
             all_versions = set()
-            while(True):
+            while True:
                 response = await session.request(method='GET', url=url)
-                response.raise_for_status()
                 resp_json = await response.json()
                 if resp_json['entries'] == []:
                     self.cache[pkg] = {}

vulnerabilities/importers/lib_oval.py vulnerabilities/lib_oval.pyvulnerabilities/importers/lib_oval.py renamed to vulnerabilities/lib_oval.py

@@ -30,7 +30,7 @@
 
 from dephell_specifier import RangeSpecifier
 
-from vulnerabilities.importers.lib_oval import (
+from vulnerabilities.lib_oval import (
     OvalDefinition, OvalDocument, OvalTest, OvalObject, OvalState)
 
 

vulnerabilities/importers/oval_parser.py vulnerabilities/oval_parser.pyvulnerabilities/importers/oval_parser.py renamed to vulnerabilities/oval_parser.py

@@ -5,7 +5,7 @@
 from dephell_specifier import RangeSpecifier
 
 
-from vulnerabilities.importers.oval_parser import OvalParser
+from vulnerabilities.oval_parser import OvalParser
 
 
 BASE_DIR = os.path.dirname(os.path.abspath(__file__))

vulnerabilities/tests/test_suse.py

@@ -8,7 +8,7 @@
 from dephell_specifier import RangeSpecifier
 from packageurl import PackageURL
 
-from vulnerabilities.importers.oval_parser import OvalParser
+from vulnerabilities.oval_parser import OvalParser
 from vulnerabilities.importers.ubuntu import UbuntuDataSource
 from vulnerabilities.data_source import Advisory
 
@@ -265,6 +265,6 @@ 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))
+        data = set(ubuntu_data_src.get_data_from_xml_doc(xml_doc,{"type":"deb"}))
 
         assert expected_data == data