Merge pull request #34 from Grunny/context-and-user-scans

Add context commands and the ability to scan as a user

README.rst

@@ -62,6 +62,7 @@ ZAP CLI can then be used with the following commands:
       active-scan  Run an Active Scan.
       ajax-spider  Run the AJAX Spider against a URL.
       alerts       Show alerts at the given alert level.
+      context      Manage contexts for the current session.
       exclude      Exclude a pattern from all scanners.
       open-url     Open a URL using the ZAP proxy.
       policies     Enable or list a set of policies.
@@ -162,3 +163,35 @@ Or to run the same scan with the API key disabled:
 ::
 
     $ zap-cli quick-scan -sc -o '-config api.disablekey=true' -s xss http://127.0.0.1/
+
+Running scans as authenticated users
+------------------------------------
+In order to run a scan as an authenticated user, first configure the authentication method and users for
+a context using the ZAP UI (see the `ZAP help page <https://github.com/zaproxy/zap-core-help/wiki/HelpStartConceptsAuthentication>`_
+for more information). Once the authentication method and users are prepared, you can then export the context
+with the configured authentication method so it can be imported and used to run authenticated scans with ZAP CLI.
+
+You can export a context with the authentication method and users configured either through the ZAP UI or using the
+``context export`` ZAP CLI command. For example, to export a context with the name DevTest to a file, you could run:
+
+::
+
+    $ zap-cli context export --name DevTest --file-path /home/user/DevTest.context
+
+To import the saved context for use with ZAP CLI later, you could run:
+
+::
+
+    $ zap-cli context import /home/user/DevTest.context
+
+After importing the context with the configured authentication method and users, you can then provide the context name
+and user name to the ``spider``, ``active-scan``, and ``quick-scan`` commands to run the scans while authenticated as
+the given user. For example:
+
+::
+
+    $ zap-cli context import /home/user/DevTest.context
+    $ zap-cli open-url "http://localhost/"
+    $ zap-cli spider --context-name DevTest --user-name SomeUser "http://localhost"
+    $ zap-cli active-scan --recursive -c DevTest -u SomeUser "http://localhost"
+    $ zap-cli quick-scan --recursive --spider -c DevTest -u SomeUser "http://localhost"

setup.py

