diff options
| -rw-r--r-- | mitmproxy/builtins/__init__.py | 10 | ||||
| -rw-r--r-- | mitmproxy/builtins/streambodies.py | 34 | ||||
| -rw-r--r-- | mitmproxy/console/master.py | 2 | ||||
| -rw-r--r-- | mitmproxy/console/statusbar.py | 4 | ||||
| -rw-r--r-- | mitmproxy/dump.py | 1 | ||||
| -rw-r--r-- | mitmproxy/flow/__init__.py | 4 | ||||
| -rw-r--r-- | mitmproxy/flow/master.py | 15 | ||||
| -rw-r--r-- | mitmproxy/flow/modules.py | 15 | ||||
| -rw-r--r-- | mitmproxy/models/http.py | 3 | ||||
| -rw-r--r-- | test/mitmproxy/builtins/test_streambodies.py | 29 | ||||
| -rw-r--r-- | test/mitmproxy/test_server.py | 12 |
11 files changed, 76 insertions, 53 deletions
diff --git a/mitmproxy/builtins/__init__.py b/mitmproxy/builtins/__init__.py index 4bd2e56d..455f8bd8 100644 --- a/mitmproxy/builtins/__init__.py +++ b/mitmproxy/builtins/__init__.py @@ -1,13 +1,14 @@ from mitmproxy.builtins import anticache from mitmproxy.builtins import anticomp +from mitmproxy.builtins import clientplayback from mitmproxy.builtins import filestreamer -from mitmproxy.builtins import stickyauth -from mitmproxy.builtins import stickycookie -from mitmproxy.builtins import script from mitmproxy.builtins import replace +from mitmproxy.builtins import script from mitmproxy.builtins import setheaders from mitmproxy.builtins import serverplayback -from mitmproxy.builtins import clientplayback +from mitmproxy.builtins import stickyauth +from mitmproxy.builtins import stickycookie +from mitmproxy.builtins import streambodies def default_addons(): @@ -18,6 +19,7 @@ def default_addons(): stickycookie.StickyCookie(), script.ScriptLoader(), filestreamer.FileStreamer(), + streambodies.StreamBodies(), replace.Replace(), setheaders.SetHeaders(), serverplayback.ServerPlayback(), diff --git a/mitmproxy/builtins/streambodies.py b/mitmproxy/builtins/streambodies.py new file mode 100644 index 00000000..b3e5d6b2 --- /dev/null +++ b/mitmproxy/builtins/streambodies.py @@ -0,0 +1,34 @@ +from netlib.http import http1 +from netlib import exceptions +from mitmproxy import ctx + + +class StreamBodies: + def __init__(self): + self.max_size = None + + def configure(self, options, updated): + self.max_size = options.stream_large_bodies + + def run(self, f, is_request): + if self.max_size: + r = f.request if is_request else f.response + try: + expected_size = http1.expected_http_body_size( + f.request, f.response if not is_request else None + ) + except exceptions.HTTPException: + f.reply.kill() + return + if expected_size and not r.raw_content and not (0 <= expected_size <= self.max_size): + # r.stream may already be a callable, which we want to preserve. + r.stream = r.stream or True + # FIXME: make message generic when we add rquest streaming + ctx.log.info("Streaming response from %s" % f.request.host) + + # FIXME! Request streaming doesn't work at the moment. + def requestheaders(self, f): + self.run(f, True) + + def responseheaders(self, f): + self.run(f, False) diff --git a/mitmproxy/console/master.py b/mitmproxy/console/master.py index 148f23bc..1a413990 100644 --- a/mitmproxy/console/master.py +++ b/mitmproxy/console/master.py @@ -236,8 +236,6 @@ class ConsoleMaster(flow.FlowMaster): if options.filter: self.set_view_filter(options.filter) - self.set_stream_large_bodies(options.stream_large_bodies) - self.palette = options.palette self.palette_transparent = options.palette_transparent diff --git a/mitmproxy/console/statusbar.py b/mitmproxy/console/statusbar.py index e281af3b..e7944a9e 100644 --- a/mitmproxy/console/statusbar.py +++ b/mitmproxy/console/statusbar.py @@ -196,10 +196,10 @@ class StatusBar(urwid.WidgetWrap): opts.append("no-upstream-cert") if self.master.state.follow_focus: opts.append("following") - if self.master.stream_large_bodies: + if self.master.options.stream_large_bodies: opts.append( "stream:%s" % human.pretty_size( - self.master.stream_large_bodies.max_size + self.master.options.stream_large_bodies ) ) diff --git a/mitmproxy/dump.py b/mitmproxy/dump.py index 35f6e1d3..d3f894fa 100644 --- a/mitmproxy/dump.py +++ b/mitmproxy/dump.py @@ -40,7 +40,6 @@ class DumpMaster(flow.FlowMaster): self.addons.add(dumper.Dumper()) # This line is just for type hinting self.options = self.options # type: Options - self.set_stream_large_bodies(options.stream_large_bodies) if not self.options.no_server and server: self.add_log( diff --git a/mitmproxy/flow/__init__.py b/mitmproxy/flow/__init__.py index 9c1cb477..9d243f49 100644 --- a/mitmproxy/flow/__init__.py +++ b/mitmproxy/flow/__init__.py @@ -2,12 +2,12 @@ from mitmproxy.flow import export, modules from mitmproxy.flow.io import FlowWriter, FilteredFlowWriter, FlowReader, read_flows_from_paths from mitmproxy.flow.master import FlowMaster from mitmproxy.flow.modules import ( - AppRegistry, StreamLargeBodies + AppRegistry ) from mitmproxy.flow.state import State, DummyState, FlowView __all__ = [ "export", "modules", "FlowWriter", "FilteredFlowWriter", "FlowReader", "read_flows_from_paths", - "FlowMaster", "AppRegistry", "StreamLargeBodies", "DummyState", "State", "FlowView", + "FlowMaster", "AppRegistry", "DummyState", "State", "FlowView", ] diff --git a/mitmproxy/flow/master.py b/mitmproxy/flow/master.py index 14ba39e2..02e5246f 100644 --- a/mitmproxy/flow/master.py +++ b/mitmproxy/flow/master.py @@ -3,7 +3,6 @@ import sys from typing import Optional # noqa -import netlib.exceptions from netlib import http from mitmproxy import controller from mitmproxy import exceptions @@ -53,18 +52,11 @@ class FlowMaster(controller.Master): if server: self.add_server(server) self.state = state - self.stream_large_bodies = None # type: Optional[modules.StreamLargeBodies] self.apps = modules.AppRegistry() def start_app(self, host, port): self.apps.add(app.mapp, host, port) - def set_stream_large_bodies(self, max_size): - if max_size is not None: - self.stream_large_bodies = modules.StreamLargeBodies(max_size) - else: - self.stream_large_bodies = False - def duplicate_flow(self, f): """ Duplicate flow, and insert it into state without triggering any of @@ -238,12 +230,7 @@ class FlowMaster(controller.Master): @controller.handler def responseheaders(self, f): - try: - if self.stream_large_bodies: - self.stream_large_bodies.run(f, False) - except netlib.exceptions.HttpException: - f.reply.kill() - return + pass @controller.handler def response(self, f): diff --git a/mitmproxy/flow/modules.py b/mitmproxy/flow/modules.py index ec86ba9f..4bc4d08b 100644 --- a/mitmproxy/flow/modules.py +++ b/mitmproxy/flow/modules.py @@ -1,6 +1,5 @@ from netlib import wsgi from netlib import version -from netlib.http import http1 class AppRegistry: @@ -28,17 +27,3 @@ class AppRegistry: if "host" in request.headers: host = request.headers["host"] return self.apps.get((host, request.port), None) - - -class StreamLargeBodies: - 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 - expected_size = http1.expected_http_body_size( - flow.request, flow.response if not is_request else None - ) - if not r.raw_content and not (0 <= expected_size <= self.max_size): - # r.stream may already be a callable, which we want to preserve. - r.stream = r.stream or True diff --git a/mitmproxy/models/http.py b/mitmproxy/models/http.py index 1d7f9eb8..91263b95 100644 --- a/mitmproxy/models/http.py +++ b/mitmproxy/models/http.py @@ -53,6 +53,7 @@ class HTTPRequest(http.Request): # Is this request replayed? self.is_replay = is_replay + self.stream = None def get_state(self): state = super().get_state() @@ -125,7 +126,7 @@ class HTTPResponse(http.Response): # Is this request replayed? self.is_replay = is_replay - self.stream = False + self.stream = None @classmethod def wrap(self, response): diff --git a/test/mitmproxy/builtins/test_streambodies.py b/test/mitmproxy/builtins/test_streambodies.py new file mode 100644 index 00000000..0e8a82f2 --- /dev/null +++ b/test/mitmproxy/builtins/test_streambodies.py @@ -0,0 +1,29 @@ +from .. import tutils, mastertest +from mitmproxy.flow import state +from mitmproxy.flow import master +from mitmproxy import options + +from mitmproxy.builtins import streambodies + + +class TestStreamBodies(mastertest.MasterTest): + def test_simple(self): + s = state.DummyState() + o = options.Options(stream_large_bodies = 10) + m = master.FlowMaster(o, None, s) + sa = streambodies.StreamBodies() + m.addons.add(sa) + + f = tutils.tflow() + f.request.content = b"" + f.request.headers["Content-Length"] = "1024" + assert not f.request.stream + m.requestheaders(f) + assert f.request.stream + + f = tutils.tflow(resp=True) + f.response.content = b"" + f.response.headers["Content-Length"] = "1024" + assert not f.response.stream + m.responseheaders(f) + assert f.response.stream diff --git a/test/mitmproxy/test_server.py b/test/mitmproxy/test_server.py index a175d193..e9db210e 100644 --- a/test/mitmproxy/test_server.py +++ b/test/mitmproxy/test_server.py @@ -288,18 +288,6 @@ class TestHTTP(tservers.HTTPProxyTest, CommonMixin, AppMixin): resp = p.request("get:'http://foo':h':foo'='bar'") assert resp.status_code == 400 - def test_stream(self): - self.master.set_stream_large_bodies(1024 * 2) - - self.pathod("200:b@1k") - assert not self.master.state.view[-1].response.stream - assert len(self.master.state.view[-1].response.content) == 1024 * 1 - - self.pathod("200:b@3k") - assert self.master.state.view[-1].response.stream - assert self.master.state.view[-1].response.content is None - self.master.set_stream_large_bodies(None) - def test_stream_modify(self): s = script.Script( tutils.test_data.path("data/addonscripts/stream_modify.py") |