From 58e1b3a47f392a5f4f16e30318820f163568f54e Mon Sep 17 00:00:00 2001 From: Aldo Cortesi Date: Sun, 12 Jan 2014 12:49:19 +1300 Subject: Start refactoring scripts - Move ScriptContext into script module - Use mock module instead of hand-rolled mock objects in tests --- libmproxy/console/__init__.py | 5 ++++- libmproxy/flow.py | 37 ++----------------------------------- libmproxy/proxy.py | 6 +++--- libmproxy/script.py | 39 ++++++++++++++++++++++++++++++++++++--- 4 files changed, 45 insertions(+), 42 deletions(-) (limited to 'libmproxy') diff --git a/libmproxy/console/__init__.py b/libmproxy/console/__init__.py index 210e2b95..b8d4a105 100644 --- a/libmproxy/console/__init__.py +++ b/libmproxy/console/__init__.py @@ -781,6 +781,9 @@ class ConsoleMaster(flow.FlowMaster): else: self.view_flowlist() + def edit_scripts(self, *args, **kwargs): + pass + def loop(self): changed = True try: @@ -883,7 +886,7 @@ class ConsoleMaster(flow.FlowMaster): grideditor.ScriptEditor( self, [[i.argv[0]] for i in self.scripts], - None + self.edit_scripts ) ) #if self.scripts: diff --git a/libmproxy/flow.py b/libmproxy/flow.py index 32306513..e88bf985 100644 --- a/libmproxy/flow.py +++ b/libmproxy/flow.py @@ -143,39 +143,6 @@ class SetHeaders: f.request.headers.add(header, value) -class ScriptContext: - def __init__(self, master): - self._master = master - - def log(self, *args, **kwargs): - """ - Logs an event. - - How this is handled depends on the front-end. mitmdump will display - events if the eventlog flag ("-e") was passed. mitmproxy sends - output to the eventlog for display ("v" keyboard shortcut). - """ - self._master.add_event(*args, **kwargs) - - def duplicate_flow(self, f): - """ - Returns a duplicate of the specified flow. The flow is also - injected into the current state, and is ready for editing, replay, - etc. - """ - self._master.pause_scripts = True - f = self._master.duplicate_flow(f) - self._master.pause_scripts = False - return f - - def replay_request(self, f): - """ - Replay the request on the current flow. The response will be added - to the flow object. - """ - self._master.replay_request(f) - - class decoded(object): """ @@ -1431,14 +1398,14 @@ class FlowMaster(controller.Master): """ Returns an (error, script) tuple. """ - s = script.Script(script_argv, ScriptContext(self)) + s = script.Script(script_argv, self) try: s.load() except script.ScriptError, v: return (v.args[0], None) return (None, s) - def unload_script(self,script): + def unload_script(self, script): script.unload() self.scripts.remove(script) diff --git a/libmproxy/proxy.py b/libmproxy/proxy.py index 15f75b8d..1894f7f0 100644 --- a/libmproxy/proxy.py +++ b/libmproxy/proxy.py @@ -110,15 +110,15 @@ class HandleSNI: self.handler, self.client_conn, self.host, self.port = handler, client_conn, host, port self.cert, self.key = cert, key - def __call__(self, connection): + def __call__(self, client_connection): try: - sn = connection.get_servername() + sn = client_connection.get_servername() if sn: self.handler.get_server_connection(self.client_conn, "https", self.host, self.port, sn) new_context = SSL.Context(SSL.TLSv1_METHOD) new_context.use_privatekey_file(self.key) new_context.use_certificate(self.cert.x509) - connection.set_context(new_context) + client_connection.set_context(new_context) self.handler.sni = sn.decode("utf8").encode("idna") # An unhandled exception in this method will core dump PyOpenSSL, so # make dang sure it doesn't happen. diff --git a/libmproxy/script.py b/libmproxy/script.py index f8a0d085..0f1056f6 100644 --- a/libmproxy/script.py +++ b/libmproxy/script.py @@ -5,6 +5,39 @@ class ScriptError(Exception): pass +class ScriptContext: + def __init__(self, master): + self._master = master + + def log(self, *args, **kwargs): + """ + Logs an event. + + How this is handled depends on the front-end. mitmdump will display + events if the eventlog flag ("-e") was passed. mitmproxy sends + output to the eventlog for display ("v" keyboard shortcut). + """ + self._master.add_event(*args, **kwargs) + + def duplicate_flow(self, f): + """ + Returns a duplicate of the specified flow. The flow is also + injected into the current state, and is ready for editing, replay, + etc. + """ + self._master.pause_scripts = True + f = self._master.duplicate_flow(f) + self._master.pause_scripts = False + return f + + def replay_request(self, f): + """ + Replay the request on the current flow. The response will be added + to the flow object. + """ + self._master.replay_request(f) + + class Script: """ The instantiator should do something along this vein: @@ -12,9 +45,9 @@ class Script: s = Script(argv, master) s.load() """ - def __init__(self, argv, ctx): + def __init__(self, argv, master): self.argv = argv - self.ctx = ctx + self.ctx = ScriptContext(master) self.ns = None def load(self): @@ -82,4 +115,4 @@ def concurrent(fn): def _concurrent(ctx, conn): _handle_concurrent_reply(fn, conn, [ctx, conn]) return _concurrent - raise NotImplementedError("Concurrent decorator not supported for this method.") \ No newline at end of file + raise NotImplementedError("Concurrent decorator not supported for this method.") -- cgit v1.2.3