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..51670bb80 100644 --- a/tests/sc/test_files.py +++ b/tests/sc/test_files.py @@ -1,18 +1,25 @@ ''' test file for testing various scenarios in file functionality ''' -import os - +from io import BytesIO +import responses 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': '', + 'response': {'filename': 'something'} + }, + ) + + assert 'something' == tsc.files.upload(fake) + + #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')