Grasscutters/mitmproxy
3
0
Fork 0
mirror of https://github.com/Grasscutters/mitmproxy.git synced 2024-11-23 00:01:36 +00:00
Code Issues Projects Releases Wiki Activity

improve inline scripts, fix #683, fix #684

Browse Source
This commit is contained in:
Maximilian Hils 2015-07-24 02:57:56 +02:00
parent 9b3fe80697
commit 7bf8088d80
7 changed files with 155 additions and 138 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

25
libmproxy/flow.py
View File

@ -659,9 +659,12 @@ class FlowMaster(controller.Master):
for s in self.scripts[:]:
self.unload_script(s)
def unload_script(self, script):
script.unload()
self.scripts.remove(script)
def unload_script(self, script_obj):
try:
script_obj.unload()
except script.ScriptError as e:
self.add_event("Script error:\n" + str(e), "error")
self.scripts.remove(script_obj)
def load_script(self, command):
"""
@ -674,16 +677,16 @@ class FlowMaster(controller.Master):
return v.args[0]
self.scripts.append(s)
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_single_script_hook(self, script_obj, name, *args, **kwargs):
if script_obj and not self.pause_scripts:
try:
script_obj.run(name, *args, **kwargs)
except script.ScriptError as e:
self.add_event("Script error:\n" + str(e), "error")
def run_script_hook(self, name, *args, **kwargs):
for script in self.scripts:
self.run_single_script_hook(script, name, *args, **kwargs)
for script_obj in self.scripts:
self._run_single_script_hook(script_obj, name, *args, **kwargs)
def get_ignore_filter(self):
return self.server.config.check_ignore.patterns

61
libmproxy/script.py
View File

@ -3,7 +3,7 @@ import os
import traceback
import threading
import shlex
from . import controller
import sys
class ScriptError(Exception):
@ -55,21 +55,17 @@ class ScriptContext:
class Script:
"""
The instantiator should do something along this vein:
s = Script(argv, master)
s.load()
Script object representing an inline script.
"""
def __init__(self, command, master):
self.command = command
self.argv = self.parse_command(command)
self.args = self.parse_command(command)
self.ctx = ScriptContext(master)
self.ns = None
self.load()
@classmethod
def parse_command(klass, command):
def parse_command(cls, command):
if not command or not command.strip():
raise ScriptError("Empty script command.")
if os.name == "nt": # Windows: escape all backslashes in the path.
@ -89,42 +85,52 @@ class Script:
def load(self):
"""
Loads a module.
Loads an inline script.
Raises ScriptError on failure, with argument equal to an error
message that may be a formatted traceback.
Returns:
The return value of self.run("start", ...)
Raises:
ScriptError on failure
"""
if self.ns is not None:
self.unload()
ns = {}
script_dir = os.path.dirname(os.path.abspath(self.args[0]))
sys.path.append(script_dir)
try:
execfile(self.argv[0], ns, ns)
except Exception as v:
raise ScriptError(traceback.format_exc(v))
execfile(self.args[0], ns, ns)
except Exception as e:
# Python 3: use exception chaining, https://www.python.org/dev/peps/pep-3134/
raise ScriptError(traceback.format_exc(e))
sys.path.pop()
self.ns = ns
r = self.run("start", self.argv)
if not r[0] and r[1]:
raise ScriptError(r[1][1])
return self.run("start", self.args)
def unload(self):
return self.run("done")
ret = self.run("done")
self.ns = None
return ret
def run(self, name, *args, **kwargs):
"""
Runs a plugin method.
Runs an inline script hook.
Returns:
The return value of the method.
None, if the script does not provide the method.
(True, retval) on success.
(False, None) on nonexistent method.
(False, (exc, traceback string)) if there was an exception.
Raises:
ScriptError if there was an exception.
"""
f = self.ns.get(name)
if f:
try:
return (True, f(self.ctx, *args, **kwargs))
except Exception as v:
return (False, (v, traceback.format_exc(v)))
return f(self.ctx, *args, **kwargs)
except Exception as e:
raise ScriptError(traceback.format_exc(e))
else:
return (False, None)
return None
class ReplyProxy(object):
@ -176,6 +182,7 @@ def concurrent(fn):
"clientdisconnect"):
def _concurrent(ctx, obj):
_handle_concurrent_reply(fn, obj, ctx, obj)
return _concurrent
raise NotImplementedError(
"Concurrent decorator not supported for this method.")
"Concurrent decorator not supported for '%s' method." % fn.func_name)

5
test/scripts/a.py
View File

@ -1,7 +1,4 @@
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--var', type=int)
from a_helper import parser
var = 0

4
test/scripts/a_helper.py Normal file
View File

@ -0,0 +1,4 @@
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--var', type=int)

2
test/scripts/unloaderr.py Normal file
View File

@ -0,0 +1,2 @@
def done(ctx):
raise RuntimeError()

2
test/test_examples.py
View File

@ -22,7 +22,7 @@ def test_load_scripts():
try:
s = script.Script(f, tmaster) # Loads the script file.
except Exception as v:
if not "ImportError" in str(v):
if "ImportError" not in str(v):
raise
else:
s.unload()

194
test/test_script.py
View File

@ -1,120 +1,124 @@
from libmproxy import script, flow
import tutils
import shlex
import os
import time
import mock
from libmproxy import script, flow
import tutils
class TestScript:
def test_simple(self):
s = flow.State()
fm = flow.FlowMaster(None, s)
sp = tutils.test_data.path("scripts/a.py")
p = script.Script("%s --var 40" % sp, fm)
def test_simple():
s = flow.State()
fm = flow.FlowMaster(None, s)
sp = tutils.test_data.path("scripts/a.py")
p = script.Script("%s --var 40" % sp, fm)
assert "here" in p.ns
assert p.run("here") == (True, 41)
assert p.run("here") == (True, 42)
assert "here" in p.ns
assert p.run("here") == 41
assert p.run("here") == 42
ret = p.run("errargs")
assert not ret[0]
assert len(ret[1]) == 2
tutils.raises(script.ScriptError, p.run, "errargs")
# Check reload
p.load()
assert p.run("here") == (True, 41)
# Check reload
p.load()
assert p.run("here") == 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"))
f = tutils.tflow()
fm.handle_request(f)
assert fm.state.flow_count() == 2
assert not fm.state.view[0].request.is_replay
assert fm.state.view[1].request.is_replay
def test_err(self):
s = flow.State()
fm = flow.FlowMaster(None, s)
def test_duplicate_flow():
s = flow.State()
fm = flow.FlowMaster(None, s)
fm.load_script(tutils.test_data.path("scripts/duplicate_flow.py"))
f = tutils.tflow()
fm.handle_request(f)
assert fm.state.flow_count() == 2
assert not fm.state.view[0].request.is_replay
assert fm.state.view[1].request.is_replay
tutils.raises(
"not found",
script.Script, "nonexistent", fm
)
tutils.raises(
"not a file",
script.Script, tutils.test_data.path("scripts"), fm
)
def test_err():
s = flow.State()
fm = flow.FlowMaster(None, s)
tutils.raises(
script.ScriptError,
script.Script, tutils.test_data.path("scripts/syntaxerr.py"), fm
)
tutils.raises(
"not found",
script.Script, "nonexistent", fm
)
tutils.raises(
script.ScriptError,
script.Script, tutils.test_data.path("scripts/loaderr.py"), fm
)
tutils.raises(
"not a file",
script.Script, tutils.test_data.path("scripts"), fm
)
def test_concurrent(self):
s = flow.State()
fm = flow.FlowMaster(None, s)
fm.load_script(tutils.test_data.path("scripts/concurrent_decorator.py"))
tutils.raises(
script.ScriptError,
script.Script, tutils.test_data.path("scripts/syntaxerr.py"), fm
)
with mock.patch("libmproxy.controller.DummyReply.__call__") as m:
f1, f2 = tutils.tflow(), tutils.tflow()
t_start = time.time()
fm.handle_request(f1)
f1.reply()
fm.handle_request(f2)
f2.reply()
tutils.raises(
script.ScriptError,
script.Script, tutils.test_data.path("scripts/loaderr.py"), fm
)
# Two instantiations
assert m.call_count == 0 # No calls yet.
assert (time.time() - t_start) < 0.09
scr = script.Script(tutils.test_data.path("scripts/unloaderr.py"), fm)
tutils.raises(script.ScriptError, scr.unload)
def test_concurrent2(self):
s = flow.State()
fm = flow.FlowMaster(None, s)
s = script.Script(
tutils.test_data.path("scripts/concurrent_decorator.py"),
fm)
s.load()
m = mock.Mock()
class Dummy:
def __init__(self):
self.response = self
self.error = self
self.reply = m
def test_concurrent():
s = flow.State()
fm = flow.FlowMaster(None, s)
fm.load_script(tutils.test_data.path("scripts/concurrent_decorator.py"))
with mock.patch("libmproxy.controller.DummyReply.__call__") as m:
f1, f2 = tutils.tflow(), tutils.tflow()
t_start = time.time()
fm.handle_request(f1)
f1.reply()
fm.handle_request(f2)
f2.reply()
for hook in ("clientconnect",
"serverconnect",
"response",
"error",
"clientconnect"):
d = Dummy()
assert s.run(hook, d)[0]
d.reply()
while (time.time() - t_start) < 20 and m.call_count <= 5:
if m.call_count == 5:
return
time.sleep(0.001)
assert False
# Two instantiations
assert m.call_count == 0 # No calls yet.
assert (time.time() - t_start) < 0.09
def test_concurrent_err(self):
s = flow.State()
fm = flow.FlowMaster(None, s)
tutils.raises(
"decorator not supported for this method",
script.Script,
tutils.test_data.path("scripts/concurrent_decorator_err.py"),
fm)
def test_concurrent2():
s = flow.State()
fm = flow.FlowMaster(None, s)
s = script.Script(
tutils.test_data.path("scripts/concurrent_decorator.py"),
fm)
s.load()
m = mock.Mock()
class Dummy:
def __init__(self):
self.response = self
self.error = self
self.reply = m
t_start = time.time()
for hook in ("clientconnect",
"serverconnect",
"response",
"error",
"clientconnect"):
d = Dummy()
s.run(hook, d)
d.reply()
while (time.time() - t_start) < 20 and m.call_count <= 5:
if m.call_count == 5:
return
time.sleep(0.001)
assert False
def test_concurrent_err():
s = flow.State()
fm = flow.FlowMaster(None, s)
tutils.raises(
"Concurrent decorator not supported for 'start' method",
script.Script,
tutils.test_data.path("scripts/concurrent_decorator_err.py"),
fm)
def test_command_parsing():
@ -122,4 +126,4 @@ def test_command_parsing():
fm = flow.FlowMaster(None, s)
absfilepath = os.path.normcase(tutils.test_data.path("scripts/a.py"))
s = script.Script(absfilepath, fm)
assert os.path.isfile(s.argv[0])
assert os.path.isfile(s.args[0])