@@ -35,7 +35,7 @@
     ],
     install_requires=[
         'click==4.0',
-        'python-owasp-zap-v2.4==0.0.11',
+        'python-owasp-zap-v2.4==0.0.12',
         'requests==2.13.0',
         'tabulate==0.7.5',
         'termcolor==1.1.0',

tests/cli_test.py

@@ -110,7 +110,7 @@ def test_open_url_no_url(self, helper_mock):
     def test_spider_url(self, helper_mock):
         """Test spider URL method."""
         result = self.runner.invoke(cli.cli, ['--boring', '--api-key', '', 'spider', 'http://localhost/'])
-        helper_mock.assert_called_with('http://localhost/')
+        helper_mock.assert_called_with('http://localhost/', None, None)
         self.assertEqual(result.exit_code, 0)
 
     @patch('zapcli.zap_helper.ZAPHelper.run_spider')
@@ -311,6 +311,32 @@ def test_html_report(self, report_mock):
         report_mock.assert_called_with('foo.html')
         self.assertEqual(result.exit_code, 0)
 
+    @patch('zapcli.zap_helper.ZAPHelper.include_in_context')
+    def test_context_include(self, helper_mock):
+        """Testing including a regex in a given context."""
+        result = self.runner.invoke(cli.cli, ['--boring', '--api-key', '', '--verbose', 'context',
+                                              'include', '--name', 'Test', '--pattern', 'zap-cli'])
+        self.assertEqual(result.exit_code, 0)
+
+    def test_context_include_error(self):
+        """Testing that an error is reported when providing an invalid regex."""
+        result = self.runner.invoke(cli.cli, ['--boring', '--api-key', '', '--verbose', 'context',
+                                              'include', '--name', 'Test', '--pattern', '['])
+        self.assertEqual(result.exit_code, 1)
+
+    @patch('zapcli.zap_helper.ZAPHelper.exclude_from_context')
+    def test_context_exclude(self, helper_mock):
+        """Testing excluding a regex from a given context."""
+        result = self.runner.invoke(cli.cli, ['--boring', '--api-key', '', '--verbose', 'context',
+                                              'exclude', '--name', 'Test', '--pattern', 'zap-cli'])
+        self.assertEqual(result.exit_code, 0)
+
+    def test_context_exclude_error(self):
+        """Testing that an error is reported when providing an invalid regex."""
+        result = self.runner.invoke(cli.cli, ['--boring', '--api-key', '', '--verbose', 'context',
+                                              'exclude', '--name', 'Test', '--pattern', '['])
+        self.assertEqual(result.exit_code, 1)
+
 
 if __name__ == '__main__':
     unittest.main()

tests/zap_helper_test.py

@@ -27,6 +27,7 @@ class ZAPHelperTestCase(unittest.TestCase):
 
     def setUp(self):
         self.zap_helper = zap_helper.ZAPHelper(api_key=self.api_key)
+        self.zap_helper._status_check_sleep = 0
 
     @data(
         (
@@ -218,7 +219,7 @@ def status_result():
         class_mock.status = status
         self.zap_helper.zap.spider = class_mock
 
-        self.zap_helper.run_spider('http://localhost', status_check_sleep=0)
+        self.zap_helper.run_spider('http://localhost')
 
     def test_run_spider_error(self):
         """Test running the spider when an error occurs."""
@@ -227,7 +228,44 @@ def test_run_spider_error(self):
         self.zap_helper.zap.spider = class_mock
 
         with self.assertRaises(ZAPError):
-            self.zap_helper.run_spider('http://localhost', status_check_sleep=0)
+            self.zap_helper.run_spider('http://localhost')
+
+    def test_run_spider_as_user(self):
+        """Test running the spider as a given user."""
+        def status_result():
+            """Return value of the status property."""
+            if status.call_count > 2:
+                return '100'
+            return '50'
+
+        class_mock = MagicMock()
+        class_mock.scan_as_user.return_value = '1'
+        status = Mock(side_effect=status_result)
+        class_mock.status = status
+        self.zap_helper.zap.spider = class_mock
+        self.zap_helper.zap.context.context = Mock(return_value={'id': '1'})
+        self.zap_helper.zap.users.users_list = Mock(return_value=[{'name': 'Test', 'id': '1'}])
+
+        self.zap_helper.run_spider('http://localhost', 'Test', 'Test')
+
+    def test_run_spider_as_user_error(self):
+        """Test running the spider as a given user when an error occurs."""
+        def status_result():
+            """Return value of the status property."""
+            if status.call_count > 2:
+                return '100'
+            return '50'
+
+        class_mock = MagicMock()
+        class_mock.scan_as_user.return_value = '1'
+        status = Mock(side_effect=status_result)
+        class_mock.status = status
+        self.zap_helper.zap.spider = class_mock
+        self.zap_helper.zap.context.context = Mock(return_value={'id': '1'})
+        self.zap_helper.zap.users.users_list = Mock(return_value=[])
+
+        with self.assertRaises(ZAPError):
+            self.zap_helper.run_spider('http://localhost', 'Test', 'Test')
 
     def test_run_active_scan(self):
         """Test running an active scan."""
@@ -243,7 +281,25 @@ def status_result():
         class_mock.status = status
         self.zap_helper.zap.ascan = class_mock
 
-        self.zap_helper.run_active_scan('http://localhost', status_check_sleep=0)
+        self.zap_helper.run_active_scan('http://localhost')
+
+    def test_run_active_scan_as_user(self):
+        """Test running an active scan as a given user."""
+        def status_result():
+            """Return value of the status property."""
+            if status.call_count > 2:
+                return '100'
+            return '50'
+
+        class_mock = MagicMock()
+        class_mock.scan_as_user.return_value = '1'
+        status = Mock(side_effect=status_result)
+        class_mock.status = status
+        self.zap_helper.zap.ascan = class_mock
+        self.zap_helper.zap.context.context = Mock(return_value={'id': '1'})
+        self.zap_helper.zap.users.users_list = Mock(return_value=[{'name': 'Test', 'id': '1'}])
+
+        self.zap_helper.run_active_scan('http://localhost', False, 'Test', 'Test')
 
     def test_run_active_scan_error(self):
         """Test running an active scan."""
@@ -276,7 +332,7 @@ def status_result():
         type(class_mock).status = status
         self.zap_helper.zap.ajaxSpider = class_mock
 
-        self.zap_helper.run_ajax_spider('http://localhost', status_check_sleep=0)
+        self.zap_helper.run_ajax_spider('http://localhost')
 
     @data(
         (
@@ -658,6 +714,110 @@ def test_html_report(self, htmlreport_mock):
             report_str = report_str.encode('utf-8')
         file_open_mock().write.assert_called_with(report_str)
 
+    @patch('zapv2.context.new_context')
+    def test_new_context(self, context_mock):
+        """Test creating a new context."""
+        self.zap_helper.new_context('Test')
+        context_mock.assert_called_with(contextname='Test', apikey=self.api_key)
+
+    @patch('zapv2.context.include_in_context')
+    def test_include_in_context(self, context_mock):
+        """Test adding a regex for URLs to include in the context."""
+        include_pattern = r"\/zapcli.+"
+        context_name = 'Test'
+        context_mock.return_value = 'OK'
+
+        self.zap_helper.include_in_context(context_name, include_pattern)
+
+        context_mock.assert_called_with(contextname=context_name, regex=include_pattern, apikey=self.api_key)
+
+    @patch('zapv2.context.include_in_context')
+    def test_include_in_context_return_error(self, context_mock):
+        """Test that an error is raised when the API returns an unexpected response."""
+        include_pattern = r"\/zapcli.+"
+        context_name = 'Test'
+        context_mock.return_value = 'Error'
+
+        with self.assertRaises(ZAPError):
+            self.zap_helper.include_in_context(context_name, include_pattern)
+
+    def test_include_in_context_regex_error(self):
+        """Test that an error is raised when an invalid regex is supplied."""
+        include_pattern = '['
+        context_name = 'Test'
+
+        with self.assertRaises(ZAPError):
+            self.zap_helper.include_in_context(context_name, include_pattern)
+
+    @patch('zapv2.context.exclude_from_context')
+    def test_exclude_from_context(self, context_mock):
+        """Test adding a regex for URLs to exclude from the context."""
+        exclude_pattern = r"\/zapcli.+"
+        context_name = 'Test'
+        context_mock.return_value = 'OK'
+
+        self.zap_helper.exclude_from_context(context_name, exclude_pattern)
+
+        context_mock.assert_called_with(contextname=context_name, regex=exclude_pattern, apikey=self.api_key)
+
+    @patch('zapv2.context.exclude_from_context')
+    def test_exclude_from_context_return_error(self, context_mock):
+        """Test that an error is raised when the API returns an unexpected response."""
+        exclude_pattern = r"\/zapcli.+"
+        context_name = 'Test'
+        context_mock.return_value = 'Error'
+
+        with self.assertRaises(ZAPError):
+            self.zap_helper.exclude_from_context(context_name, exclude_pattern)
+
+    def test_exclude_from_context_regex_error(self):
+        """Test that an error is raised when an invalid regex is supplied."""
+        exclude_pattern = '['
+        context_name = 'Test'
+
+        with self.assertRaises(ZAPError):
+            self.zap_helper.exclude_from_context(context_name, exclude_pattern)
+
+    @patch('zapv2.context.import_context')
+    def test_import_context(self, context_mock):
+        """Test importing a context."""
+        file_path = '/tmp/Test'
+        context_mock.return_value = '1'
+
+        self.zap_helper.import_context(file_path)
+
+        context_mock.assert_called_with(file_path, apikey=self.api_key)
+
+    @patch('zapv2.context.import_context')
+    def test_import_context_return_error(self, context_mock):
+        """Test that an error is raised when the API returns an unexpected response."""
+        file_path = '/tmp/Test'
+        context_mock.return_value = 'Error'
+
+        with self.assertRaises(ZAPError):
+            self.zap_helper.import_context(file_path)
+
+    @patch('zapv2.context.export_context')
+    def test_export_context(self, context_mock):
+        """Test exporting a context."""
+        file_path = '/tmp/Test'
+        context_name = 'Test'
+        context_mock.return_value = 'OK'
+
+        self.zap_helper.export_context(context_name, file_path)
+
+        context_mock.assert_called_with(context_name, file_path, apikey=self.api_key)
+
+    @patch('zapv2.context.export_context')
+    def test_export_context_return_error(self, context_mock):
+        """Test that an error is raised when the API returns an unexpected response."""
+        file_path = '/tmp/Test'
+        context_name = 'Test'
+        context_mock.return_value = 'Error'
+
+        with self.assertRaises(ZAPError):
+            self.zap_helper.export_context(context_name, file_path)
+
 
 if __name__ == '__main__':
     unittest.main()

zapcli/cli.py

@@ -8,6 +8,7 @@
 
 from zapcli import __version__
 from zapcli import helpers
+from zapcli.commands.context import context_group
 from zapcli.commands.policies import policies_group
 from zapcli.commands.scanners import scanner_group
 from zapcli.commands.scripts import scripts_group
@@ -103,12 +104,16 @@ def open_url(zap_helper, url):
 
 @cli.command('spider')
 @click.argument('url')
+@click.option('--context-name', '-c', type=str, help='Context to use if provided.')
+@click.option('--user-name', '-u', type=str,
+              help='Run scan as this user if provided. If this option is used, the context parameter must also ' +
+              'be provided.')
 @click.pass_obj
-def spider_url(zap_helper, url):
+def spider_url(zap_helper, url, context_name, user_name):
     """Run the spider against a URL."""
     console.info('Running spider...')
     with helpers.zap_error_handler():
-        zap_helper.run_spider(url)
+        zap_helper.run_spider(url, context_name, user_name)
 
 
 @cli.command('ajax-spider')
@@ -127,8 +132,12 @@ def ajax_spider_url(zap_helper, url):
               'subcommand to get a list of IDs. Available groups are: {0}.'.format(
                   ', '.join(['all'] + list(ZAPHelper.scanner_group_map.keys()))))
 @click.option('--recursive', '-r', is_flag=True, default=False, help='Make scan recursive.')
+@click.option('--context-name', '-c', type=str, help='Context to use if provided.')
+@click.option('--user-name', '-u', type=str,
+              help='Run scan as this user if provided. If this option is used, the context parameter must also ' +
+              'be provided.')
 @click.pass_obj
-def active_scan(zap_helper, url, scanners, recursive):
+def active_scan(zap_helper, url, scanners, recursive, context_name, user_name):
     """
     Run an Active Scan against a URL.
 
@@ -141,7 +150,7 @@ def active_scan(zap_helper, url, scanners, recursive):
         if scanners:
             zap_helper.set_enabled_scanners(scanners)
 
-        zap_helper.run_active_scan(url, recursive=recursive)
+        zap_helper.run_active_scan(url, recursive, context_name, user_name)
 
 
 @cli.command('alerts')
@@ -183,6 +192,10 @@ def show_alerts(zap_helper, alert_level, output_format, exit_code):
               ' e.g. "-config api.key=12345"')
 @click.option('--output-format', '-f', default='table', type=click.Choice(['table', 'json']),
               help='Output format to print the alerts.')
+@click.option('--context-name', '-c', type=str, help='Context to use if provided.')
+@click.option('--user-name', '-u', type=str,
+              help='Run scan as this user if provided. If this option is used, the context parameter must also ' +
+              'be provided.')
 @click.pass_obj
 def quick_scan(zap_helper, url, **options):
     """
@@ -209,12 +222,12 @@ def quick_scan(zap_helper, url, **options):
         zap_helper.open_url(url)
 
         if options['spider']:
-            zap_helper.run_spider(url)
+            zap_helper.run_spider(url, options['context_name'], options['user_name'])
 
         if options['ajax_spider']:
             zap_helper.run_ajax_spider(url)
 
-        zap_helper.run_active_scan(url, recursive=options['recursive'])
+        zap_helper.run_active_scan(url, options['recursive'], options['context_name'], options['user_name'])
 
     alerts = zap_helper.alerts(options['alert_level'])
 
@@ -255,6 +268,7 @@ def report(zap_helper, output, output_format):
 
 
 # Add subcommand groups
+cli.add_command(context_group)
 cli.add_command(policies_group)
 cli.add_command(scanner_group)
 cli.add_command(scripts_group)