From 0b24623779e95f1f75aa6fe0a9082689a47a2c0f Mon Sep 17 00:00:00 2001 From: Kushal Das Date: Wed, 19 Jun 2019 15:33:43 +0530 Subject: [PATCH] Fixes #93 raises the TimeOutErrors properly --- sdclientapi/__init__.py | 126 +++++++++++++++++++--------------------- tests/test_api.py | 17 ++++++ tests/test_apiproxy.py | 31 ++++++++++ 3 files changed, 108 insertions(+), 66 deletions(-) diff --git a/sdclientapi/__init__.py b/sdclientapi/__init__.py index 7d964b1..a869814 100644 --- a/sdclientapi/__init__.py +++ b/sdclientapi/__init__.py @@ -570,41 +570,38 @@ def download_submission( if os.path.exists(path) and not os.path.isdir(path): raise BaseError("Please provide a vaild directory to save.") - try: - data, status_code, headers = self._send_json_request( - method, - path_query, - headers=self.req_headers, - timeout=timeout or self.default_download_timeout, - ) + data, status_code, headers = self._send_json_request( + method, + path_query, + headers=self.req_headers, + timeout=timeout or self.default_download_timeout, + ) - if status_code == 404: - raise WrongUUIDError("Missing submission {}".format(submission.uuid)) + if status_code == 404: + raise WrongUUIDError("Missing submission {}".format(submission.uuid)) - # Get the headers - headers = headers - - if not self.proxy: - # This is where we will save our downloaded file - filepath = os.path.join(path, submission.filename) - with open(filepath, "wb") as fobj: - for chunk in data.iter_content( - chunk_size=1024 - ): # Getting 1024 in each chunk - if chunk: - fobj.write(chunk) - - else: - filepath = os.path.join( - "/home/user/QubesIncoming/", self.proxy_vm_name, data["filename"] - ) - # Return the tuple of sha256sum, filepath - # Returning empty string instead of sha256sum due to this - # SecureDrop server bug: - # https://github.com/freedomofpress/securedrop/issues/3877 - return "", filepath - except Exception as err: - raise BaseError(err) + # Get the headers + headers = headers + + if not self.proxy: + # This is where we will save our downloaded file + filepath = os.path.join(path, submission.filename) + with open(filepath, "wb") as fobj: + for chunk in data.iter_content( + chunk_size=1024 + ): # Getting 1024 in each chunk + if chunk: + fobj.write(chunk) + + else: + filepath = os.path.join( + "/home/user/QubesIncoming/", self.proxy_vm_name, data["filename"] + ) + # Return the tuple of sha256sum, filepath + # Returning empty string instead of sha256sum due to this + # SecureDrop server bug: + # https://github.com/freedomofpress/securedrop/issues/3877 + return "", filepath def flag_source(self, source: Source) -> bool: """ @@ -823,41 +820,38 @@ def download_reply(self, reply: Reply, path: str = "") -> Tuple[str, str]: if os.path.exists(path) and not os.path.isdir(path): raise BaseError("Please provide a valid directory to save.") - try: - data, status_code, headers = self._send_json_request( - method, - path_query, - headers=self.req_headers, - timeout=self.default_request_timeout, - ) + data, status_code, headers = self._send_json_request( + method, + path_query, + headers=self.req_headers, + timeout=self.default_request_timeout, + ) - if status_code == 404: - raise WrongUUIDError("Missing reply {}".format(reply.uuid)) + if status_code == 404: + raise WrongUUIDError("Missing reply {}".format(reply.uuid)) + + # Get the headers + headers = headers - # Get the headers - headers = headers - - if not self.proxy: - # This is where we will save our downloaded file - filepath = os.path.join(path, reply.filename) - with open(filepath, "wb") as fobj: - for chunk in data.iter_content( - chunk_size=1024 - ): # Getting 1024 in each chunk - if chunk: - fobj.write(chunk) - - else: - filepath = os.path.join( - "/home/user/QubesIncoming/", self.proxy_vm_name, data["filename"] - ) - # Return the tuple of sha256sum, filepath - # Returning empty string instead of sha256sum due to this - # SecureDrop server bug: - # https://github.com/freedomofpress/securedrop/issues/3877 - return "", filepath - except Exception as err: - raise BaseError(err) + if not self.proxy: + # This is where we will save our downloaded file + filepath = os.path.join(path, reply.filename) + with open(filepath, "wb") as fobj: + for chunk in data.iter_content( + chunk_size=1024 + ): # Getting 1024 in each chunk + if chunk: + fobj.write(chunk) + + else: + filepath = os.path.join( + "/home/user/QubesIncoming/", self.proxy_vm_name, data["filename"] + ) + # Return the tuple of sha256sum, filepath + # Returning empty string instead of sha256sum due to this + # SecureDrop server bug: + # https://github.com/freedomofpress/securedrop/issues/3877 + return "", filepath def delete_reply(self, reply: Reply) -> bool: """ diff --git a/tests/test_api.py b/tests/test_api.py index f73d457..94c80bd 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -18,6 +18,7 @@ from sdclientapi.sdlocalobjects import Reply from sdclientapi.sdlocalobjects import ReplyError from sdclientapi.sdlocalobjects import Source +from sdclientapi.sdlocalobjects import Submission from sdclientapi.sdlocalobjects import WrongUUIDError NUM_REPLIES_PER_SOURCE = 2 @@ -341,3 +342,19 @@ def test_request_read_timeout(mocker): mocker.patch("sdclientapi.requests.request", side_effect=ReadTimeout) with pytest.raises(RequestTimeoutError): api.authenticate() + + +def test_download_reply_timeout(mocker): + api = API("mock", "mock", "mock", "mock", proxy=False) + mocker.patch("sdclientapi.requests.request", side_effect=RequestTimeoutError) + with pytest.raises(RequestTimeoutError): + r = Reply(uuid="humanproblem", filename="secret.txt") + api.download_reply(r) + + +def test_download_submission_timeout(mocker): + api = API("mock", "mock", "mock", "mock", proxy=False) + mocker.patch("sdclientapi.requests.request", side_effect=RequestTimeoutError) + with pytest.raises(RequestTimeoutError): + s = Submission(uuid="climateproblem") + api.download_submission(s) diff --git a/tests/test_apiproxy.py b/tests/test_apiproxy.py index 8fdbdb4..19a3837 100644 --- a/tests/test_apiproxy.py +++ b/tests/test_apiproxy.py @@ -16,6 +16,7 @@ ReplyError, Reply, Source, + Submission, ) from utils import load_auth, save_auth, dastollervey_datasaver @@ -329,3 +330,33 @@ def communicate(self, *nargs, **kwargs) -> None: with pytest.raises(RequestTimeoutError): api.authenticate() + + +def test_download_reply_timeout(mocker): + class MockedPopen: + def __init__(self, *nargs, **kwargs) -> None: + self.stdin = mocker.MagicMock() + + def communicate(self, *nargs, **kwargs) -> None: + raise TimeoutExpired(["mock"], 123) + + api = API("mock", "mock", "mock", "mock", proxy=True) + mocker.patch("sdclientapi.Popen", MockedPopen) + with pytest.raises(RequestTimeoutError): + r = Reply(uuid="humanproblem", filename="secret.txt") + api.download_reply(r) + + +def test_download_submission_timeout(mocker): + class MockedPopen: + def __init__(self, *nargs, **kwargs) -> None: + self.stdin = mocker.MagicMock() + + def communicate(self, *nargs, **kwargs) -> None: + raise TimeoutExpired(["mock"], 123) + + api = API("mock", "mock", "mock", "mock", proxy=True) + mocker.patch("sdclientapi.Popen", MockedPopen) + with pytest.raises(RequestTimeoutError): + s = Submission(uuid="climateproblem") + api.download_submission(s)