diff options
Diffstat (limited to 'libmproxy')
| -rw-r--r-- | libmproxy/cmdline.py | 9 | ||||
| -rw-r--r-- | libmproxy/console/__init__.py | 2 | ||||
| -rw-r--r-- | libmproxy/dump.py | 11 | ||||
| -rw-r--r-- | libmproxy/flow.py | 53 | ||||
| -rw-r--r-- | libmproxy/proxy.py | 5 | ||||
| -rw-r--r-- | libmproxy/script.py | 20 |
6 files changed, 55 insertions, 45 deletions
diff --git a/libmproxy/cmdline.py b/libmproxy/cmdline.py index a394d7f3..53926028 100644 --- a/libmproxy/cmdline.py +++ b/libmproxy/cmdline.py @@ -1,6 +1,8 @@ import proxy import re, filt import argparse +import shlex +import os APP_DOMAIN = "mitm" APP_IP = "1.1.1.1" @@ -143,7 +145,7 @@ def get_common_options(options): replacements = reps, setheaders = setheaders, server_replay = options.server_replay, - script = options.script, + scripts = options.scripts, stickycookie = stickycookie, stickyauth = stickyauth, showhost = options.showhost, @@ -201,8 +203,9 @@ def common_options(parser): ) parser.add_argument( "-s", - action="store", dest="script", default=None, - help="Run a script." + action="append", type=lambda x: shlex.split(x,posix=(os.name != "nt")), dest="scripts", default=[], + metavar='"script.py --bar"', + help="Run a script. Surround with quotes to pass script arguments. Can be passed multiple times." ) parser.add_argument( "-t", diff --git a/libmproxy/console/__init__.py b/libmproxy/console/__init__.py index 37f5d646..04e240d3 100644 --- a/libmproxy/console/__init__.py +++ b/libmproxy/console/__init__.py @@ -459,7 +459,7 @@ class ConsoleMaster(flow.FlowMaster): self._run_script_method("response", s, f) if f.error: self._run_script_method("error", s, f) - s.run("done") + s.unload() self.refresh_flow(f) self.state.last_script = path diff --git a/libmproxy/dump.py b/libmproxy/dump.py index 7a49cee0..8b9c9813 100644 --- a/libmproxy/dump.py +++ b/libmproxy/dump.py @@ -24,7 +24,7 @@ class Options(object): "rheaders", "setheaders", "server_replay", - "script", + "scripts", "showhost", "stickycookie", "stickyauth", @@ -109,8 +109,9 @@ class DumpMaster(flow.FlowMaster): not options.keepserving ) - if options.script: - err = self.load_script(options.script) + scripts = options.scripts or [] + for script_argv in scripts: + err = self.load_script(script_argv) if err: raise DumpError(err) @@ -221,8 +222,8 @@ class DumpMaster(flow.FlowMaster): def run(self): # pragma: no cover if self.o.rfile and not self.o.keepserving: - if self.script: - self.load_script(None) + for script in self.scripts: + self.unload_script(script) return try: return flow.FlowMaster.run(self) diff --git a/libmproxy/flow.py b/libmproxy/flow.py index 84836d11..da946b0b 100644 --- a/libmproxy/flow.py +++ b/libmproxy/flow.py @@ -1349,7 +1349,7 @@ class FlowMaster(controller.Master): self.server_playback = None self.client_playback = None self.kill_nonreplay = False - self.script = None + self.scripts = [] self.pause_scripts = False self.stickycookie_state = False @@ -1385,37 +1385,43 @@ class FlowMaster(controller.Master): """ pass - def get_script(self, path): + def get_script(self, script_argv): """ Returns an (error, script) tuple. """ - s = script.Script(path, ScriptContext(self)) + s = script.Script(script_argv, ScriptContext(self)) try: s.load() except script.ScriptError, v: return (v.args[0], None) - ret = s.run("start") - if not ret[0] and ret[1]: - return ("Error in script start:\n\n" + ret[1][1], None) return (None, s) - def load_script(self, path): + def unload_script(self,script): + script.unload() + self.scripts.remove(script) + + def load_script(self, script_argv): """ Loads a script. Returns an error description if something went - wrong. If path is None, the current script is terminated. + wrong. """ - if path is None: - self.run_script_hook("done") - self.script = None + r = self.get_script(script_argv) + if r[0]: + return r[0] else: - r = self.get_script(path) - if r[0]: - return r[0] - else: - if self.script: - self.run_script_hook("done") - self.script = r[1] + self.scripts.append(r[1]) + + def run_single_script_hook(self, script, name, *args, **kwargs): + if script and not self.pause_scripts: + ret = script.run(name, *args, **kwargs) + if not ret[0] and ret[1]: + e = "Script error:\n" + ret[1][1] + self.add_event(e, "error") + def run_script_hook(self, name, *args, **kwargs): + for script in self.scripts: + self.run_single_script_hook(script, name, *args, **kwargs) + def set_stickycookie(self, txt): if txt: flt = filt.parse(txt) @@ -1565,13 +1571,6 @@ class FlowMaster(controller.Master): if block: rt.join() - def run_script_hook(self, name, *args, **kwargs): - if self.script and not self.pause_scripts: - ret = self.script.run(name, *args, **kwargs) - if not ret[0] and ret[1]: - e = "Script error:\n" + ret[1][1] - self.add_event(e, "error") - def handle_clientconnect(self, cc): self.run_script_hook("clientconnect", cc) cc.reply() @@ -1613,8 +1612,8 @@ class FlowMaster(controller.Master): return f def shutdown(self): - if self.script: - self.load_script(None) + for script in self.scripts: + self.unload_script(script) controller.Master.shutdown(self) if self.stream: for i in self.state._flow_list: diff --git a/libmproxy/proxy.py b/libmproxy/proxy.py index b0c15463..5dcf6dc1 100644 --- a/libmproxy/proxy.py +++ b/libmproxy/proxy.py @@ -467,8 +467,9 @@ class ProxyHandler(tcp.BaseHandler): self.wfile.write("Server: %s\r\n"%self.server_version) self.wfile.write("Content-type: text/html\r\n") self.wfile.write("Content-Length: %d\r\n"%len(html_content)) - for key, value in headers.items(): - self.wfile.write("%s: %s\r\n"%(key, value)) + if headers: + for key, value in headers.items(): + self.wfile.write("%s: %s\r\n"%(key, value)) self.wfile.write("Connection: close\r\n") self.wfile.write("\r\n") self.wfile.write(html_content) diff --git a/libmproxy/script.py b/libmproxy/script.py index 95fd5be2..623f2b92 100644 --- a/libmproxy/script.py +++ b/libmproxy/script.py @@ -8,12 +8,12 @@ class Script: """ The instantiator should do something along this vein: - s = Script(path, master) + s = Script(argv, master) s.load() - s.run("start") """ - def __init__(self, path, ctx): - self.path, self.ctx = path, ctx + def __init__(self, argv, ctx): + self.argv = argv + self.ctx = ctx self.ns = None def load(self): @@ -23,17 +23,23 @@ class Script: Raises ScriptError on failure, with argument equal to an error message that may be a formatted traceback. """ - path = os.path.expanduser(self.path) + path = os.path.expanduser(self.argv[0]) if not os.path.exists(path): - raise ScriptError("No such file: %s"%self.path) + raise ScriptError("No such file: %s" % path) if not os.path.isfile(path): - raise ScriptError("Not a file: %s"%self.path) + raise ScriptError("Not a file: %s" % path) ns = {} try: execfile(path, ns, ns) except Exception, v: raise ScriptError(traceback.format_exc(v)) self.ns = ns + r = self.run("start", self.argv) + if not r[0] and r[1]: + raise ScriptError(r[1][1]) + + def unload(self): + return self.run("done") def run(self, name, *args, **kwargs): """ |