Break load_toml_from_md into smaller methods

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

Author: sbs2001

vulnerabilities/importers/rust.py

@@ -83,13 +83,13 @@ def _load_advisories(self, files) -> Set[Advisory]:
     def collect_packages(self, paths):
         packages = set()
         for path in paths:
-            record = load_toml_from_md(path)
+            record = get_advisory_data(path)
             packages.add(record["advisory"]["package"])
 
         return packages
 
     def _load_advisory(self, path: str) -> Optional[Advisory]:
-        record = load_toml_from_md(path)
+        record = get_advisory_data(path)
         advisory = record.get("advisory", {})
         crate_name = advisory["package"]
         references = []
@@ -174,7 +174,54 @@ def categorize_versions(
     return unaffected, affected
 
 
-def load_toml_from_md(md_path):
+def get_toml_lines(lines):
+    """
+    Yield lines of TOML extracted from an iterable of text ``lines``.
+    The lines are expected to be from a RustSec Markdown advisory file with
+    embedded TOML metadata.
+
+    For example::
+
+    >>> text = '''
+    ... ```toml
+    ... [advisory]
+    ... id = "RUST-001"
+    ...
+    ... [versions]
+    ... patch = [">= 1.2.1"]
+    ... ```
+    ... # Use-after-free with objects returned by `Stream`'s `get_format_info`
+    ...
+    ... Affected versions contained a pair of use-after-free issues with the objects.
+    ... '''
+    >>> list(get_toml_lines(text.splitlines()))
+    ['', '[advisory]', 'id = "RUST-001"', '', '[versions]', 'patch = [">= 1.2.1"]']
+    """
+
+    for line in lines:
+        line = line.strip()
+        if line.startswith("```toml"):
+            continue
+        elif line.endswith("```"):
+            break
+        else:
+            yield line
+
+
+def data_from_toml_lines(lines):
+    """
+    Return a mapping of data from an iterable of TOML text ``lines``.
+
+    For example::
+
+    >>> lines = ['[advisory]', 'id = "RUST1"', '', '[versions]', 'patch = [">= 1"]']
+    >>> data_from_toml_lines(lines)
+    {'advisory': {'id': 'RUST1'}, 'versions': {'patch': ['>= 1']}}
+    """
+    return toml.loads("\n".join(lines))
+
+
+def get_advisory_data(location):
     """
     Return a mapping of vulnerability data from a RustSec advisory file at
     ``location``.
@@ -183,11 +230,7 @@ def load_toml_from_md(md_path):
     https://github.com/RustSec/advisory-db#advisory-format:
         Advisories are formatted in Markdown with TOML "front matter".
     """
-    parsed_data = {}
-    with open(md_path) as f:
-        result = re.findall("^```toml(.*?)^```", f.read(), re.DOTALL | re.MULTILINE)
-        # This regex matches everything that's between ```toml and ```
-        for i in result:
-            parsed_toml = toml.loads("".join(i))
-            parsed_data.update(parsed_toml)
-    return parsed_data
+
+    with open(location) as lines:
+        toml_lines = get_toml_lines(lines)
+        return data_from_toml_lines(toml_lines)

vulnerabilities/tests/test_rust.py

@@ -34,7 +34,7 @@
 from vulnerabilities.importers.rust import categorize_versions
 from vulnerabilities.import_runner import ImportRunner
 from vulnerabilities.package_managers import VersionAPI
-from vulnerabilities.importers.rust import load_toml_from_md
+from vulnerabilities.importers.rust import get_advisory_data
 
 BASE_DIR = os.path.dirname(os.path.abspath(__file__))
 TEST_DATA = os.path.join(BASE_DIR, "test_data/rust")
@@ -187,7 +187,7 @@ def test_import(self, _):
     def test_load_toml_from_md(self, _):
         md_path = os.path.join(TEST_DATA, "CVE-2019-16760.md")
         # print(list(os.walk(self.tempdir)))
-        loaded_data = load_toml_from_md(md_path)
+        loaded_data = get_advisory_data(md_path)
         expected_data = {
             "advisory": {
                 "aliases": ["GHSA-phjm-8x66-qw4r"],
@@ -196,8 +196,6 @@ def test_load_toml_from_md(self, _):
                 "package": "cargo",
                 "url": "https://groups.google.com/forum/#!topic/rustlang-security-announcements/rVQ5e3TDnpQ",  # nopep8
             },
-            "serde": {"features": ["derive"], "version": "1.0"},
-            "serde1": {"features": ["derive"], "version": "1.0"},
             "versions": {"patched": [">= 1.26.0"]},
         }