From 33e4cb3a4082b8ef119c52dc8ebd10b10fd71f37 Mon Sep 17 00:00:00 2001 From: Steve McGrath Date: Fri, 10 May 2024 15:09:57 -0500 Subject: [PATCH] Updated SC Files method as per #762 --- requirements.txt | 1 + setup.py | 1 + tenable/sc/files.py | 17 ++++++++++---- test-requirements.txt | 1 + tests/sc/cassettes/sc_login.yaml | 2 +- tests/sc/conftest.py | 17 +++++++++++++- tests/sc/test_files.py | 39 ++++++++++++++++++++++---------- 7 files changed, 59 insertions(+), 19 deletions(-) diff --git a/requirements.txt b/requirements.txt index bff87ffc4..f42e93c13 100644 --- a/requirements.txt +++ b/requirements.txt @@ -8,3 +8,4 @@ python-box>=4.0 defusedxml>=0.5.0 urllib3==1.26.18 typing-extensions>=3.10.0.2 +requests-toolbelt==1.0.0 diff --git a/setup.py b/setup.py index e36e75f63..57e5f22c7 100644 --- a/setup.py +++ b/setup.py @@ -48,5 +48,6 @@ 'urllib3>=1.26.18', 'typing-extensions>=4.0.1', 'dataclasses>=0.8;python_version=="3.6"', + 'requests-toolbelt>=1.0.0' ], ) diff --git a/tenable/sc/files.py b/tenable/sc/files.py index aaad0fe1c..e01b6be4a 100644 --- a/tenable/sc/files.py +++ b/tenable/sc/files.py @@ -11,9 +11,13 @@ .. autoclass:: FileAPI :members: ''' +from requests_toolbelt.multipart.encoder import MultipartEncoder from .base import SCEndpoint class FileAPI(SCEndpoint): + _path = 'file' + _box = True + def upload(self, fobj): ''' Uploads a file into SecurityCenter and returns the file identifier @@ -29,8 +33,11 @@ def upload(self, fobj): The filename identifier to use for subsequent calls in Tenable Security Center. ''' - return self._api.post('file/upload', files={ - 'Filedata': fobj}).json()['response']['filename'] + encoder = MultipartEncoder({'Filedata': ('filedata', fobj)}) + return self._post('upload', + data=encoder, + headers={'Content-Type': encoder.content_type} + ).response.filename def clear(self, filename): ''' @@ -45,6 +52,6 @@ def clear(self, filename): :obj:`str`: The file location on disk that was removed. ''' - return self._api.post('file/clear', json={ - 'filename': self._check('filename', filename, str) - }).json()['response']['filename'] + return self._post('clear', + json={'filename': filename} + ).response.filename diff --git a/test-requirements.txt b/test-requirements.txt index e9fedcf44..f80851da6 100644 --- a/test-requirements.txt +++ b/test-requirements.txt @@ -18,3 +18,4 @@ docker>=3.7.2 certifi>=2023.7.22 # not directly required, pinned by Snyk to avoid a vulnerability requests>=2.31.0 # not directly required, pinned by Snyk to avoid a vulnerability black>=24.3.0 # not directly required, pinned by Snyk to avoid a vulnerability +requests-toolbelt==1.0.0 diff --git a/tests/sc/cassettes/sc_login.yaml b/tests/sc/cassettes/sc_login.yaml index 9a15d0fcf..4d0d85444 100644 --- a/tests/sc/cassettes/sc_login.yaml +++ b/tests/sc/cassettes/sc_login.yaml @@ -14,7 +14,7 @@ interactions: "enabled":"true","attributeSets":[]},{"name":"CSV","type":"csv","enabled":"true","attributeSets":[]}, {"name":"RTF","type":"rtf","enabled":"true","attributeSets":[]},{"name":"DISA ARF","type":"arf","enabled":"false", "attributeSets":["arf"]},{"name":"DISA ASR","type":"asr","enabled":"false","attributeSets":[]}, - {"name":"CyberScope","type":"lasr","enabled":"true","attributeSets":["lasr"]}],"version":"5.8.0", + {"name":"CyberScope","type":"lasr","enabled":"true","attributeSets":["lasr"]}],"version":"6.4.0", "buildID":"201810303290","banner":"","releaseID":"201810303290","uuid":"3ae050e2-486c-5c1a-a6a5-45305557cec3", "logo":"assets\/mariana\/images\/sc4icon.png","serverAuth":"any","serverClassification":"None", "sessionTimeout":"3600","licenseStatus":"Valid","ACAS":"false","freshInstall":"no","headerText":"", diff --git a/tests/sc/conftest.py b/tests/sc/conftest.py index 9ecf03f1b..ae4fc9424 100644 --- a/tests/sc/conftest.py +++ b/tests/sc/conftest.py @@ -2,7 +2,7 @@ configuration which would be used for testing sc. ''' import os - +import responses import pytest from tenable.errors import APIError, NotFoundError @@ -112,3 +112,18 @@ def teardown(): request.addfinalizer(teardown) return grp + + +@pytest.fixture +def tsc(): + with responses.RequestsMock() as rsps: + rsps.get('https://nourl/rest/system', + json={ + 'error_code': None, + 'response': {'version': '6.4.0'} + } + ) + return TenableSC(url='https://nourl', + access_key='SOMETHING', + secret_key='SECRET' + ) diff --git a/tests/sc/test_files.py b/tests/sc/test_files.py index 478d3d998..b8efb1cf7 100644 --- a/tests/sc/test_files.py +++ b/tests/sc/test_files.py @@ -1,18 +1,33 @@ ''' test file for testing various scenarios in file functionality ''' -import os - +from io import BytesIO +import responses +from responses.matchers import json_params_matcher import pytest -@pytest.mark.vcr() -def test_files_upload_clear_success(security_center): - ''' - test files upload and clear for success - ''' - with open('1000003.xml', 'w+') as file: - response = security_center.files.upload(fobj=file) - filename = security_center.files.clear(response) - assert filename == 'filename' - os.remove('1000003.xml') +@responses.activate +def test_files_upload(tsc): + fake = BytesIO(b'Test File') + responses.post('https://nourl/rest/file/upload', + json={ + 'error_code': 0, + 'response': {'filename': 'something'} + }, + ) + + assert 'something' == tsc.files.upload(fake) + +@responses.activate +def test_files_clear(tsc): + responses.post('https://nourl/rest/file/delete', + match=[ + json_params_matcher({'filename': 'testfile'}) + ], + json={ + 'error_code': 0, + 'response': {'filename': 'testfile'} + } + ) + assert 'testfile' == tsc.files.clear('testfile')