add --stream options, various fixes

This commit is contained in:
Maximilian Hils 2014-07-21 21:06:55 +02:00
parent 562ac9e721
commit 4b4a18a2e4
7 changed files with 111 additions and 61 deletions

5
examples/stream.py Normal file
View File

@ -0,0 +1,5 @@
def responseheaders(ctx, flow):
"""
Enables streaming for all responses.
"""
flow.response.stream = True

View File

@ -27,6 +27,15 @@ def request(ctx, flow):
"""
ctx.log("request")
def responseheaders(ctx, flow):
"""
Called when the response headers for a server response have been received,
but the response body has not been processed yet. Can be used to tell mitmproxy
to stream the response.
"""
ctx.log("responseheaders")
def response(ctx, flow):
"""
Called when a server response has been received.

View File

@ -3,7 +3,7 @@ import re
import argparse
from argparse import ArgumentTypeError
from netlib import http
from . import proxy, filt
from . import filt, utils
from .proxy import config
APP_HOST = "mitm.it"
@ -98,7 +98,6 @@ def parse_setheader(s):
def parse_server_spec(url):
normalized_url = re.sub("^https?2", "", url)
p = http.parse_url(normalized_url)
@ -125,6 +124,8 @@ def get_common_options(options):
if options.stickyauth_filt:
stickyauth = options.stickyauth_filt
stream_large_bodies = utils.parse_size(options.stream_large_bodies)
reps = []
for i in options.replace:
try:
@ -143,7 +144,6 @@ def get_common_options(options):
raise ArgumentTypeError("Could not read replace file: %s" % path)
reps.append((patt, rex, v))
setheaders = []
for i in options.setheader:
try:
@ -172,6 +172,7 @@ def get_common_options(options):
scripts=options.scripts,
stickycookie=stickycookie,
stickyauth=stickyauth,
stream_large_bodies=stream_large_bodies,
showhost=options.showhost,
wfile=options.wfile,
verbosity=options.verbose,
@ -243,7 +244,13 @@ def common_options(parser):
help="Byte size limit of HTTP request and response bodies." \
" Understands k/m/g suffixes, i.e. 3m for 3 megabytes."
)
parser.add_argument(
"--stream",
action="store", dest="stream_large_bodies", default=None,
metavar="SIZE",
help="Stream data to the client if response body exceeds the given threshold. "
"If streamed, the body will not be stored in any way. Understands k/m/g suffixes, i.e. 3m for 3 megabytes."
)
group = parser.add_argument_group("Proxy Options")
# We could make a mutually exclusive group out of -R, -U, -T, but we don't do that because
@ -280,7 +287,6 @@ def common_options(parser):
help="Set transparent proxy mode."
)
group = parser.add_argument_group(
"Advanced Proxy Options",
"""
@ -304,7 +310,6 @@ def common_options(parser):
help="Override the destination server all requests are sent to: http[s][2http[s]]://host[:port]"
)
group = parser.add_argument_group("Web App")
group.add_argument(
"-a",
@ -329,7 +334,6 @@ def common_options(parser):
help="Serve the app outside of the proxy."
)
group = parser.add_argument_group("Client Replay")
group.add_argument(
"-c",
@ -367,7 +371,6 @@ def common_options(parser):
"This makes it possible to replay same response multiple times."
)
group = parser.add_argument_group(
"Replacements",
"""
@ -389,7 +392,6 @@ def common_options(parser):
help="Replacement pattern, where the replacement clause is a path to a file."
)
group = parser.add_argument_group(
"Set Headers",
"""
@ -405,7 +407,6 @@ def common_options(parser):
help="Header set pattern."
)
group = parser.add_argument_group(
"Proxy Authentication",
"""
@ -434,5 +435,4 @@ def common_options(parser):
help="Allow access to users specified in an Apache htpasswd file."
)
config.ssl_option_group(parser)

View File

@ -343,6 +343,7 @@ class Options(object):
"server_replay",
"stickycookie",
"stickyauth",
"stream_large_bodies",
"verbosity",
"wfile",
"nopop",
@ -391,6 +392,8 @@ class ConsoleMaster(flow.FlowMaster):
print >> sys.stderr, "Sticky auth error:", r
sys.exit(1)
self.set_stream_large_bodies(options.stream_large_bodies)
self.refresh_server_playback = options.refresh_server_playback
self.anticache = options.anticache
self.anticomp = options.anticomp

View File

