Merge pull request #57 from haikoschol/import-cli

Provide CLI for data import

Author: pombredanne

README.md

@@ -65,24 +65,7 @@ DJANGO_DEV=1 python manage.py test vulnerabilities/tests
 ## Data import
 
 ```
-DJANGO_DEV=1 python manage.py shell
-```
-
-```
-from vulnerabilities.scraper import archlinux, debian, ubuntu
-from vulnerabilities.data_dump import archlinux_dump, debian_dump, ubuntu_dump
-
-# May be needed on macOS
-# import ssl; ssl._create_default_https_context = ssl._create_unverified_context
-
-ubuntu_cves = ubuntu.scrape_cves()
-ubuntu_dump(ubuntu_cves)
-
-debian_vulnerabilities = debian.scrape_vulnerabilities()
-debian_dump(debian_vulnerabilities)
-
-archlinux_vulnerabilities = archlinux.scrape_vulnerabilities()
-archlinux_dump(archlinux_vulnerabilities)
+DJANGO_DEV=1 python manage.py import --all
 ```
 
 ## API

vulnerabilities/management/commands/import.py

@@ -0,0 +1,78 @@
+#
+# Copyright (c) 2017 nexB Inc. and others. All rights reserved.
+# http://nexb.com and https://github.com/nexB/vulnerablecode/
+# The VulnerableCode software is licensed under the Apache License version 2.0.
+# Data generated with VulnerableCode require an acknowledgment.
+#
+# You may not use this software except in compliance with the License.
+# You may obtain a copy of the License at: http://apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software distributed
+# under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
+# CONDITIONS OF ANY KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations under the License.
+#
+# When you publish or redistribute any data created with VulnerableCode or any VulnerableCode
+# derivative work, you must accompany this data with the following acknowledgment:
+#
+#  Generated with VulnerableCode and provided on an "AS IS" BASIS, WITHOUT WARRANTIES
+#  OR CONDITIONS OF ANY KIND, either express or implied. No content created from
+#  VulnerableCode should be considered or used as legal advice. Consult an Attorney
+#  for any legal advice.
+#  VulnerableCode is a free software code scanning tool from nexB Inc. and others.
+#  Visit https://github.com/nexB/vulnerablecode/ for support and download.
+
+from django.core.management.base import BaseCommand, CommandError
+
+from vulnerabilities import data_dump as dd
+from vulnerabilities.scraper import debian, ubuntu, archlinux
+
+IMPORTERS = {
+    'debian': lambda: dd.debian_dump(debian.scrape_vulnerabilities()),
+    'ubuntu': lambda: dd.ubuntu_dump(ubuntu.scrape_cves()),
+    'archlinux': lambda: dd.archlinux_dump(archlinux.scrape_vulnerabilities()),
+}
+
+
+class Command(BaseCommand):
+    help = 'Import vulnerability data'
+
+    def add_arguments(self, parser):
+        parser.add_argument('--list', action='store_true', help='List available data sources')
+
+        parser.add_argument('--all', action='store_true',
+                            help='Import data from all available sources')
+
+        parser.add_argument('sources', nargs='*',
+                            help='Data sources from which to import')
+
+    def handle(self, *args, **options):
+        if options['list']:
+            self.list_sources()
+            return
+
+        if options['all']:
+            self.import_data(IMPORTERS.keys())
+            return
+
+        sources = options['sources']
+        if not sources:
+            raise CommandError(
+                'Please provide at least one data source to import from or use "--all".')
+
+        self.validate_sources(sources)
+        self.import_data(sources)
+
+    def validate_sources(self, sources):
+        unknown = ', '.join([s for s in sources if s not in IMPORTERS.keys()])
+        if unknown:
+            raise CommandError(f'Unknown data sources: {unknown}')
+
+    def list_sources(self):
+        self.stdout.write('Vulnerability data can be imported from the following sources:')
+        self.stdout.write(', '.join(IMPORTERS.keys()))
+
+    def import_data(self, sources):
+        for src in sources:
+            self.stdout.write(f'Importing data from {src}')
+            IMPORTERS[src]()
+            self.stdout.write(self.style.SUCCESS(f'Successfully imported data from {src}'))

vulnerabilities/scraper/__init__.py

@@ -0,0 +1,56 @@
+#
+# Copyright (c) 2017 nexB Inc. and others. All rights reserved.
+# http://nexb.com and https://github.com/nexB/vulnerablecode/
+# The VulnerableCode software is licensed under the Apache License version 2.0.
+# Data generated with VulnerableCode require an acknowledgment.
+#
+# You may not use this software except in compliance with the License.
+# You may obtain a copy of the License at: http://apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software distributed
+# under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
+# CONDITIONS OF ANY KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations under the License.
+#
+# When you publish or redistribute any data created with VulnerableCode or any VulnerableCode
+# derivative work, you must accompany this data with the following acknowledgment:
+#
+#  Generated with VulnerableCode and provided on an "AS IS" BASIS, WITHOUT WARRANTIES
+#  OR CONDITIONS OF ANY KIND, either express or implied. No content created from
+#  VulnerableCode should be considered or used as legal advice. Consult an Attorney
+#  for any legal advice.
+#  VulnerableCode is a free software code scanning tool from nexB Inc. and others.
+#  Visit https://github.com/nexB/vulnerablecode/ for support and download.
+
+from io import StringIO
+
+from django.core.management import call_command
+from django.core.management.base import CommandError
+from django.test import TestCase
+
+
+class ImportCommandTest(TestCase):
+    def test_list_sources(self):
+        buf = StringIO()
+        
+        call_command('import', '--list', stdout=buf)
+        
+        out = buf.getvalue()
+        self.assertIn('debian', out)
+        self.assertIn('ubuntu', out)
+        self.assertIn('archlinux', out)
+
+    def test_missing_sources(self):
+        with self.assertRaises(CommandError) as cm: 
+           call_command('import', stdout=StringIO())
+
+        err = str(cm.exception)
+        self.assertIn('Please provide at least one data source', err)
+
+    def test_unknown_sources(self):
+        with self.assertRaises(CommandError) as cm:
+            call_command('import', 'debian', 'foo', 'bar', stdout=StringIO())
+
+        err = str(cm.exception)
+        self.assertIn('bar', err)
+        self.assertIn('foo', err)
+        self.assertNotIn('debian', err)