Grasscutters/mitmproxy
3
0
Fork 0
mirror of https://github.com/Grasscutters/mitmproxy.git synced 2024-11-29 19:08:44 +00:00
Code Issues Projects Releases Wiki Activity

add support for multiple scripts and script arguments. refs #76

Browse Source
This commit is contained in:
Maximilian Hils 2013-06-13 16:04:04 +02:00
parent d3beaa7382
commit 2b4af8d475
12 changed files with 101 additions and 63 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
doc-src/scripting/inlinescripts.html
View File

@ -25,7 +25,7 @@ The new header will be added to all responses passing through the proxy.
## Events
### start(ScriptContext)
### start(ScriptContext, argv)
Called once on startup, before any other events.

3
examples/stub.py
View File

@ -1,8 +1,7 @@
"""
This is a script stub, with definitions for all events.
"""
def start(ctx):
def start(ctx, argv):
"""
Called once on script startup, before any other events.
"""

9
libmproxy/cmdline.py
View File

@ -16,6 +16,8 @@
import proxy
import re, filt
import argparse
import shlex
import os
class ParseException(Exception): pass
class OptionException(Exception): pass
@ -151,7 +153,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,
@ -209,8 +211,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",

2
libmproxy/console/__init__.py
View File

@ -460,7 +460,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

10
libmproxy/dump.py
View File

@ -36,7 +36,7 @@ class Options(object):
"rheaders",
"setheaders",
"server_replay",
"script",
"scripts",
"showhost",
"stickycookie",
"stickyauth",
@ -121,8 +121,8 @@ class DumpMaster(flow.FlowMaster):
not options.keepserving
)
if options.script:
err = self.load_script(options.script)
for script_argv in options.scripts:
err = self.load_script(script_argv)
if err:
raise DumpError(err)
@ -230,8 +230,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)

49
libmproxy/flow.py
View File

@ -1357,7 +1357,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
@ -1381,36 +1381,42 @@ 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
else:
r = self.get_script(path)
r = self.get_script(script_argv)
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:
@ -1561,13 +1567,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()
@ -1609,8 +1608,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:

20
libmproxy/script.py
View File