@ -2,6 +2,7 @@ from __future__ import absolute_import
import sys, os
import netlib.utils
from . import flow, filt, utils
from .protocol import http
class DumpError(Exception): pass
@ -30,6 +31,7 @@ class Options(object):
"showhost",
"stickycookie",
"stickyauth",
"stream_large_bodies",
"verbosity",
"wfile",
]
@ -69,6 +71,8 @@ class DumpMaster(flow.FlowMaster):
self.showhost = options.showhost
self.refresh_server_playback = options.refresh_server_playback
self.set_stream_large_bodies(options.stream_large_bodies)
if filtstr:
self.filt = filt.parse(filtstr)
else:
@ -80,6 +84,7 @@ class DumpMaster(flow.FlowMaster):
if options.stickyauth:
self.set_stickyauth(options.stickyauth)
if options.wfile:
path = os.path.expanduser(options.wfile)
try:
@ -157,12 +162,17 @@ class DumpMaster(flow.FlowMaster):
if f.response:
if self.o.flow_detail > 0:
if f.response.content == http.CONTENT_MISSING:
sz = "(content missing)"
else:
sz = utils.pretty_size(len(f.response.content))
result = " << %s %s"%(str_response(f.response), sz)
if self.o.flow_detail > 1:
result = result + "\n\n" + self.indent(4, f.response.headers)
if self.o.flow_detail > 2:
if utils.isBin(f.response.content):
if f.response.content == http.CONTENT_MISSING:
cont = self.indent(4, "(content missing)")
elif utils.isBin(f.response.content):
d = netlib.utils.hexdump(f.response.content)
d = "\n".join("%s\t%s %s"%i for i in d)
cont = self.indent(4, d)

View File

@ -145,6 +145,17 @@ class SetHeaders:
f.request.headers.add(header, value)
class StreamLargeBodies(object):
def __init__(self, max_size):
self.max_size = max_size
def run(self, flow, is_request):
r = flow.request if is_request else flow.response
code = flow.response.code if flow.response else None
expected_size = netlib.http.expected_http_body_size(r.headers, is_request, flow.request.method, code)
if not (0 <= expected_size <= self.max_size):
r.stream = True
class ClientPlaybackState:
def __init__(self, flows, exit):
self.flows, self.exit = flows, exit
@ -437,6 +448,7 @@ class FlowMaster(controller.Master):
self.anticache = False
self.anticomp = False
self.stream_large_bodies = False
self.refresh_server_playback = False
self.replacehooks = ReplaceHooks()
self.setheaders = SetHeaders()
@ -522,6 +534,12 @@ class FlowMaster(controller.Master):
self.stickycookie_state = None
self.stickycookie_txt = None
def set_stream_large_bodies(self, max_size):
if max_size is not None:
self.stream_large_bodies = StreamLargeBodies(max_size)
else:
self.stream_large_bodies = False
def set_stickyauth(self, txt):
if txt:
flt = filt.parse(txt)
@ -708,6 +726,10 @@ class FlowMaster(controller.Master):
def handle_responseheaders(self, f):
self.run_script_hook("responseheaders", f)
if self.stream_large_bodies:
self.stream_large_bodies.run(f, False)
f.reply()
return f

View File

@ -924,7 +924,9 @@ class HTTPHandler(ProtocolHandler, TemporaryServerChangeMixin):
self.c.channel.ask("responseheaders", flow)
# now get the rest of the request body, if body still needs to be read but not streaming this response
if not flow.response.stream and flow.response.content is None:
if flow.response.stream:
flow.response.content = CONTENT_MISSING
else:
flow.response.content = http.read_http_body(self.c.server_conn.rfile, flow.response.headers,
self.c.config.body_size_limit,
flow.request.method, flow.response.code, False)
@ -937,20 +939,19 @@ class HTTPHandler(ProtocolHandler, TemporaryServerChangeMixin):
if response_reply is None or response_reply == KILL:
return False
if flow.response.content is not None:
# if not streaming or there is no body to be read, we'll already have the body, just send it
if not flow.response.stream:
# no streaming:
# we already received the full response from the server and can send it to the client straight away.
self.c.client_conn.send(flow.response._assemble())
else:
# if streaming, we still need to read the body and stream its bits back to the client
# start with head
# streaming:
# First send the body and then transfer the response incrementally:
h = flow.response._assemble_head(preserve_transfer_encoding=True)
self.c.client_conn.send(h)
for chunk in http.read_http_body_chunked(self.c.server_conn.rfile,
flow.response.headers,
self.c.config.body_size_limit, "GET", 200, False, 4096):
self.c.config.body_size_limit, flow.request.method,
flow.response.code, False, 4096):
for part in chunk:
self.c.client_conn.wfile.write(part)
self.c.client_conn.wfile.flush()