From 013164dd1aa80ee59272c565c207abe622f42460 Mon Sep 17 00:00:00 2001 From: Alan Fung Date: Fri, 26 May 2017 13:58:06 -0700 Subject: [PATCH] WorkbenchHelper method to export workbench reports Resolves #30 --- doc/source/tenable_io.api.rst | 24 ++++++++++++ doc/source/tenable_io.helpers.rst | 8 ++++ doc/source/tenable_io.parser.rst | 22 +++++++++++ doc/source/tenable_io.rst | 9 +++++ examples/workbench.py | 60 +++++++++++++++++++----------- tenable_io/helpers/workbench.py | 35 +++++++++++++++++ tests/integration/test_examples.py | 4 +- 7 files changed, 139 insertions(+), 23 deletions(-) create mode 100644 doc/source/tenable_io.parser.rst diff --git a/doc/source/tenable_io.api.rst b/doc/source/tenable_io.api.rst index 04e2605..e14fcad 100644 --- a/doc/source/tenable_io.api.rst +++ b/doc/source/tenable_io.api.rst @@ -4,6 +4,22 @@ tenable_io.api package Submodules ---------- +tenable_io.api.agent_groups module +---------------------------------- + +.. automodule:: tenable_io.api.agent_groups + :members: + :undoc-members: + :show-inheritance: + +tenable_io.api.agents module +---------------------------- + +.. automodule:: tenable_io.api.agents + :members: + :undoc-members: + :show-inheritance: + tenable_io.api.base module -------------------------- @@ -124,6 +140,14 @@ tenable_io.api.users module :undoc-members: :show-inheritance: +tenable_io.api.workbenches module +--------------------------------- + +.. automodule:: tenable_io.api.workbenches + :members: + :undoc-members: + :show-inheritance: + Module contents --------------- diff --git a/doc/source/tenable_io.helpers.rst b/doc/source/tenable_io.helpers.rst index d1bb38a..6d0732f 100644 --- a/doc/source/tenable_io.helpers.rst +++ b/doc/source/tenable_io.helpers.rst @@ -28,6 +28,14 @@ tenable_io.helpers.scan module :undoc-members: :show-inheritance: +tenable_io.helpers.workbench module +----------------------------------- + +.. automodule:: tenable_io.helpers.workbench + :members: + :undoc-members: + :show-inheritance: + Module contents --------------- diff --git a/doc/source/tenable_io.parser.rst b/doc/source/tenable_io.parser.rst new file mode 100644 index 0000000..c5c5c01 --- /dev/null +++ b/doc/source/tenable_io.parser.rst @@ -0,0 +1,22 @@ +tenable_io.parser package +========================= + +Submodules +---------- + +tenable_io.parser.workbenches module +------------------------------------ + +.. automodule:: tenable_io.parser.workbenches + :members: + :undoc-members: + :show-inheritance: + + +Module contents +--------------- + +.. automodule:: tenable_io.parser + :members: + :undoc-members: + :show-inheritance: diff --git a/doc/source/tenable_io.rst b/doc/source/tenable_io.rst index d571fe7..13a0278 100644 --- a/doc/source/tenable_io.rst +++ b/doc/source/tenable_io.rst @@ -8,6 +8,7 @@ Subpackages tenable_io.api tenable_io.helpers + tenable_io.parser Submodules ---------- @@ -36,6 +37,14 @@ tenable_io.exceptions module :undoc-members: :show-inheritance: +tenable_io.log module +--------------------- + +.. automodule:: tenable_io.log + :members: + :undoc-members: + :show-inheritance: + tenable_io.util module ---------------------- diff --git a/examples/workbench.py b/examples/workbench.py index a58e70a..4d4411a 100644 --- a/examples/workbench.py +++ b/examples/workbench.py @@ -1,7 +1,12 @@ +import os + from tenable_io.client import TenableIOClient -def example(): +def example(test_file): + + test_asset_file = test_file(u'asset.nessus') + test_vulnerability_file = test_file(u'vulnerability.nessus') date_range = 7 @@ -24,23 +29,36 @@ def example(): vulnerabilities_iter = client.workbench_helper.vulnerabilities(date_range) vulnerabilities = [v for page in vulnerabilities_iter for v in page] - if len(vulnerabilities) > 0: - ''' - Get recent assets found for a plugin. - ''' - plugin_id = vulnerabilities[0].plugin_id - vulnerability_assets_iter = client.workbench_helper.assets(date_range, plugin_id=plugin_id) - vulnerability_assets = [a for page in vulnerability_assets_iter for a in page] - - assert len(vulnerability_assets) > 0 - assert plugin_id in [v.plugin_id for v in vulnerability_assets[0].vulnerabilities] - - ''' - Get recent vulnerabilities found for an asset. - ''' - asset_id = vulnerability_assets[0].asset.host_uuid - asset_vulnerabilities_iter = client.workbench_helper.vulnerabilities(date_range, asset_id=asset_id) - asset_vulnerabilities = [v for page in asset_vulnerabilities_iter for v in page] - - assert len(asset_vulnerabilities) > 0 - assert plugin_id in [v.plugin_id for v in asset_vulnerabilities] + if len(vulnerabilities) < 1: + return + + ''' + Get recent assets found for a plugin. + ''' + plugin_id = vulnerabilities[0].plugin_id + vulnerability_assets_iter = client.workbench_helper.assets(date_range, plugin_id=plugin_id) + vulnerability_assets = [a for page in vulnerability_assets_iter for a in page] + + assert len(vulnerability_assets) > 0 + assert plugin_id in [v.plugin_id for v in vulnerability_assets[0].vulnerabilities] + + ''' + Get recent vulnerabilities found for an asset. + ''' + asset_id = vulnerability_assets[0].asset.host_uuid + asset_vulnerabilities_iter = client.workbench_helper.vulnerabilities(date_range, asset_id=asset_id) + asset_vulnerabilities = [v for page in asset_vulnerabilities_iter for v in page] + + assert len(asset_vulnerabilities) > 0 + assert plugin_id in [v.plugin_id for v in asset_vulnerabilities] + + ''' + Export reports for a given asset or vulnerability. + ''' + client.workbench_helper.export(test_asset_file, asset_id=assets[0].asset.host_uuid) + client.workbench_helper.export(test_vulnerability_file, plugin_id=vulnerabilities[0].plugin_id) + + assert os.path.isfile(test_asset_file) + assert os.path.isfile(test_vulnerability_file) + os.remove(test_asset_file) + os.remove(test_vulnerability_file) diff --git a/tenable_io/helpers/workbench.py b/tenable_io/helpers/workbench.py index bd4b1f2..53faa80 100644 --- a/tenable_io/helpers/workbench.py +++ b/tenable_io/helpers/workbench.py @@ -188,6 +188,41 @@ def vulnerabilities_parse(self, date_range=1, asset_id=None, page_size=DEFAULT_P if len(vulnerabilities) > 0: yield vulnerabilities + def export( + self, + path, + format=WorkbenchesApi.FORMAT_NESSUS, + report=WorkbenchesApi.REPORT_VULNERABILITIES, + chapter=WorkbenchesApi.CHAPTER_VULN_BY_ASSET, + file_open_mode='wb', + **kwargs + ): + """Download a workbench report. + + :param format: The file format. Default to WorkbenchesApi.FORMAT_NESSUS. + :param report: The type of workbench report. Default to WorkbenchesApi.REPORT_VULNERABILITIES. + :param chapter: Chapter to include. Default to WorkbenchesApi.CHAPTER_VULN_BY_ASSET. + :param file_open_mode: Chapter to include, WorkbenchesApi.CHAPTER_VULN_BY_ASSET. + :param **kwargs: Additional keyword arguments are the same as + :class:`tenable_io.api.workbenches.WorkbenchesApi.export_request` + :return: The same WorkbenchHelper instance. + """ + file_id = self._client.workbenches_api.export_request( + format, + report, + chapter, + **kwargs + ) + + wait_until(lambda: self._client.workbenches_api.export_status(file_id) == WorkbenchesApi.STATUS_EXPORT_READY) + + iter_content = self._client.workbenches_api.export_download(file_id) + with open(path, file_open_mode) as fd: + for chunk in iter_content: + fd.write(chunk) + + return self + class AssetVulnerabilities(object): diff --git a/tests/integration/test_examples.py b/tests/integration/test_examples.py index bf4c307..8a272c6 100644 --- a/tests/integration/test_examples.py +++ b/tests/integration/test_examples.py @@ -15,6 +15,6 @@ def test_scans(self, app): from examples.scans import example example(app.session_name, app.session_file_output) - def test_workbench(self): + def test_workbench(self, app): from examples.workbench import example - example() + example(app.session_file_output)