@ -23,12 +23,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):
@ -38,17 +38,21 @@ 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)
self.ns = ns
self.run("start", self.argv)
except Exception, v:
raise ScriptError(traceback.format_exc(v))
self.ns = ns
def unload(self):
return self.run("done")
def run(self, name, *args, **kwargs):
"""

8
test/scripts/a.py
View File

@ -1,5 +1,13 @@
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--var')
var = 0
def __init__(ctx, argv):
global var
var = parser.parse_args(argv).var
def here(ctx):
global var
var += 1

2
test/scripts/starterr.py
View File

@ -1,3 +1,3 @@
def start(ctx):
def start(ctx, argv):
raise ValueError

13
test/test_cmdline.py
View File

@ -1,6 +1,7 @@
import argparse
from libmproxy import cmdline
import tutils
import os.path
def test_parse_replace_hook():
@ -39,6 +40,18 @@ def test_parse_setheaders():
x = cmdline.parse_setheader("/foo/bar/voing")
assert x == ("foo", "bar", "voing")
def test_shlex():
"""
shlex.split assumes posix=True by default, we do manual detection for windows.
Test whether script paths are parsed correctly
"""
absfilepath = os.path.normcase(os.path.abspath(__file__))
parser = argparse.ArgumentParser()
cmdline.common_options(parser)
opts = parser.parse_args(args=["-s",absfilepath])
assert os.path.isfile(opts.scripts[0][0])
def test_common():
parser = argparse.ArgumentParser()

21
test/test_flow.py
View File

@ -544,9 +544,11 @@ class TestFlowMaster:
fm = flow.FlowMaster(None, s)
assert not fm.load_script(tutils.test_data.path("scripts/a.py"))
assert not fm.load_script(tutils.test_data.path("scripts/a.py"))
assert not fm.load_script(None)
assert not fm.unload_script(fm.scripts[0])
assert not fm.unload_script(fm.scripts[0])
assert fm.load_script("nonexistent")
assert "ValueError" in fm.load_script(tutils.test_data.path("scripts/starterr.py"))
assert len(fm.scripts) == 0
def test_replay(self):
s = flow.State()
@ -572,20 +574,27 @@ class TestFlowMaster:
assert not fm.load_script(tutils.test_data.path("scripts/all.py"))
req = tutils.treq()
fm.handle_clientconnect(req.client_conn)
assert fm.script.ns["log"][-1] == "clientconnect"
assert fm.scripts[0].ns["log"][-1] == "clientconnect"
f = fm.handle_request(req)
assert fm.script.ns["log"][-1] == "request"
assert fm.scripts[0].ns["log"][-1] == "request"
resp = tutils.tresp(req)
fm.handle_response(resp)
assert fm.script.ns["log"][-1] == "response"
assert fm.scripts[0].ns["log"][-1] == "response"
#load second script
assert not fm.load_script(tutils.test_data.path("scripts/all.py"))
assert len(fm.scripts) == 2
dc = flow.ClientDisconnect(req.client_conn)
dc.reply = controller.DummyReply()
fm.handle_clientdisconnect(dc)
assert fm.script.ns["log"][-1] == "clientdisconnect"
assert fm.scripts[0].ns["log"][-1] == "clientdisconnect"
assert fm.scripts[1].ns["log"][-1] == "clientdisconnect"
#unload first script
fm.unload_script(fm.scripts[0])
assert len(fm.scripts) == 1
err = flow.Error(f.request, "msg")
err.reply = controller.DummyReply()
fm.handle_error(err)
assert fm.script.ns["log"][-1] == "error"
assert fm.scripts[0].ns["log"][-1] == "error"
def test_duplicate_flow(self):
s = flow.State()

21
test/test_script.py
View File

@ -1,5 +1,7 @@
from libmproxy import script, flow
import tutils
import shlex
import os
class TestScript:
def test_simple(self):
@ -7,11 +9,12 @@ class TestScript:
fm = flow.FlowMaster(None, s)
ctx = flow.ScriptContext(fm)
p = script.Script(tutils.test_data.path("scripts/a.py"), ctx)
p = script.Script(shlex.split(tutils.test_data.path("scripts/a.py")+" --var 40", posix=(os.name != "nt")), ctx)
p.load()
assert "here" in p.ns
assert p.run("here") == (True, 1)
assert p.run("here") == (True, 2)
assert p.run("here") == (True, 41)
assert p.run("here") == (True, 42)
ret = p.run("errargs")
assert not ret[0]
@ -19,12 +22,12 @@ class TestScript:
# Check reload
p.load()
assert p.run("here") == (True, 1)
assert p.run("here") == (True, 41)
def test_duplicate_flow(self):
s = flow.State()
fm = flow.FlowMaster(None, s)
fm.load_script(tutils.test_data.path("scripts/duplicate_flow.py"))
fm.load_script([tutils.test_data.path("scripts/duplicate_flow.py")])
r = tutils.treq()
fm.handle_request(r)
assert fm.state.flow_count() == 2
@ -37,25 +40,25 @@ class TestScript:
ctx = flow.ScriptContext(fm)
s = script.Script("nonexistent", ctx)
s = script.Script(["nonexistent"], ctx)
tutils.raises(
"no such file",
s.load
)
s = script.Script(tutils.test_data.path("scripts"), ctx)
s = script.Script([tutils.test_data.path("scripts")], ctx)
tutils.raises(
"not a file",
s.load
)
s = script.Script(tutils.test_data.path("scripts/syntaxerr.py"), ctx)
s = script.Script([tutils.test_data.path("scripts/syntaxerr.py")], ctx)
tutils.raises(
script.ScriptError,
s.load
)
s = script.Script(tutils.test_data.path("scripts/loaderr.py"), ctx)
s = script.Script([tutils.test_data.path("scripts/loaderr.py")], ctx)
tutils.raises(
script.ScriptError,
s.load