Skip to content

Commit

Permalink
Fixed nits
Browse files Browse the repository at this point in the history
  • Loading branch information
Leeren Chang committed Feb 13, 2018
1 parent 9b02a77 commit 4e4b537
Showing 1 changed file with 22 additions and 3 deletions.
25 changes: 22 additions & 3 deletions threat_intel/util/http.py
Original file line number Diff line number Diff line change
Expand Up @@ -101,9 +101,15 @@ def _cull(self):

class AvailabilityLimiter(object):

"""Limits the total number of requests issued for a session"""
"""Limits the total number of requests issued for a session."""

def __init__(self, total_retries):
""" Wrapper object for managing total session retry limit.
Args:
total_retries: Total request attempts to be made per sesssion.
This is shared between all request objects.
"""
self.total_retries = total_retries

def map_with_retries(self, requests, responses_for_requests, *args, **kwargs):
Expand All @@ -115,10 +121,13 @@ def map_with_retries(self, requests, responses_for_requests, *args, **kwargs):
:param exception_handler: Callback function, called when exception occured. Params: Request, Exception
:param args: Additional arguments to pass into a retry mapping call
:param kwargs: Keyword arguments passed into the grequests wrapper, e.g. exception_handler
"""
_exception_handler = kwargs.pop('exception_handler', None)
retries = []


def exception_handler(request, exception, *args, **kwargs):
if self.total_retries > 0:
self.total_retries -= 1
Expand All @@ -130,19 +139,22 @@ def exception_handler(request, exception, *args, **kwargs):

for request, response in zip(requests, responses):
if response is not None and response.status_code == 403:
logging.debug('Request was received with a 403 response status code.')
raise InvalidRequestError('Access forbidden')
if response:
responses_for_requests[request] = response

# Recursively retry failed requests with the modified total retry count
if retries:
self.map_with_retries(retries, responses_for_requests, *args, exception_handler=_exception_handler, **kwargs)
self.map_with_retries(retries, responses_for_requests,
*args, exception_handler=_exception_handler, **kwargs)

class MultiRequest(object):

"""Wraps grequests to make simultaneous HTTP requests.
Can use a RateLimiter to limit # of outstanding requests.
Can also use AvailabilityLimiter to limit total # of request issuance threshold.
`multi_get` and `multi_post` try to be smart about how many requests to issue:
* One url & one param - One request will be made.
Expand All @@ -153,14 +165,21 @@ class MultiRequest(object):
_VERB_GET = 'GET'
_VERB_POST = 'POST'

def __init__(self, default_headers=None, max_requests=20, rate_limit=0, req_timeout=None, max_retry=10, total_retry=100):
def __init__(self, default_headers=None, max_requests=20,
rate_limit=0, req_timeout=None, max_retry=10, total_retry=100):
"""Create the MultiRequest.
Args:
default_headers - A dict of headers which will be added to every request
max_requests - Maximum number of requests to issue at once
rate_limit - Maximum number of requests to issue per second
req_timeout - Maximum number of seconds to wait without reading a response byte before deciding an error has occurred
max_retry - The total number of attempts to retry a single batch of requests
total_retry - The total number of request retries that can be made through the entire session
Note there is a difference between `max_retry` and `total_retry`:
- `max_retry` refers to how many times a batch of requests will be re-issued collectively
- `total_retry` refers to a limit on the total number of outstanding requests made
Once the latter is exhausted, no failed request within the whole session will be retried.
"""
self._default_headers = default_headers
self._max_requests = max_requests
Expand Down

0 comments on commit 4e4b537

Please sign in to comment.