WorkbenchHelper method to export workbench reports

Resolves #30

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
 ---------------

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
 ---------------

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:

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
 ----------------------
 

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)

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):
 

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)