diff --git a/sherlock/sherlock.py b/sherlock/sherlock.py index 46be623d..7a2ca676 100644 --- a/sherlock/sherlock.py +++ b/sherlock/sherlock.py @@ -34,30 +34,56 @@ global proxy_list proxy_list = [] -class ElapsedFuturesSession(FuturesSession): - """ - Extends FutureSession to add a response time metric to each request. - - This is taken (almost) directly from here: https://github.com/ross/requests-futures#working-in-the-background - """ - +class SherlockFuturesSession(FuturesSession): def request(self, method, url, hooks={}, *args, **kwargs): + """Request URL. + + This extends the FuturesSession request method to calculate a response + time metric to each request. + + It is taken (almost) directly from the following StackOverflow answer: + https://github.com/ross/requests-futures#working-in-the-background + + Keyword Arguments: + self -- This object. + method -- String containing method desired for request. + url -- String containing URL for request. + hooks -- Dictionary containing hooks to execute after + request finishes. + args -- Arguments. + kwargs -- Keyword arguments. + + Return Value: + Request object. + """ start = time() - def timing(r, *args, **kwargs): + def response_time(resp, *args, **kwargs): elapsed_sec = time() - start - r.elapsed = round(elapsed_sec * 1000) + resp.elapsed = round(elapsed_sec * 1000) + #Install hook to execute when response completes. + #Make sure that the time measurement hook is first, so we will not + #track any later hook's execution time. try: - if isinstance(hooks['response'], (list, tuple)): - # needs to be first so we don't time other hooks execution - hooks['response'].insert(0, timing) + if isinstance(hooks['response'], list): + hooks['response'].insert(0, response_time) + elif isinstance(hooks['response'], tuple): + #Convert tuple to list and insert time measurement hook first. + hooks['response'] = list(hooks['response']) + hooks['response'].insert(0, response_time) else: - hooks['response'] = [timing, hooks['response']] + #Must have previously contained a single hook function, + #so convert to list. + hooks['response'] = [response_time, hooks['response']] except KeyError: - hooks['response'] = timing + #No response hook was already defined, so install it ourselves. + hooks['response'] = [response_time] - return super(ElapsedFuturesSession, self).request(method, url, hooks=hooks, *args, **kwargs) + return super(SherlockFuturesSession, self).request(method, + url, + hooks=hooks, + *args, **kwargs) def print_info(title, info): @@ -177,9 +203,10 @@ def sherlock(username, site_data, verbose=False, tor=False, unique_tor=False, pr underlying_session = requests.session() underlying_request = requests.Request() - # Create multi-threaded session for all requests. Use our custom FuturesSession that exposes response time - session = ElapsedFuturesSession( - executor=executor, session=underlying_session) + #Create multi-threaded session for all requests. + session = SherlockFuturesSession(executor=executor, + session=underlying_session) + # Results from analysis of all sites results_total = {}