from __future__ import absolute_import, print_function import sys import os import traceback import click import itertools from netlib.http import CONTENT_MISSING import netlib.utils from . import flow, filt, contentviews from .exceptions import ContentViewException from .models import HTTPRequest class DumpError(Exception): pass class Options(object): attributes = [ "app", "app_host", "app_port", "anticache", "anticomp", "client_replay", "filtstr", "flow_detail", "keepserving", "kill", "no_server", "nopop", "refresh_server_playback", "replacements", "rfile", "rheaders", "setheaders", "server_replay", "scripts", "showhost", "stickycookie", "stickyauth", "stream_large_bodies", "verbosity", "outfile", "replay_ignore_content", "replay_ignore_params", "replay_ignore_payload_params", "replay_ignore_host" ] def __init__(self, **kwargs): for k, v in kwargs.items(): setattr(self, k, v) for i in self.attributes: if not hasattr(self, i): setattr(self, i, None) class DumpMaster(flow.FlowMaster): def __init__(self, server, options, outfile=None): flow.FlowMaster.__init__(self, server, flow.State()) self.outfile = outfile self.o = options self.anticache = options.anticache self.anticomp = options.anticomp self.showhost = options.showhost self.replay_ignore_params = options.replay_ignore_params self.replay_ignore_content = options.replay_ignore_content self.replay_ignore_host = options.replay_ignore_host self.refresh_server_playback = options.refresh_server_playback self.replay_ignore_payload_params = options.replay_ignore_payload_params self.set_stream_large_bodies(options.stream_large_bodies) if options.filtstr: self.filt = filt.parse(options.filtstr) else: self.filt = None if options.stickycookie: self.set_stickycookie(options.stickycookie) if options.stickyauth: self.set_stickyauth(options.stickyauth) if options.outfile: path = os.path.expanduser(options.outfile[0]) try: f = open(path, options.outfile[1]) self.start_stream(f, self.filt) except IOError as v: raise DumpError(v.strerror) if options.replacements: for i in options.replacements: self.replacehooks.add(*i) if options.setheaders: for i in options.setheaders: self.setheaders.add(*i) if options.server_replay: self.start_server_playback( self._readflow(options.server_replay), options.kill, options.rheaders, not options.keepserving, options.nopop, options.replay_ignore_params, options.replay_ignore_content, options.replay_ignore_payload_params, options.replay_ignore_host ) if options.client_replay: self.start_client_playback( self._readflow(options.client_replay), not options.keepserving ) scripts = options.scripts or [] for command in scripts: err = self.load_script(command) if err: raise DumpError(err) if options.rfile: try: self.load_flows_file(options.rfile) except flow.FlowReadError as v: self.add_event("Flow file corrupted.", "error") raise DumpError(v) if self.o.app: self.start_app(self.o.app_host, self.o.app_port) def _readflow(self, paths): """ Utitility function that reads a list of flows or raises a DumpError if that fails. """ try: return flow.read_flows_from_paths(paths) except flow.FlowReadError as e: raise DumpError(e.strerror) def add_event(self, e, level="info"): needed = dict(error=0, info=1, debug=2).get(level, 1) if self.o.verbosity >= needed: self.echo( e, fg="red" if level == "error" else None, dim=(level == "debug"), err=(level == "error") ) @staticmethod def indent(n, text): l = str(text).strip().splitlines() pad = " " * n return "\n".join(pad + i for i in l) def echo(self, text, indent=None, **style): if indent: text = self.indent(indent, text) click.secho(text, file=self.outfile, **style) def _echo_message(self, message): if self.o.flow_detail >= 2: headers = "\r\n".join( "{}: {}".format( click.style(k, fg="blue", bold=True), click.style(v, fg="blue")) for k, v in message.headers.fields ) self.echo(headers, indent=4) if self.o.flow_detail >= 3: if message.content == CONTENT_MISSING: self.echo("(content missing)", indent=4) elif message.content: self.echo("") try: type, lines = contentviews.get_content_view( contentviews.get("Auto"), message.content, headers=message.headers ) except ContentViewException: s = "Content viewer failed: \n" + traceback.format_exc() self.add_event(s, "debug") type, lines = contentviews.get_content_view( contentviews.get("Raw"), message.content, headers=message.headers ) styles = dict( highlight=dict(bold=True), offset=dict(fg="blue"), header=dict(fg="green", bold=True), text=dict(fg="green") ) def colorful(line): yield u" " # we can already indent here for (style, text) in line: yield click.style(text, **styles.get(style, {})) if self.o.flow_detail == 3: lines_to_echo = itertools.islice(lines, 70) else: lines_to_echo = lines lines_to_echo = list(lines_to_echo) content = u"\r\n".join( u"".join(colorful(line)) for line in lines_to_echo ) self.echo(content) if next(lines, None): self.echo("(cut off)", indent=4, dim=True) if self.o.flow_detail >= 2: self.echo("") def _echo_request_line(self, flow): if flow.request.stickycookie: stickycookie = click.style("[stickycookie] ", fg="yellow", bold=True) else: stickycookie = "" if flow.client_conn: client = click.style(flow.client_conn.address.host, bold=True) else: client = click.style("[replay]", fg="yellow", bold=True) method = flow.request.method method_color=dict( GET="green", DELETE="red" ).get(method.upper(), "magenta") method = click.style(method, fg=method_color, bold=True) if self.showhost: url = flow.request.pretty_url else: url = flow.request.url url = click.style(url, bold=True) line = "{stickycookie}{client} {method} {url}".format( stickycookie=stickycookie, client=client, method=method, url=url ) self.echo(line) def _echo_response_line(self, flow): if flow.response.is_replay: replay = click.style("[replay] ", fg="yellow", bold=True) else: replay = "" code = flow.response.status_code code_color = None if 200 <= code < 300: code_color = "green" elif 300 <= code < 400: code_color = "magenta" elif 400 <= code < 600: code_color = "red" code = click.style(str(code), fg=code_color, bold=True, blink=(code == 418)) reason = click.style(flow.response.reason, fg=code_color, bold=True) if flow.response.content == CONTENT_MISSING: size = "(content missing)" else: size = netlib.utils.pretty_size(len(flow.response.content)) size = click.style(size, bold=True) arrows = click.style("<<", bold=True) line = "{replay} {arrows} {code} {reason} {size}".format( replay=replay, arrows=arrows, code=code, reason=reason, size=size ) self.echo(line) def echo_flow(self, f): if self.o.flow_detail == 0: return if f.request: self._echo_request_line(f) self._echo_message(f.request) if f.response: self._echo_response_line(f) self._echo_message(f.response) if f.error: self.echo(" << {}".format(f.error.msg), bold=True, fg="red") if self.outfile: self.outfile.flush() def _process_flow(self, f): self.state.delete_flow(f) if self.filt and not f.match(self.filt): return self.echo_flow(f) def handle_request(self, f): flow.FlowMaster.handle_request(self, f) if f: f.reply() return f def handle_response(self, f): flow.FlowMaster.handle_response(self, f) if f: f.reply() self._process_flow(f) return f def handle_error(self, f): flow.FlowMaster.handle_error(self, f) if f: self._process_flow(f) return f def shutdown(self): # pragma: no cover return flow.FlowMaster.shutdown(self) def run(self): # pragma: no cover if self.o.rfile and not self.o.keepserving: self.shutdown() return try: return flow.FlowMaster.run(self) except BaseException: self.shutdown() raise