def load_script(zap_helper, **options):
"""Load a script from a file."""
with zap_error_handler():
if not os.path.isfile(options['file_path']):
raise ZAPError(
'No file found at "{0}", cannot load script.'.format(
options['file_path']))
if not _is_valid_script_engine(zap_helper.zap, options['engine']):
engines = zap_helper.zap.script.list_engines
raise ZAPError(
'Invalid script engine provided. Valid engines are: {0}'.
format(', '.join(engines)))
console.debug('Loading script "{0}" from "{1}"'.format(
options['name'], options['file_path']))
result = zap_helper.zap.script.load(
options['name'],
options['script_type'],
options['engine'],
options['file_path'],
scriptdescription=options['description'])
if result != 'OK':
raise ZAPError('Error loading script: {0}'.format(result))
console.info('Script "{0}" loaded'.format(options['name']))
def load_script(self,
name,
script_type,
engine,
file_path,
description=''):
"""Load a given script."""
if not os.path.isfile(file_path):
raise ZAPError(
'No file found at "{0}", cannot load script.'.format(
file_path))
if not self.is_valid_script_engine(engine):
engines = self.zap.script.list_engines
raise ZAPError(
'Invalid script engine provided. Valid engines are: {0}'.
format(', '.join(engines)))
self.logger.debug('Loading script "{0}" from "{1}"'.format(
name, file_path))
result = self.zap.script.load(name,
script_type,
engine,
file_path,
scriptdescription=description,
apikey=self.api_key)
if result != 'OK':
raise ZAPError('Error loading script: {0}'.format(result))
def run_active_scan(self, target_url, recursive=False, context_name=None, user_name=None):
"""Run an active scan against a URL."""
self.logger.debug('Scanning target {0}...'.format(target_url))
context_id, user_id = self._get_context_and_user_ids(context_name, user_name)
if user_id:
self.logger.debug('Scanning in context {0} as user {1}'.format(context_id, user_id))
scan_id = self.zap.ascan.scan_as_user(target_url, context_id, user_id, recursive, apikey=self.api_key)
else:
scan_id = self.zap.ascan.scan(target_url, recurse=recursive, apikey=self.api_key)
if not scan_id:
raise ZAPError('Error running active scan.')
elif not scan_id.isdigit():
raise ZAPError(('Error running active scan: "{0}". Make sure the URL is in the site ' +
'tree by using the open-url or scanner commands before running an active ' +
'scan.').format(scan_id))
self.logger.debug('Started scan with ID {0}...'.format(scan_id))
while int(self.zap.ascan.status()) < 100:
self.logger.debug('Scan progress %: {0}'.format(self.zap.ascan.status()))
time.sleep(self._status_check_sleep)
self.logger.debug('Scan #{0} completed'.format(scan_id))
def run_spider(self, target_url, context_name=None, user_name=None):
"""Run spider against a URL."""
self.logger.debug('Spidering target {0}...'.format(target_url))
context_id, user_id = self._get_context_and_user_ids(
context_name, user_name)
if user_id:
self.logger.debug(
'Running spider in context {0} as user {1}'.format(
context_id, user_id))
scan_id = self.zap.spider.scan_as_user(context_id,
user_id,
target_url,
apikey=self.api_key)
else:
scan_id = self.zap.spider.scan(target_url, apikey=self.api_key)
if not scan_id:
raise ZAPError('Error running spider.')
elif not scan_id.isdigit():
raise ZAPError('Error running spider: "{0}"'.format(scan_id))
self.logger.debug('Started spider with ID {0}...'.format(scan_id))
while int(self.zap.spider.status()) < 100:
self.logger.debug('Spider progress %: {0}'.format(
self.zap.spider.status()))
time.sleep(self._status_check_sleep)
self.logger.debug('Spider #{0} completed'.format(scan_id))
def exclude_from_context(self, context_name, regex):
"""Add exclude regex to context."""
try:
re.compile(regex)
except re.error:
raise ZAPError('Invalid regex "{0}" provided'.format(regex))
result = self.zap.context.exclude_from_context(contextname=context_name, regex=regex, apikey=self.api_key)
if result != 'OK':
raise ZAPError('Excluding regex from context failed: {}'.format(result))
def test_quick_scan_start_error(self, helper_mock):
"""Testing quick scan."""
helper_mock.side_effect = ZAPError('error')
result = self.runner.invoke(cli.cli, ['--boring', '--api-key', '', '--verbose', 'quick-scan',
'http://localhost/', '--self-contained'])
self.assertEqual(result.exit_code, 2)
def remove_script(self, script_name):
"""Remove a given script."""
self.logger.debug('Removing script "{0}"'.format(script_name))
result = self.zap.script.remove(script_name, apikey=self.api_key)
if result != 'OK':
raise ZAPError('Error removing script: {0}'.format(result))
def test_exclude_from_scanners_error(self, helper_mock):
"""Test exclude from scanners command with error raised."""
helper_mock.side_effect = ZAPError('error')
result = self.runner.invoke(
cli.cli, ['--boring', '--api-key', '', 'exclude', '['])
helper_mock.assert_called_with('[')
self.assertEqual(result.exit_code, 1)
def wait_for_zap(self, timeout):
"""Wait for ZAP to be ready to receive API calls."""
timeout_time = time.time() + timeout
while not self.is_running():
if time.time() > timeout_time:
raise ZAPError('Timed out waiting for ZAP to start.')
time.sleep(2)
def test_shutdown_zap_daemon_exception(self, helper_mock):
"""Test command to shutdown ZAP daemon has an exit code of 1 when an exception is raised."""
helper_mock.side_effect = ZAPError('error')
result = self.runner.invoke(cli.cli,
['--boring', '--api-key', '', 'shutdown'])
helper_mock.assert_called_with()
self.assertEqual(result.exit_code, 1)
def test_check_status_timeout(self, running_mock, wait_mock):
"""Test the status command with a timeout."""
running_mock.return_value = False
wait_mock.side_effect = ZAPError('error')
result = self.runner.invoke(
cli.cli, ['--boring', '--api-key', '', 'status', '-t', '0'])
self.assertEqual(result.exit_code, 1)
def test_load_script_error(self, helper_mock):
"""Test command to load a script with error raised."""
helper_mock.side_effect = ZAPError('error')
result = self.runner.invoke(cli.cli, ['--boring', '--api-key', '', 'scripts', 'load',
'--name', 'Foo.js', '--script-type', 'proxy',
'--engine', 'Oracle Nashorn', '--file-path', 'Foo.js'])
self.assertEqual(result.exit_code, 1)
def get_context_info(self, context_name):
"""Get the context ID for a given context name."""
context_info = self.zap.context.context(context_name)
if not isinstance(context_info, dict):
raise ZAPError('Context with name "{0}" wasn\'t found'.format(context_name))
return context_info
def disable_script(self, script_name):
"""Disable a given script."""
self.logger.debug('Disabling script "{0}"'.format(script_name))
result = self.zap.script.disable(script_name, apikey=self.api_key)
if result != 'OK':
raise ZAPError('Error disabling script: {0}'.format(result))
def start(self, options=None):
"""Start the ZAP Daemon."""
if self.is_running():
self.logger.warn('ZAP is already running on port {0}'.format(self.port))
return
if platform.system() == 'Windows' or platform.system().startswith('CYGWIN'):
executable = 'zap.bat'
else:
executable = 'zap.sh'
executable_path = os.path.join(self.zap_path, executable)
if not os.path.isfile(executable_path):
raise ZAPError(('ZAP was not found in the path "{0}". You can set the path to where ZAP is ' +
'installed on your system using the --zap-path command line parameter or by ' +
'default using the ZAP_PATH environment variable.').format(self.zap_path))
zap_command = [executable_path, '-daemon', '-port', str(self.port)]
if options:
extra_options = shlex.split(options)
zap_command += extra_options
log_path = os.path.join(self.zap_path, 'zap.log')
self.logger.debug('Starting ZAP process with command: {0}.'.format(' '.join(zap_command)))
self.logger.debug('Logging to {0}'.format(log_path))
with open(log_path, 'w+') as log_file:
subprocess.Popen(
zap_command, cwd=self.zap_path, stdout=log_file,
stderr=subprocess.STDOUT)
self.wait_for_zap(self.timeout)
self.logger.debug('ZAP started successfully.')
def set_scanner_attack_strength(self, scanner_ids, attack_strength):
"""Set the attack strength for the given scanners."""
for scanner_id in scanner_ids:
self.logger.debug('Setting strength for scanner {0} to {1}'.format(scanner_id, attack_strength))
result = self.zap.ascan.set_scanner_attack_strength(scanner_id, attack_strength,
apikey=self.api_key)
if result != 'OK':
raise ZAPError('Error setting strength for scanner with ID {0}: {1}'.format(scanner_id, result))
def _get_user_id_from_name(self, context_id, user_name):
"""Get a user ID from the user name."""
users = self.zap.users.users_list(context_id)
for user in users:
if user['name'] == user_name:
return user['id']
raise ZAPError('No user with the name "{0}"" was found for context {1}'.format(user_name, context_id))
def set_policy_alert_threshold(self, policy_ids, alert_threshold):
"""Set the alert theshold for the given policies."""
for policy_id in policy_ids:
self.logger.debug('Setting alert threshold for policy {0} to {1}'.format(policy_id, alert_threshold))
result = self.zap.ascan.set_policy_alert_threshold(policy_id, alert_threshold,
apikey=self.api_key)
if result != 'OK':
raise ZAPError('Error setting alert threshold for policy with ID {0}: {1}'.format(policy_id, result))
def set_policy_attack_strength(self, policy_ids, attack_strength):
"""Set the attack strength for the given policies."""
for policy_id in policy_ids:
self.logger.debug('Setting strength for policy {0} to {1}'.format(policy_id, attack_strength))
result = self.zap.ascan.set_policy_attack_strength(policy_id, attack_strength,
apikey=self.api_key)
if result != 'OK':
raise ZAPError('Error setting strength for policy with ID {0}: {1}'.format(policy_id, result))
def test_remove_script_error(self, helper_mock):
"""Test command to remove a script with error raised."""
helper_mock.side_effect = ZAPError('error')
result = self.runner.invoke(
cli.cli,
['--boring', '--api-key', '', 'scripts', 'remove', 'Foo.js'])
helper_mock.assert_called_with('Foo.js')
self.assertEqual(result.exit_code, 1)
def load_session(self, file_path):
"""Load a given session."""
if not os.path.isfile(file_path):
raise ZAPError(
'No file found at "{0}", cannot load session.'.format(
file_path))
self.logger.debug('Loading session from "{0}"'.format(file_path))
self.zap.core.load_session(file_path, apikey=self.api_key)
def import_context(self, file_path):
"""Import a context from a file."""
result = self.zap.context.import_context(file_path,
apikey=self.api_key)
if not result.isdigit():
raise ZAPError(
'Importing context from file failed: {}'.format(result))
def run_spider(self, target_url, status_check_sleep=10):
"""Run spider against a URL."""
self.logger.debug('Spidering target {0}...'.format(target_url))
scan_id = self.zap.spider.scan(target_url, apikey=self.api_key)
if not scan_id:
raise ZAPError('Error running spider.')
elif not scan_id.isdigit():
raise ZAPError('Error running spider: "{0}"'.format(scan_id))
self.logger.debug('Started spider with ID {0}...'.format(scan_id))
while int(self.zap.spider.status()) < 100:
self.logger.debug('Spider progress %: {0}'.format(self.zap.spider.status()))
time.sleep(status_check_sleep)
self.logger.debug('Spider #{0} completed'.format(scan_id))
def test_quick_scan_exclude_from_all_error(self, helper_mock):
"""Testing quick scan."""
instance = helper_mock.return_value
instance.alerts.return_value = []
instance.exclude_from_all.side_effect = ZAPError('error')
result = self.runner.invoke(cli.cli, ['--boring', '--api-key', '', '--verbose', 'quick-scan',
'http://localhost/', '--exclude', 'pattern'])
self.assertEqual(result.exit_code, 2)
def load_session(zap_helper, file_path):
"""Load a given session."""
with zap_error_handler():
if not os.path.isfile(file_path):
raise ZAPError(
'No file found at "{0}", cannot load session.'.format(
file_path))
console.debug('Loading session from "{0}"'.format(file_path))
zap_helper.zap.core.load_session(file_path)
def export_context(self, context_name, file_path):
"""Export a given context to a file."""
result = self.zap.context.export_context(context_name,
file_path,
apikey=self.api_key)
if result != 'OK':
raise ZAPError(
'Exporting context to file failed: {}'.format(result))
def test_quick_scan_shutdown_error(self, helper_mock):
"""Testing quick scan."""
instance = helper_mock.return_value
instance.alerts.return_value = []
instance.shutdown.side_effect = ZAPError('error')
result = self.runner.invoke(cli.cli, ['--boring', '--api-key', '', '--verbose', 'quick-scan',
'http://localhost/', '--self-contained'])
self.assertEqual(result.exit_code, 1)
def disable_script(zap_helper, script_name):
"""Disable a script."""
with zap_error_handler():
console.debug('Disabling script "{0}"'.format(script_name))
result = zap_helper.zap.script.disable(script_name)
if result != 'OK':
raise ZAPError('Error disabling script: {0}'.format(result))
console.info('Script "{0}" disabled'.format(script_name))
def context_export(zap_helper, name, file_path):
"""Export a given context to a file."""
with zap_error_handler():
result = zap_helper.zap.context.export_context(name, file_path)
if result != 'OK':
raise ZAPError(
'Exporting context to file failed: {}'.format(result))
console.info('Exported context {0} to {1}'.format(name, file_path))
def remove_script(zap_helper, script_name):
"""Remove a script."""
with zap_error_handler():
console.debug('Removing script "{0}"'.format(script_name))
result = zap_helper.zap.script.remove(script_name)
if result != 'OK':
raise ZAPError('Error removing script: {0}'.format(result))
console.info('Script "{0}" removed'.format(script_name))