from io import BytesIO try: import httplib except ImportError: import http.client as httplib orig_connection = httplib.HTTPConnection from contextlib import contextmanager import ssl from array import array from time import sleep BUFF_SIZE = 8192 # ============================================================================ class RecordingStream(object): def __init__(self, fp, recorder): self.fp = fp self.recorder = recorder self.incomplete = False if hasattr(self.fp, 'unread'): self.unread = self.fp.unread if hasattr(self.fp, 'tell'): self.tell = self.fp.tell def read(self, *args, **kwargs): buff = self.fp.read(*args, **kwargs) self.recorder.write_response_buff(buff) return buff def readinto(self, buff): res = self.fp.readinto(buff) self.recorder.write_response_buff(buff) return res def readline(self, maxlen=None): line = self.fp.readline(maxlen) self.recorder.write_response_header_line(line) return line def flush(self): self.fp.flush() def close(self): try: self.recorder.finish_response(self.incomplete) except Exception as e: import traceback traceback.print_exc() res = self.fp.close() return res # ============================================================================ class RecordingHTTPResponse(httplib.HTTPResponse): def __init__(self, recorder, *args, **kwargs): httplib.HTTPResponse.__init__(self, *args, **kwargs) self.fp = RecordingStream(self.fp, recorder) def mark_incomplete(self): self.fp.incomplete = True # ============================================================================ class RecordingHTTPConnection(httplib.HTTPConnection): global_recorder_maker = None def __init__(self, *args, **kwargs): orig_connection.__init__(self, *args, **kwargs) if not self.global_recorder_maker: self.recorder = None else: self.recorder = self.global_recorder_maker() def make_recording_response(*args, **kwargs): return RecordingHTTPResponse(self.recorder, *args, **kwargs) self.response_class = make_recording_response def send(self, data): if not self.recorder: orig_connection.send(self, data) return if hasattr(data,'read') and not isinstance(data, array): url = None while True: buff = data.read(BUFF_SIZE) if not buff: break orig_connection.send(self, buff) self.recorder.write_request(url, buff) else: orig_connection.send(self, data) self.recorder.write_request(self, data) def get_url(self, data): try: buff = BytesIO(data) line = buff.readline() path = line.split(' ', 2)[1] host = self.host port = self.port scheme = 'https' if isinstance(self.sock, ssl.SSLSocket) else 'http' url = scheme + '://' + host if (scheme == 'https' and port != '443') and (scheme == 'http' and port != '80'): url += ':' + port url += path except Exception as e: raise return url def request(self, *args, **kwargs): #if self.recorder: # self.recorder.start_request(self) res = orig_connection.request(self, *args, **kwargs) if self.recorder: self.recorder.finish_request(self.sock) return res # ============================================================================ class BaseRecorder(object): def write_request(self, conn, buff): #url = conn.get_url() pass def write_response_header_line(self, line): pass def write_response_buff(self, buff): pass def finish_request(self, socket): pass def finish_response(self, incomplete=False): pass #================================================================= class ReadFullyStream(object): def __init__(self, stream): self.stream = stream def read(self, *args, **kwargs): try: return self.stream.read(*args, **kwargs) except: self.mark_incomplete() raise def readline(self, *args, **kwargs): try: return self.stream.readline(*args, **kwargs) except: self.mark_incomplete() raise def mark_incomplete(self): if (hasattr(self.stream, '_fp') and hasattr(self.stream._fp, 'mark_incomplete')): self.stream._fp.mark_incomplete() def close(self): try: while True: buff = self.stream.read(BUFF_SIZE) sleep(0) if not buff: break except Exception as e: import traceback traceback.print_exc() self.mark_incomplete() finally: self.stream.close() # ============================================================================ httplib.HTTPConnection = RecordingHTTPConnection # ============================================================================ class DefaultRecorderMaker(object): def __call__(self): return BaseRecorder() class FixedRecorder(object): def __init__(self, recorder): self.recorder = recorder def __call__(self): return self.recorder @contextmanager def record_requests(url, recorder_maker): RecordingHTTPConnection.global_recorder_maker = recorder_maker yield RecordingHTTPConnection.global_recorder_maker = None @contextmanager def orig_requests(): httplib.HTTPConnection = orig_connection yield httplib.HTTPConnection = RecordingHTTPConnection import requests as patched_requests def request(url, method='GET', recorder=None, recorder_maker=None, session=patched_requests, **kwargs): if kwargs.get('skip_recording'): recorder_maker = None elif recorder: recorder_maker = FixedRecorder(recorder) elif not recorder_maker: recorder_maker = DefaultRecorderMaker() with record_requests(url, recorder_maker): kwargs['allow_redirects'] = False r = session.request(method=method, url=url, **kwargs) return r