From b151c5831d9019fb0a8f6db943256c9716999d41 Mon Sep 17 00:00:00 2001 From: joarleymoraes Date: Fri, 5 May 2017 16:19:00 -0300 Subject: [PATCH 1/2] added support to set timeout for TAXII requests --- cabby/abstract.py | 8 ++++++-- cabby/dispatcher.py | 13 ++++++++----- tests/test_common.py | 15 +++++++++++++++ 3 files changed, 29 insertions(+), 7 deletions(-) diff --git a/cabby/abstract.py b/cabby/abstract.py index 38cb16e..f987e53 100644 --- a/cabby/abstract.py +++ b/cabby/abstract.py @@ -47,6 +47,7 @@ def __init__(self, host=None, discovery_path=None, port=None, self.jwt_token = None self.headers = headers or {} + self.timeout = None self.log = logging.getLogger( "{}.{}".format(self.__module__, self.__class__.__name__)) @@ -161,6 +162,7 @@ def _execute_request(self, request, uri=None, service_type=None): services by ``service_type``. ''' + timeout = self.timeout if not uri and not service_type: raise NoURIProvidedError('URI or service_type needed') elif not uri: @@ -193,13 +195,15 @@ def _execute_request(self, request, uri=None, service_type=None): 'key_file': self.key_file, 'key_password': self.key_password, 'ca_cert': self.ca_cert - }) + }, + timeout=timeout) else: message = dispatcher.send_taxii_request( session, self._prepare_url(uri), request, - taxii_binding=self.taxii_binding) + taxii_binding=self.taxii_binding, + timeout=timeout) return message diff --git a/cabby/dispatcher.py b/cabby/dispatcher.py index 2f8963f..383c4c0 100644 --- a/cabby/dispatcher.py +++ b/cabby/dispatcher.py @@ -34,7 +34,7 @@ def raise_http_error(status_code, response_stream=None): def send_taxii_request(session, url, request, taxii_binding=None, - tls_details=None): + tls_details=None, timeout=None): ''' Send XML message to a TAXII service and parse a response. ''' @@ -55,7 +55,7 @@ def send_taxii_request(session, url, request, taxii_binding=None, # https://github.com/kennethreitz/requests/issues/2519 is fixed try: response = get_response_using_key_pass( - url, request_body, session, **tls_details) + url, request_body, session, timeout=timeout, **tls_details) except urllib.error.HTTPError as e: log.error( "Error while connecting to {}".format(url), @@ -64,7 +64,7 @@ def send_taxii_request(session, url, request, taxii_binding=None, stream, headers = response, response.headers else: - response = session.post(url, data=request_body, stream=True) + response = session.post(url, data=request_body, stream=True, timeout=timeout) if not response.ok: raise_http_error(response.status_code, response.raw) @@ -369,7 +369,7 @@ def obtain_jwt_token(session, jwt_url, username, password): def get_response_using_key_pass(url, data, session, cert_file, key_file, - key_password, ca_cert=None): + key_password, ca_cert=None, timeout=None): if sys.version_info < (2, 7, 9): raise ValueError( @@ -405,4 +405,7 @@ def get_response_using_key_pass(url, data, session, cert_file, key_file, request = urllib.request.Request(url, data, headers) - return opener.open(request) + if timeout: + return opener.open(request, timeout=timeout) + else: + return opener.open(request) diff --git a/tests/test_common.py b/tests/test_common.py index 9c9b4f4..8fa1370 100644 --- a/tests/test_common.py +++ b/tests/test_common.py @@ -3,6 +3,7 @@ import json import gzip import sys +import requests from six import StringIO @@ -222,3 +223,17 @@ def test_gzip_response(version): httpretty.disable() httpretty.reset() + + +@pytest.mark.parametrize("version", [11, 10]) +def test_timeout(version): + + client = make_client(version) + client.timeout = 1 + # + # Connect to an existing host but to a port that is blocked by the firewall + # If no timeout is set, this would take forever to return + # + with pytest.raises(requests.exceptions.Timeout): + client.discover_services(uri='http://httpbin.org:81/') + From bf53cf7ba175c0c157307e3119a193fbb7f0f2c1 Mon Sep 17 00:00:00 2001 From: joarleymoraes Date: Mon, 15 May 2017 17:18:52 -0300 Subject: [PATCH 2/2] addressing the issues by @traut --- cabby/abstract.py | 9 ++++----- cabby/dispatcher.py | 3 ++- tests/test_common.py | 34 +++++++++++++++++++++++++--------- 3 files changed, 31 insertions(+), 15 deletions(-) diff --git a/cabby/abstract.py b/cabby/abstract.py index f987e53..4577a61 100644 --- a/cabby/abstract.py +++ b/cabby/abstract.py @@ -25,7 +25,7 @@ class AbstractClient(object): taxii_version = None def __init__(self, host=None, discovery_path=None, port=None, - use_https=False, headers=None): + use_https=False, headers=None, timeout=None): self.host = host self.port = port @@ -47,7 +47,7 @@ def __init__(self, host=None, discovery_path=None, port=None, self.jwt_token = None self.headers = headers or {} - self.timeout = None + self.timeout = timeout self.log = logging.getLogger( "{}.{}".format(self.__module__, self.__class__.__name__)) @@ -162,7 +162,6 @@ def _execute_request(self, request, uri=None, service_type=None): services by ``service_type``. ''' - timeout = self.timeout if not uri and not service_type: raise NoURIProvidedError('URI or service_type needed') elif not uri: @@ -196,14 +195,14 @@ def _execute_request(self, request, uri=None, service_type=None): 'key_password': self.key_password, 'ca_cert': self.ca_cert }, - timeout=timeout) + timeout=self.timeout) else: message = dispatcher.send_taxii_request( session, self._prepare_url(uri), request, taxii_binding=self.taxii_binding, - timeout=timeout) + timeout=self.timeout) return message diff --git a/cabby/dispatcher.py b/cabby/dispatcher.py index 383c4c0..b4dff46 100644 --- a/cabby/dispatcher.py +++ b/cabby/dispatcher.py @@ -64,7 +64,8 @@ def send_taxii_request(session, url, request, taxii_binding=None, stream, headers = response, response.headers else: - response = session.post(url, data=request_body, stream=True, timeout=timeout) + response = session.post(url, data=request_body, stream=True, + timeout=timeout) if not response.ok: raise_http_error(response.status_code, response.raw) diff --git a/tests/test_common.py b/tests/test_common.py index 8fa1370..85c53d0 100644 --- a/tests/test_common.py +++ b/tests/test_common.py @@ -4,6 +4,7 @@ import gzip import sys import requests +from time import sleep from six import StringIO @@ -56,7 +57,6 @@ def get_sent_message(version): @pytest.mark.parametrize("version", [11, 10]) def test_set_headers(version): - httpretty.reset() httpretty.enable() @@ -87,7 +87,6 @@ def test_set_headers(version): @pytest.mark.parametrize("version", [11, 10]) def test_invalid_response(version): - httpretty.reset() httpretty.enable() @@ -113,7 +112,6 @@ def test_invalid_response(version): @pytest.mark.parametrize("version", [11, 10]) def test_invalid_response_status(version): - httpretty.reset() httpretty.enable() @@ -133,7 +131,6 @@ def test_invalid_response_status(version): @pytest.mark.parametrize("version", [11, 10]) def test_jwt_auth_response(version): - httpretty.reset() httpretty.enable() @@ -203,7 +200,6 @@ def compress(text): @pytest.mark.parametrize("version", [11, 10]) def test_gzip_response(version): - httpretty.reset() httpretty.enable() @@ -227,13 +223,33 @@ def test_gzip_response(version): @pytest.mark.parametrize("version", [11, 10]) def test_timeout(version): + httpretty.reset() + httpretty.enable() + + timeout_in_sec = 1 client = make_client(version) - client.timeout = 1 # - # Connect to an existing host but to a port that is blocked by the firewall - # If no timeout is set, this would take forever to return + # configure to raise the error before the timeout # + client.timeout = timeout_in_sec / 2.0 + + def timeout_request_callback(request, uri, headers): + sleep(timeout_in_sec) + + return 200, headers, {'result': 'success'} + + uri = get_fix(version).DISCOVERY_URI_HTTP + + httpretty.register_uri( + httpretty.POST, + uri, + body=timeout_request_callback, + content_type='application/json' + ) + with pytest.raises(requests.exceptions.Timeout): - client.discover_services(uri='http://httpbin.org:81/') + client.discover_services(uri=uri) + httpretty.disable() + httpretty.reset()