Add test for ruby importer

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

Author: sbs2001

vulnerabilities/data_dump.py

@@ -195,17 +195,10 @@ def npm_dump(extract_data):
 def ruby_dump(extract_data):
     for package_data in extract_data:
 
-        if(package_data['severity']):
-            vulnerability, _ = Vulnerability.objects.get_or_create(
-                cve_id=package_data['cve_id'],
-                summary=package_data['summary'],
-                cvss=package_data['severity']
-            )
-        else:
-            vulnerability, _ = Vulnerability.objects.get_or_create(
-                cve_id=package_data['cve_id'],
-                summary=package_data['summary']
-            )
+        vulnerability, _ = Vulnerability.objects.get_or_create(
+            cve_id=package_data['cve_id'],
+            summary=package_data['summary'],
+        )
 
         VulnerabilityReference.objects.get_or_create(
             vulnerability=vulnerability,
@@ -215,6 +208,7 @@ def ruby_dump(extract_data):
         for version in package_data['affected_versions']:
             affected_package = Package.objects.create(
                 name=package_data['package_name'],
+                type='gem',
                 version=version
             )
             ImpactedPackage.objects.create(
@@ -225,6 +219,7 @@ def ruby_dump(extract_data):
         for version in package_data['fixed_versions']:
             unaffected_package = Package.objects.create(
                 name=package_data['package_name'],
+                type='gem',
                 version=version
             )
             ResolvedPackage.objects.create(

vulnerabilities/scraper/ruby.py

@@ -1,39 +1,21 @@
 import os
 import urllib.request
-import saneyaml
-
-from shutil import rmtree
 from urllib.error import HTTPError
 from zipfile import ZipFile
-from itertools import chain
+from io import BytesIO
+import saneyaml
 from dephell_specifier import RangeSpecifier
+from urllib.request import urlopen
 
-RUBYCVE_LINK = 'https://github.com/rubysec/ruby-advisory-db/archive/master.zip'
-DOWNLOAD_PATH = os.path.dirname(os.path.realpath(__file__))
-
+RUBYSEC_DB_URL = 'https://github.com/rubysec/ruby-advisory-db/archive/master.zip'
 
-def get_rubycve_db():
-    pathToZip, _ = urllib.request.urlretrieve(
-        RUBYCVE_LINK, os.path.join(
-            DOWNLOAD_PATH, 'ruby.zip'))
-    ZipFile(pathToZip).extractall(DOWNLOAD_PATH)
-    os.remove(pathToZip)
 
-
-def path_of_yaml_of_all_packages():
-    gemPath = os.path.join(DOWNLOAD_PATH, 'ruby-advisory-db-master', 'gems')
-    rubiesPath = os.path.join(
-        DOWNLOAD_PATH,
-        'ruby-advisory-db-master',
-        'rubies')
-    for (
-            packagePath,
-            _,
-            yamlNames) in chain(
-            os.walk(gemPath),
-            os.walk(rubiesPath)):
-        for yamlName in yamlNames:
-            yield os.path.join(packagePath, yamlName)
+def rubygem_advisories(url, prefix='ruby-advisory-db-master/gems/'):
+    with urlopen(url) as response:
+        with ZipFile(BytesIO(response.read())) as zf:
+            for path in zf.namelist():
+                if path.startswith(prefix) and path.endswith('.yml'):
+                    yield saneyaml.load(zf.open(path))
 
 
 def get_all_versions_of_package(package_name):
@@ -48,57 +30,55 @@ def get_all_versions_of_package(package_name):
 
 
 def get_patched_range(spec_list):
-    if spec_list:
-        def remove_space(string): return string.replace(' ', '')
-        spec_list = list(map(remove_space, spec_list))
-        for spec in spec_list:
-            if 'rc' in spec:
-                continue
-            yield RangeSpecifier(spec)
-    else:
-        return None
+    spec_list = [string.replace(' ', '') for string in spec_list]
+    for spec in spec_list:
+        if 'rc' in spec:
+            continue
+        yield RangeSpecifier(spec)
 
 
 def import_vulnerabilities():
-    get_rubycve_db()
     ids = set()
-    vulnerability_to_package_map = []
-    for vulnerability_path in path_of_yaml_of_all_packages():
-        with open(vulnerability_path) as yamlFile:
-            vulnerability = saneyaml.load(yamlFile)
-            package_name = vulnerability.get(
-                'engine', vulnerability.get('gem'))
-            summary = vulnerability.get('description', '')
-            if 'cve' in vulnerability:
-                vulnerability_id = 'CVE-{}'.format(vulnerability['cve'])
-                if vulnerability_id in ids:
-                    continue
-                ids.add(vulnerability_id)
-            else:
+    vulnerability_package_dicts = []
+    for vulnerability in rubygem_advisories(RUBYSEC_DB_URL):
+
+        package_name = vulnerability.get(
+            'gem')
+
+        if not package_name:
+            continue
+
+        summary = vulnerability.get('description', '')
+
+        if 'cve' in vulnerability:
+            vulnerability_id = 'CVE-{}'.format(vulnerability['cve'])
+            if vulnerability_id in ids:
                 continue
-            severity = vulnerability.get(
-                'cvss_v3', vulnerability.get('cvss_v2'))
-            advisoryUrl = vulnerability.get('url')
-            specs = list(
-                get_patched_range(
-                    vulnerability.get('patched_versions')))
-            allVersions = set(list(get_all_versions_of_package(package_name)))
-            unaffected_versions = set()
-            if specs:
-                for version in allVersions:
-                    for spec in specs:
-                        if(version in spec):
-                            unaffected_versions.add(version)
-                            break
-            affected_versions = allVersions - unaffected_versions
-            vulnerability_to_package_map.append({
-                'package_name': package_name,
-                'summary': summary,
-                'cve_id': vulnerability_id,
-                'fixed_versions': unaffected_versions,
-                'affected_versions': affected_versions,
-                'severity': severity,
-                'advisory': advisoryUrl
-            })
-    rmtree(os.path.join(DOWNLOAD_PATH, 'ruby-advisory-db-master'))
-    return vulnerability_to_package_map
+            ids.add(vulnerability_id)
+        else:
+            continue
+
+        advisory_url = vulnerability.get('url')
+        patched_version_ranges = list(
+            get_patched_range(
+                vulnerability.get('patched_versions', [])))
+        all_versions = set(get_all_versions_of_package(package_name))
+        unaffected_versions = set()
+
+        if patched_version_ranges:
+            for version in all_versions:
+                for spec in patched_version_ranges:
+                    if version in spec:
+                        unaffected_versions.add(version)
+                        break
+
+        affected_versions = all_versions - unaffected_versions
+        vulnerability_package_dicts.append({
+            'package_name': package_name,
+            'summary': summary,
+            'cve_id': vulnerability_id,
+            'fixed_versions': unaffected_versions,
+            'affected_versions': affected_versions,
+            'advisory': advisory_url
+        })
+    return vulnerability_package_dicts

vulnerabilities/tests/test_ruby.py

@@ -0,0 +1,16 @@
+from dephell_specifier import RangeSpecifier
+
+from vulnerabilities.scraper.ruby import get_all_versions_of_package
+from vulnerabilities.scraper.ruby import get_patched_range
+
+
+def test_get_all_versions_of_package():
+    versions = set(get_all_versions_of_package('actionpack'))
+    expected = set(['1.10.2', '4.0.3', '3.2.3', '4.2.0.beta3',
+                    '5.2.3', '4.1.14.2', '4.0.13'])
+    assert versions >= expected
+
+
+def test_get_patched_range():
+    expected = set([RangeSpecifier('>4'), RangeSpecifier('~>4.1')])
+    assert set(get_patched_range(['>4', '~>4.1', '>4.2rc'])) == expected