diff --git a/OTXv2.py b/OTXv2.py index 9c14f6f..ebe8b3c 100755 --- a/OTXv2.py +++ b/OTXv2.py @@ -88,6 +88,14 @@ def __str__(self): return repr(self.value) +class TimeoutError(Exception): + def __init__(self, value=None): + self.value = value or "TimeoutError" + + def __str__(self): + return repr(self.value) + + class OTXv2(object): """ Main class to interact with the AlienVault OTX API. @@ -95,7 +103,7 @@ class OTXv2(object): def __init__( self, api_key, proxy=None, proxy_https=None, server="https://otx.alienvault.com", project="SDK", - user_agent=None, verify=True, cert=None + user_agent=None, verify=True, cert=None, timeout=None ): self.key = api_key self.server = server @@ -114,6 +122,7 @@ def __init__( 'User-Agent': user_agent or 'OTX Python {}/1.5.12'.format(project), 'Content-Type': 'application/json' } + self.timeout = timeout def session(self): if self.request_session is None: @@ -181,10 +190,13 @@ def get(self, url, **kwargs): proxies=self.proxies, verify=self.verify, cert=self.cert, + timeout=self.timeout ) return self.handle_response_errors(response).json() except requests.exceptions.RetryError: raise RetryError() + except requests.exceptions.Timeout: + raise TimeoutError() def patch(self, url, body, **kwargs): """ @@ -193,16 +205,19 @@ def patch(self, url, body, **kwargs): :param body: HTTP Body to send in request :return: response as dict """ - - response = self.session().patch( - self.create_url(url, **kwargs), - data=json.dumps(body), - headers=self.headers, - proxies=self.proxies, - verify=self.verify, - cert=self.cert, - ) - return self.handle_response_errors(response).json() + try: + response = self.session().patch( + self.create_url(url, **kwargs), + data=json.dumps(body), + headers=self.headers, + proxies=self.proxies, + verify=self.verify, + cert=self.cert, + timeout=self.timeout + ) + return self.handle_response_errors(response).json() + except requests.exceptions.Timeout: + raise TimeoutError() def post(self, url, body=None, headers=None, files=None, **kwargs): """ @@ -213,17 +228,20 @@ def post(self, url, body=None, headers=None, files=None, **kwargs): :param files: (optional) list of file tuples, if posting multipart form data :return: response as dict """ - - response = self.session().post( - self.create_url(url, **kwargs), - data=json.dumps(body) if body else None, - files=files, - headers=headers or self.headers, - proxies=self.proxies, - verify=self.verify, - cert=self.cert, - ) - return self.handle_response_errors(response).json() + try: + response = self.session().post( + self.create_url(url, **kwargs), + data=json.dumps(body) if body else None, + files=files, + headers=headers or self.headers, + proxies=self.proxies, + verify=self.verify, + cert=self.cert, + timeout=self.timeout + ) + return self.handle_response_errors(response).json() + except requests.exceptions.Timeout: + raise TimeoutError() def create_pulse(self, **kwargs): """