From 5d2b7c52f9c33e84be5c4330b09b0f2a5ad869e2 Mon Sep 17 00:00:00 2001 From: Maximilian Hils Date: Fri, 8 Jul 2016 19:57:57 -0700 Subject: [PATCH] move script context to mitmproxy.ctx --- examples/har_extractor.py | 6 +++--- examples/nonblocking.py | 4 ++-- examples/proxapp.py | 4 ++-- examples/stub.py | 20 +++++++++---------- examples/tls_passthrough.py | 2 +- mitmproxy/__init__.py | 5 ----- mitmproxy/controller.py | 12 +++++------ mitmproxy/ctx.py | 4 ++++ test/mitmproxy/data/scripts/all.py | 14 ++++++------- test/mitmproxy/data/scripts/duplicate_flow.py | 4 ++-- test/mitmproxy/test_examples.py | 6 +++--- 11 files changed, 40 insertions(+), 41 deletions(-) create mode 100644 mitmproxy/ctx.py diff --git a/examples/har_extractor.py b/examples/har_extractor.py index 208a2fa8f..2a69b9afd 100644 --- a/examples/har_extractor.py +++ b/examples/har_extractor.py @@ -219,17 +219,17 @@ def done(): compressed_json_dump = context.HARLog.compress() if context.dump_file == '-': - mitmproxy.log(pprint.pformat(json.loads(json_dump))) + mitmproxy.ctx.log(pprint.pformat(json.loads(json_dump))) elif context.dump_file.endswith('.zhar'): file(context.dump_file, "w").write(compressed_json_dump) else: file(context.dump_file, "w").write(json_dump) - mitmproxy.log( + mitmproxy.ctx.log( "HAR log finished with %s bytes (%s bytes compressed)" % ( len(json_dump), len(compressed_json_dump) ) ) - mitmproxy.log( + mitmproxy.ctx.log( "Compression rate is %s%%" % str( 100. * len(compressed_json_dump) / len(json_dump) ) diff --git a/examples/nonblocking.py b/examples/nonblocking.py index 05a269210..b81478dff 100644 --- a/examples/nonblocking.py +++ b/examples/nonblocking.py @@ -5,6 +5,6 @@ from mitmproxy.script import concurrent @concurrent # Remove this and see what happens def request(flow): - mitmproxy.log("handle request: %s%s" % (flow.request.host, flow.request.path)) + mitmproxy.ctx.log("handle request: %s%s" % (flow.request.host, flow.request.path)) time.sleep(5) - mitmproxy.log("start request: %s%s" % (flow.request.host, flow.request.path)) + mitmproxy.ctx.log("start request: %s%s" % (flow.request.host, flow.request.path)) diff --git a/examples/proxapp.py b/examples/proxapp.py index ddc385448..2935b5874 100644 --- a/examples/proxapp.py +++ b/examples/proxapp.py @@ -17,9 +17,9 @@ def hello_world(): # Register the app using the magic domain "proxapp" on port 80. Requests to # this domain and port combination will now be routed to the WSGI app instance. def start(): - mitmproxy.master.apps.add(app, "proxapp", 80) + mitmproxy.ctx.master.apps.add(app, "proxapp", 80) # SSL works too, but the magic domain needs to be resolvable from the mitmproxy machine due to mitmproxy's design. # mitmproxy will connect to said domain and use serve its certificate (unless --no-upstream-cert is set) # but won't send any data. - mitmproxy.master.apps.add(app, "example.com", 443) + mitmproxy.ctx.master.apps.add(app, "example.com", 443) diff --git a/examples/stub.py b/examples/stub.py index a4f16699c..10b342831 100644 --- a/examples/stub.py +++ b/examples/stub.py @@ -8,7 +8,7 @@ def start(): """ Called once on script startup, before any other events. """ - mitmproxy.log("start") + mitmproxy.ctx.log("start") def clientconnect(root_layer): @@ -16,14 +16,14 @@ def clientconnect(root_layer): Called when a client initiates a connection to the proxy. Note that a connection can correspond to multiple HTTP requests """ - mitmproxy.log("clientconnect") + mitmproxy.ctx.log("clientconnect") def request(flow): """ Called when a client request has been received. """ - mitmproxy.log("request") + mitmproxy.ctx.log("request") def serverconnect(server_conn): @@ -31,7 +31,7 @@ def serverconnect(server_conn): Called when the proxy initiates a connection to the target server. Note that a connection can correspond to multiple HTTP requests """ - mitmproxy.log("serverconnect") + mitmproxy.ctx.log("serverconnect") def responseheaders(flow): @@ -40,14 +40,14 @@ def responseheaders(flow): but the response body has not been processed yet. Can be used to tell mitmproxy to stream the response. """ - mitmproxy.log("responseheaders") + mitmproxy.ctx.log("responseheaders") def response(flow): """ Called when a server response has been received. """ - mitmproxy.log("response") + mitmproxy.ctx.log("response") def error(flow): @@ -56,25 +56,25 @@ def error(flow): interrupted connections. This is distinct from a valid server HTTP error response, which is simply a response with an HTTP error code. """ - mitmproxy.log("error") + mitmproxy.ctx.log("error") def serverdisconnect(server_conn): """ Called when the proxy closes the connection to the target server. """ - mitmproxy.log("serverdisconnect") + mitmproxy.ctx.log("serverdisconnect") def clientdisconnect(root_layer): """ Called when a client disconnects from the proxy. """ - mitmproxy.log("clientdisconnect") + mitmproxy.ctx.log("clientdisconnect") def done(): """ Called once on script shutdown, after any other events. """ - mitmproxy.log("done") + mitmproxy.ctx.log("done") diff --git a/examples/tls_passthrough.py b/examples/tls_passthrough.py index 374020e7e..20e8f9bec 100644 --- a/examples/tls_passthrough.py +++ b/examples/tls_passthrough.py @@ -135,7 +135,7 @@ def next_layer(next_layer): next_layer.__class__ = TlsFeedback else: # We don't intercept - reply with a pass-through layer and add a "skipped" entry. - mitmproxy.log("TLS passthrough for %s" % repr(next_layer.server_conn.address), "info") + mitmproxy.ctx.log("TLS passthrough for %s" % repr(next_layer.server_conn.address), "info") next_layer_replacement = RawTCPLayer(next_layer.ctx, ignore=True) next_layer.reply.send(next_layer_replacement) tls_strategy.record_skipped(server_address) diff --git a/mitmproxy/__init__.py b/mitmproxy/__init__.py index 99ce7c9c1..e69de29bb 100644 --- a/mitmproxy/__init__.py +++ b/mitmproxy/__init__.py @@ -1,5 +0,0 @@ -from typing import Callable # noqa -from mitmproxy import flow # noqa - -master = None # type: flow.FlowMaster -log = None # type: Callable[[str], None] diff --git a/mitmproxy/controller.py b/mitmproxy/controller.py index 222ebc69a..e2be3a532 100644 --- a/mitmproxy/controller.py +++ b/mitmproxy/controller.py @@ -6,7 +6,7 @@ import contextlib from six.moves import queue -import mitmproxy +from . import ctx as mitmproxy_ctx from netlib import basethread from . import exceptions @@ -59,16 +59,16 @@ class Master(object): @contextlib.contextmanager def handlecontext(self): # Handlecontexts also have to nest - leave cleanup to the outermost - if mitmproxy.master: + if mitmproxy_ctx.master: yield return - mitmproxy.master = self - mitmproxy.log = Log(self) + mitmproxy_ctx.master = self + mitmproxy_ctx.log = Log(self) try: yield finally: - mitmproxy.master = None - mitmproxy.log = None + mitmproxy_ctx.master = None + mitmproxy_ctx.log = None def add_server(self, server): # We give a Channel to the server which can be used to communicate with the master diff --git a/mitmproxy/ctx.py b/mitmproxy/ctx.py new file mode 100644 index 000000000..fcfdfd0b0 --- /dev/null +++ b/mitmproxy/ctx.py @@ -0,0 +1,4 @@ +from typing import Callable # noqa + +master = None # type: "mitmproxy.flow.FlowMaster" +log = None # type: Callable[[str], None] diff --git a/test/mitmproxy/data/scripts/all.py b/test/mitmproxy/data/scripts/all.py index 17ffe33f0..bf8e93ecf 100644 --- a/test/mitmproxy/data/scripts/all.py +++ b/test/mitmproxy/data/scripts/all.py @@ -3,35 +3,35 @@ log = [] def clientconnect(cc): - mitmproxy.log("XCLIENTCONNECT") + mitmproxy.ctx.log("XCLIENTCONNECT") log.append("clientconnect") def serverconnect(cc): - mitmproxy.log("XSERVERCONNECT") + mitmproxy.ctx.log("XSERVERCONNECT") log.append("serverconnect") def request(f): - mitmproxy.log("XREQUEST") + mitmproxy.ctx.log("XREQUEST") log.append("request") def response(f): - mitmproxy.log("XRESPONSE") + mitmproxy.ctx.log("XRESPONSE") log.append("response") def responseheaders(f): - mitmproxy.log("XRESPONSEHEADERS") + mitmproxy.ctx.log("XRESPONSEHEADERS") log.append("responseheaders") def clientdisconnect(cc): - mitmproxy.log("XCLIENTDISCONNECT") + mitmproxy.ctx.log("XCLIENTDISCONNECT") log.append("clientdisconnect") def error(cc): - mitmproxy.log("XERROR") + mitmproxy.ctx.log("XERROR") log.append("error") diff --git a/test/mitmproxy/data/scripts/duplicate_flow.py b/test/mitmproxy/data/scripts/duplicate_flow.py index a50d2cb51..565b1845c 100644 --- a/test/mitmproxy/data/scripts/duplicate_flow.py +++ b/test/mitmproxy/data/scripts/duplicate_flow.py @@ -2,5 +2,5 @@ import mitmproxy def request(f): - f = mitmproxy.master.duplicate_flow(f) - mitmproxy.master.replay_request(f, block=True, run_scripthooks=False) + f = mitmproxy.ctx.master.duplicate_flow(f) + mitmproxy.ctx.master.replay_request(f, block=True, run_scripthooks=False) diff --git a/test/mitmproxy/test_examples.py b/test/mitmproxy/test_examples.py index 3b5ff2a26..bdadcd110 100644 --- a/test/mitmproxy/test_examples.py +++ b/test/mitmproxy/test_examples.py @@ -21,8 +21,8 @@ def example(command): yield s -@mock.patch("mitmproxy.master") -@mock.patch("mitmproxy.log") +@mock.patch("mitmproxy.ctx.master") +@mock.patch("mitmproxy.ctx.log") def test_load_scripts(log, master): scripts = glob.glob("%s/*.py" % example_dir) @@ -121,7 +121,7 @@ def test_redirect_requests(): assert flow.request.host == "mitmproxy.org" -@mock.patch("mitmproxy.log") +@mock.patch("mitmproxy.ctx.log") def test_har_extractor(log): if sys.version_info >= (3, 0): with tutils.raises("does not work on Python 3"):