diff --git a/mitmproxy/addons/__init__.py b/mitmproxy/addons/__init__.py index 97fa2dcd6..1bf89bbba 100644 --- a/mitmproxy/addons/__init__.py +++ b/mitmproxy/addons/__init__.py @@ -3,6 +3,7 @@ from mitmproxy.addons import anticomp from mitmproxy.addons import check_alpn from mitmproxy.addons import check_ca from mitmproxy.addons import clientplayback +from mitmproxy.addons import core from mitmproxy.addons import disable_h2c_upgrade from mitmproxy.addons import onboarding from mitmproxy.addons import proxyauth @@ -19,6 +20,7 @@ from mitmproxy.addons import upstream_auth def default_addons(): return [ + core.Core(), anticache.AntiCache(), anticomp.AntiComp(), check_alpn.CheckALPN(), diff --git a/mitmproxy/addons/core.py b/mitmproxy/addons/core.py new file mode 100644 index 000000000..5d2cf57b0 --- /dev/null +++ b/mitmproxy/addons/core.py @@ -0,0 +1,20 @@ +""" + The core addon is responsible for verifying core settings that are not + checked by other addons. +""" +from mitmproxy import exceptions +from mitmproxy.utils import human + + +class Core: + def configure(self, options, updated): + if "body_size_limit" in updated and options.body_size_limit: + try: + options._processed["body_size_limit"] = human.parse_size( + options.body_size_limit + ) + except ValueError as e: + raise exceptions.OptionsError( + "Invalid body size limit specification: %s" % + options.body_size_limit + ) diff --git a/mitmproxy/addons/streambodies.py b/mitmproxy/addons/streambodies.py index 3c2a153be..a10bdb934 100644 --- a/mitmproxy/addons/streambodies.py +++ b/mitmproxy/addons/streambodies.py @@ -1,6 +1,7 @@ from mitmproxy.net.http import http1 from mitmproxy import exceptions from mitmproxy import ctx +from mitmproxy.utils import human class StreamBodies: @@ -8,7 +9,11 @@ class StreamBodies: self.max_size = None def configure(self, options, updated): - self.max_size = options.stream_large_bodies + if "stream_large_bodies" in updated and options.stream_large_bodies: + try: + self.max_size = human.parse_size(options.stream_large_bodies) + except ValueError as e: + raise exceptions.OptionsError(e) def run(self, f, is_request): if self.max_size: diff --git a/mitmproxy/options.py b/mitmproxy/options.py index c91e2f88a..58907ef32 100644 --- a/mitmproxy/options.py +++ b/mitmproxy/options.py @@ -77,7 +77,14 @@ class Options(optmanager.OptManager): self.add_option("server_replay", [], Sequence[str]) self.add_option("stickycookie", None, Optional[str]) self.add_option("stickyauth", None, Optional[str]) - self.add_option("stream_large_bodies", None, Optional[int]) + self.add_option( + "stream_large_bodies", None, Optional[str], + """ + 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. + """ + ) self.add_option( "verbosity", 2, int, "Log verbosity." @@ -109,7 +116,11 @@ class Options(optmanager.OptManager): "Add all certificates of the upstream server to the certificate chain " "that will be served to the proxy client, as extras." ) - self.add_option("body_size_limit", None, Optional[int]) + self.add_option( + "body_size_limit", None, Optional[str], + "Byte size limit of HTTP request and response bodies." + " Understands k/m/g suffixes, i.e. 3m for 3 megabytes." + ) self.add_option("cadir", CA_DIR, str) self.add_option("certs", [], Sequence[Tuple[str, str]]) self.add_option("ciphers_client", DEFAULT_CLIENT_CIPHERS, str) diff --git a/mitmproxy/optmanager.py b/mitmproxy/optmanager.py index c16e29997..b9aaf6c86 100644 --- a/mitmproxy/optmanager.py +++ b/mitmproxy/optmanager.py @@ -12,7 +12,6 @@ import ruamel.yaml from mitmproxy import exceptions from mitmproxy.utils import typecheck - """ The base implementation for Options. """ @@ -92,6 +91,7 @@ class OptManager: self.__dict__["_options"] = {} self.__dict__["changed"] = blinker.Signal() self.__dict__["errored"] = blinker.Signal() + self.__dict__["_processed"] = {} def add_option( self, @@ -330,7 +330,7 @@ class OptManager: help=o.help ) parser.set_defaults(**{option: o.default}) - elif o.typespec == int: + elif o.typespec in (int, typing.Optional[int]): parser.add_argument( "--%s" % f, action="store", @@ -339,5 +339,14 @@ class OptManager: help=o.help, metavar=metavar ) + elif o.typespec in (str, typing.Optional[str]): + parser.add_argument( + "--%s" % f, + action="store", + type=str, + dest=option, + help=o.help, + metavar=metavar + ) else: raise ValueError("Unsupported option type: %s", o.typespec) diff --git a/mitmproxy/proxy/protocol/http1.py b/mitmproxy/proxy/protocol/http1.py index b1fd0ecd5..cafc26827 100644 --- a/mitmproxy/proxy/protocol/http1.py +++ b/mitmproxy/proxy/protocol/http1.py @@ -19,7 +19,7 @@ class Http1Layer(httpbase._HttpTransmissionLayer): return http1.read_body( self.client_conn.rfile, expected_size, - self.config.options.body_size_limit + self.config.options._processed.get("body_size_limit") ) def send_request(self, request): @@ -35,7 +35,7 @@ class Http1Layer(httpbase._HttpTransmissionLayer): return http1.read_body( self.server_conn.rfile, expected_size, - self.config.options.body_size_limit + self.config.options._processed.get("body_size_limit") ) def send_response_headers(self, response): diff --git a/mitmproxy/proxy/protocol/http2.py b/mitmproxy/proxy/protocol/http2.py index 014067980..a6e8a4dd9 100644 --- a/mitmproxy/proxy/protocol/http2.py +++ b/mitmproxy/proxy/protocol/http2.py @@ -183,7 +183,7 @@ class Http2Layer(base.Layer): return True def _handle_data_received(self, eid, event, source_conn): - bsl = self.config.options.body_size_limit + bsl = self.config.options._processed.get("body_size_limit") if bsl and self.streams[eid].queued_data_length > bsl: self.streams[eid].kill() self.connections[source_conn].safe_reset_stream( diff --git a/mitmproxy/proxy/protocol/http_replay.py b/mitmproxy/proxy/protocol/http_replay.py index 38f511c4e..c47f2e8fc 100644 --- a/mitmproxy/proxy/protocol/http_replay.py +++ b/mitmproxy/proxy/protocol/http_replay.py @@ -31,6 +31,7 @@ class RequestReplayThread(basethread.BaseThread): def run(self): r = self.f.request + bsl = self.config.options._processed.get("body_size_limit") first_line_format_backup = r.first_line_format server = None try: @@ -55,7 +56,7 @@ class RequestReplayThread(basethread.BaseThread): resp = http1.read_response( server.rfile, connect_request, - body_size_limit=self.config.options.body_size_limit + body_size_limit=bsl ) if resp.status_code != 200: raise exceptions.ReplayException("Upstream server refuses CONNECT request") @@ -87,7 +88,7 @@ class RequestReplayThread(basethread.BaseThread): http1.read_response( server.rfile, r, - body_size_limit=self.config.options.body_size_limit + body_size_limit=bsl ) ) if self.channel: diff --git a/mitmproxy/tools/cmdline.py b/mitmproxy/tools/cmdline.py index a2c7e2b82..def2ae31c 100644 --- a/mitmproxy/tools/cmdline.py +++ b/mitmproxy/tools/cmdline.py @@ -4,7 +4,6 @@ import os from mitmproxy import exceptions from mitmproxy import options from mitmproxy import platform -from mitmproxy.utils import human from mitmproxy.net import tcp from mitmproxy import version from mitmproxy.addons import view @@ -25,10 +24,6 @@ def get_common_options(args): if args.stickyauth_filt: stickyauth = args.stickyauth_filt - stream_large_bodies = args.stream_large_bodies - if stream_large_bodies: - stream_large_bodies = human.parse_size(stream_large_bodies) - if args.streamfile and args.streamfile[0] == args.rfile: if args.streamfile[1] == "wb": raise exceptions.OptionsError( @@ -49,15 +44,6 @@ def get_common_options(args): parts = ["*", parts[0]] certs.append(parts) - body_size_limit = args.body_size_limit - if body_size_limit: - try: - body_size_limit = human.parse_size(body_size_limit) - except ValueError as e: - raise exceptions.OptionsError( - "Invalid body size limit specification: %s" % body_size_limit - ) - # Establish proxy mode c = 0 mode, upstream_server = "regular", None @@ -117,7 +103,7 @@ def get_common_options(args): scripts=args.scripts, stickycookie=stickycookie, stickyauth=stickyauth, - stream_large_bodies=stream_large_bodies, + stream_large_bodies=args.stream_large_bodies, showhost=args.showhost, streamfile=args.streamfile[0] if args.streamfile else None, streamfile_append=True if args.streamfile and args.streamfile[1] == "a" else False, @@ -132,7 +118,7 @@ def get_common_options(args): auth_singleuser = args.auth_singleuser, auth_htpasswd = args.auth_htpasswd, add_upstream_certs_to_client_chain = args.add_upstream_certs_to_client_chain, - body_size_limit = body_size_limit, + body_size_limit = args.body_size_limit, cadir = args.cadir, certs = certs, ciphers_client = args.ciphers_client, @@ -229,23 +215,8 @@ def basic_options(parser, opts): help="Append flows to file." ) opts.make_parser(parser, "anticomp") - parser.add_argument( - "-Z", "--body-size-limit", - action="store", dest="body_size_limit", - metavar="SIZE", - 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", - 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. - """ - ) + opts.make_parser(parser, "body_size_limit", metavar="SIZE") + opts.make_parser(parser, "stream_large_bodies") def proxy_modes(parser, opts): diff --git a/test/mitmproxy/addons/test_core.py b/test/mitmproxy/addons/test_core.py new file mode 100644 index 000000000..95272716a --- /dev/null +++ b/test/mitmproxy/addons/test_core.py @@ -0,0 +1,13 @@ +from mitmproxy import exceptions +from mitmproxy.addons import core +from mitmproxy.test import taddons +import pytest + + +def test_simple(): + sa = core.Core() + with taddons.context() as tctx: + with pytest.raises(exceptions.OptionsError): + tctx.configure(sa, body_size_limit = "invalid") + tctx.configure(sa, body_size_limit = "1m") + assert tctx.options._processed["body_size_limit"] diff --git a/test/mitmproxy/addons/test_streambodies.py b/test/mitmproxy/addons/test_streambodies.py index b982c66db..c6ce5e81c 100644 --- a/test/mitmproxy/addons/test_streambodies.py +++ b/test/mitmproxy/addons/test_streambodies.py @@ -1,13 +1,16 @@ +from mitmproxy import exceptions from mitmproxy.test import tflow from mitmproxy.test import taddons - from mitmproxy.addons import streambodies +import pytest def test_simple(): sa = streambodies.StreamBodies() with taddons.context() as tctx: - tctx.configure(sa, stream_large_bodies = 10) + with pytest.raises(exceptions.OptionsError): + tctx.configure(sa, stream_large_bodies = "invalid") + tctx.configure(sa, stream_large_bodies = "10") f = tflow.tflow() f.request.content = b"" diff --git a/test/mitmproxy/proxy/protocol/test_http2.py b/test/mitmproxy/proxy/protocol/test_http2.py index 770c65501..1f695cc5f 100644 --- a/test/mitmproxy/proxy/protocol/test_http2.py +++ b/test/mitmproxy/proxy/protocol/test_http2.py @@ -499,7 +499,8 @@ class TestBodySizeLimit(_Http2Test): return True def test_body_size_limit(self): - self.config.options.body_size_limit = 20 + self.config.options.body_size_limit = "20" + self.config.options._processed["body_size_limit"] = 20 client, h2_conn = self._setup_connection()