Grasscutters/mitmproxy
3
0
Fork 0
mirror of https://github.com/Grasscutters/mitmproxy.git synced 2024-11-26 10:16:27 +00:00
Code Issues Projects Releases Wiki Activity

structure libmproxy.script

Browse Source
This commit is contained in:
Maximilian Hils 2015-11-14 04:21:38 +01:00
commit 4499ab61c0
11 changed files with 293 additions and 223 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
libmproxy/console/__init__.py
View File

@ -315,8 +315,8 @@ class ConsoleMaster(flow.FlowMaster):
signals.add_event("Running script on flow: %s" % command, "debug")
try:
s = script.Script(command, self)
except script.ScriptError as v:
s = script.Script(command, script.ScriptContext(self))
except script.ScriptException as v:
signals.status_message.send(
message = "Error loading script."
)
@ -731,3 +731,7 @@ class ConsoleMaster(flow.FlowMaster):
if f:
self.process_flow(f)
return f
def script_change(self, script):
self.masterq.put(("script_change", script))
signals.status_message.send(message="<{}> reloaded.".format(script.args[0]))

2
libmproxy/console/grideditor.py
View File

@ -636,7 +636,7 @@ class ScriptEditor(GridEditor):
def is_error(self, col, val):
try:
script.Script.parse_command(val)
except script.ScriptError as v:
except script.ScriptException as v:
return str(v)

4
libmproxy/exceptions.py
View File

@ -48,3 +48,7 @@ class ContentViewException(ProxyException):
class ReplayException(ProxyException):
pass
class ScriptException(ProxyException):
pass

18
libmproxy/flow.py
View File

@ -640,6 +640,7 @@ class FlowMaster(controller.Master):
self.stream = None
self.apps = AppRegistry()
script.script_change.connect(self.script_change)
def start_app(self, host, port):
self.apps.add(
@ -661,18 +662,18 @@ class FlowMaster(controller.Master):
def unload_script(self, script_obj):
try:
script_obj.unload()
except script.ScriptError as e:
except script.ScriptException as e:
self.add_event("Script error:\n" + str(e), "error")
self.scripts.remove(script_obj)
def load_script(self, command):
"""
Loads a script. Returns an error description if something went
wrong.
"""
try:
s = script.Script(command, self)
except script.ScriptError as v:
s = script.Script(command, script.ScriptContext(self))
except script.ScriptException as v:
return v.args[0]
self.scripts.append(s)
@ -680,7 +681,7 @@ class FlowMaster(controller.Master):
if script_obj and not self.pause_scripts:
try:
script_obj.run(name, *args, **kwargs)
except script.ScriptError as e:
except script.ScriptException as e:
self.add_event("Script error:\n" + str(e), "error")
def run_script_hook(self, name, *args, **kwargs):
@ -1019,6 +1020,9 @@ class FlowMaster(controller.Master):
def handle_accept_intercept(self, f):
self.state.update_flow(f)
def handle_script_change(self, script):
script.load()
def shutdown(self):
self.unload_scripts()
controller.Master.shutdown(self)
@ -1035,6 +1039,10 @@ class FlowMaster(controller.Master):
self.stream.fo.close()
self.stream = None
def script_change(self, script):
self.masterq.put(("script_change", script))
self.add_event("<{}> reloaded.".format(script.args[0]))
def read_flows_from_paths(paths):
"""

202
libmproxy/script.py
View File

@ -1,202 +0,0 @@
from __future__ import absolute_import
import os
import traceback
import threading
import shlex
import sys
from . import contentviews
class ScriptError(Exception):
pass
class ScriptContext:
"""
The script context should be used to interact with the global mitmproxy state from within a
script.
"""
def __init__(self, master):
self._master = master
def log(self, message, level="info"):
"""
Logs an event.
By default, only events with level "error" get displayed. This can be controlled with the "-v" switch.
How log messages are handled depends on the front-end. mitmdump will print them to stdout,
mitmproxy sends output to the eventlog for display ("e" keyboard shortcut).
"""
self._master.add_event(message, level)
def kill_flow(self, f):
"""
Kills a flow immediately. No further data will be sent to the client or the server.
"""
f.kill(self._master)
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.
"""
return self._master.replay_request(f, block=True, run_scripthooks=False)
@property
def app_registry(self):
return self._master.apps
def add_contentview(self, view_obj):
contentviews.add(view_obj)
def remove_contentview(self, view_obj):
contentviews.remove(view_obj)
class Script:
"""
Script object representing an inline script.
"""
def __init__(self, command, master):
self.command = command
self.args = self.parse_command(command)
self.ctx = ScriptContext(master)
self.ns = None
self.load()
@classmethod
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.
backslashes = shlex.split(command, posix=False)[0].count("\\")
command = command.replace("\\", "\\\\", backslashes)
args = shlex.split(command)
args[0] = os.path.expanduser(args[0])
if not os.path.exists(args[0]):
raise ScriptError(
("Script file not found: %s.\r\n"
"If your script path contains spaces, "
"make sure to wrap it in additional quotes, e.g. -s \"'./foo bar/baz.py' --args\".") %
args[0])
elif os.path.isdir(args[0]):
raise ScriptError("Not a file: %s" % args[0])
return args
def load(self):
"""
Loads an inline script.
Returns:
The return value of self.run("start", ...)
Raises:
ScriptError on failure
"""
if self.ns is not None:
self.unload()
script_dir = os.path.dirname(os.path.abspath(self.args[0]))
ns = {'__file__': os.path.abspath(self.args[0])}
sys.path.append(script_dir)
try:
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
return self.run("start", self.args)
def unload(self):
ret = self.run("done")
self.ns = None
return ret
def run(self, name, *args, **kwargs):
"""
Runs an inline script hook.
Returns:
The return value of the method.
None, if the script does not provide the method.
Raises:
ScriptError if there was an exception.
"""
f = self.ns.get(name)
if f:
try:
return f(self.ctx, *args, **kwargs)
except Exception as e:
raise ScriptError(traceback.format_exc(e))
else:
return None
class ReplyProxy(object):
def __init__(self, original_reply, script_thread):
self.original_reply = original_reply
self.script_thread = script_thread
self._ignore_call = True
self.lock = threading.Lock()
def __call__(self, *args, **kwargs):
with self.lock:
if self._ignore_call:
self.script_thread.start()
self._ignore_call = False
return
self.original_reply(*args, **kwargs)
def __getattr__(self, k):
return getattr(self.original_reply, k)
def _handle_concurrent_reply(fn, o, *args, **kwargs):
# Make first call to o.reply a no op and start the script thread.
# We must not start the script thread before, as this may lead to a nasty race condition
# where the script thread replies a different response before the normal reply, which then gets swallowed.
def run():
fn(*args, **kwargs)
# If the script did not call .reply(), we have to do it now.
reply_proxy()
script_thread = ScriptThread(target=run)
reply_proxy = ReplyProxy(o.reply, script_thread)
o.reply = reply_proxy
class ScriptThread(threading.Thread):
name = "ScriptThread"
def concurrent(fn):
if fn.func_name in (
"request",
"response",
"error",
"clientconnect",
"serverconnect",
"clientdisconnect",
"next_layer"):
def _concurrent(ctx, obj):
_handle_concurrent_reply(fn, obj, ctx, obj)
return _concurrent
raise NotImplementedError(
"Concurrent decorator not supported for '%s' method." % fn.func_name)

11
libmproxy/script/__init__.py Normal file
View File

@ -0,0 +1,11 @@
from .script import Script, script_change
from .script_context import ScriptContext
from .concurrent import concurrent
from ..exceptions import ScriptException
__all__ = [
"Script", "script_change",
"ScriptContext",
"concurrent",
"ScriptException"
]

62
libmproxy/script/concurrent.py Normal file
View File

@ -0,0 +1,62 @@
"""
This module provides a @concurrent decorator primitive to
offload computations from mitmproxy's main master thread.
"""
from __future__ import absolute_import, print_function, division
import threading
class ReplyProxy(object):
def __init__(self, original_reply, script_thread):
self.original_reply = original_reply
self.script_thread = script_thread
self._ignore_call = True
self.lock = threading.Lock()
def __call__(self, *args, **kwargs):
with self.lock:
if self._ignore_call:
self.script_thread.start()
self._ignore_call = False
return
self.original_reply(*args, **kwargs)
def __getattr__(self, k):
return getattr(self.original_reply, k)
def _handle_concurrent_reply(fn, o, *args, **kwargs):
# Make first call to o.reply a no op and start the script thread.
# We must not start the script thread before, as this may lead to a nasty race condition
# where the script thread replies a different response before the normal reply, which then gets swallowed.
def run():
fn(*args, **kwargs)
# If the script did not call .reply(), we have to do it now.
reply_proxy()
script_thread = ScriptThread(target=run)
reply_proxy = ReplyProxy(o.reply, script_thread)
o.reply = reply_proxy
class ScriptThread(threading.Thread):
name = "ScriptThread"
def concurrent(fn):
if fn.func_name in (
"request",
"response",
"error",
"clientconnect",
"serverconnect",
"clientdisconnect",
"next_layer"):
def _concurrent(ctx, obj):
_handle_concurrent_reply(fn, obj, ctx, obj)
return _concurrent
raise NotImplementedError(
"Concurrent decorator not supported for '%s' method." % fn.func_name)

123
libmproxy/script/script.py Normal file
View File

@ -0,0 +1,123 @@
"""
The script object representing mitmproxy inline scripts.
Script objects know nothing about mitmproxy or mitmproxy's API - this knowledge is provided
by the mitmproxy-specific ScriptContext.
"""
from __future__ import absolute_import, print_function, division
import os
import shlex
import traceback
import sys
import blinker
from watchdog.events import PatternMatchingEventHandler, FileModifiedEvent
from watchdog.observers import Observer
from ..exceptions import ScriptException
script_change = blinker.Signal()
class Script(object):
"""
Script object representing an inline script.
"""
def __init__(self, command, context, use_reloader=True):
self.command = command
self.args = self.parse_command(command)
self.ctx = context
self.ns = None
self.load()
if use_reloader:
self.start_observe()
@classmethod
def parse_command(cls, command):
if not command or not command.strip():
raise ScriptException("Empty script command.")
if os.name == "nt": # Windows: escape all backslashes in the path.
backslashes = shlex.split(command, posix=False)[0].count("\\")
command = command.replace("\\", "\\\\", backslashes)
args = shlex.split(command)
args[0] = os.path.expanduser(args[0])
if not os.path.exists(args[0]):
raise ScriptException(
("Script file not found: %s.\r\n"
"If your script path contains spaces, "
"make sure to wrap it in additional quotes, e.g. -s \"'./foo bar/baz.py' --args\".") %
args[0])
elif os.path.isdir(args[0]):
raise ScriptException("Not a file: %s" % args[0])
return args
def load(self):
"""
Loads an inline script.
Returns:
The return value of self.run("start", ...)
Raises:
ScriptException on failure
"""
if self.ns is not None:
self.unload()
script_dir = os.path.dirname(os.path.abspath(self.args[0]))
ns = {'__file__': os.path.abspath(self.args[0])}
sys.path.append(script_dir)
try:
execfile(self.args[0], ns, ns)
except Exception as e:
# Python 3: use exception chaining, https://www.python.org/dev/peps/pep-3134/
raise ScriptException(traceback.format_exc(e))
sys.path.pop()
self.ns = ns
return self.run("start", self.args)
def unload(self):
ret = self.run("done")
self.ns = None
return ret
def run(self, name, *args, **kwargs):
"""
Runs an inline script hook.
Returns:
The return value of the method.
None, if the script does not provide the method.
Raises:
ScriptException if there was an exception.
"""
f = self.ns.get(name)
if f:
try:
return f(self.ctx, *args, **kwargs)
except Exception as e:
raise ScriptException(traceback.format_exc(e))
else:
return None
def start_observe(self):
script_dir = os.path.dirname(self.args[0])
event_handler = ScriptModified(self)
observer = Observer()
observer.schedule(event_handler, script_dir)
observer.start()
def stop_observe(self):
raise NotImplementedError() # FIXME
class ScriptModified(PatternMatchingEventHandler):
def __init__(self, script):
# We could enumerate all relevant *.py files (as werkzeug does it),
# but our case looks like it isn't as simple as enumerating sys.modules.
# This should be good enough for now.
super(ScriptModified, self).__init__(ignore_directories=True, patterns=["*.py"])
self.script = script
def on_modified(self, event=FileModifiedEvent):
script_change.send(self.script)

59
libmproxy/script/script_context.py Normal file
View File

@ -0,0 +1,59 @@
"""
The mitmproxy script context provides an API to inline scripts.
"""
from __future__ import absolute_import, print_function, division
from .. import contentviews
class ScriptContext(object):
"""
The script context should be used to interact with the global mitmproxy state from within a
script.
"""
def __init__(self, master):
self._master = master
def log(self, message, level="info"):
"""
Logs an event.
By default, only events with level "error" get displayed. This can be controlled with the "-v" switch.
How log messages are handled depends on the front-end. mitmdump will print them to stdout,
mitmproxy sends output to the eventlog for display ("e" keyboard shortcut).
"""
self._master.add_event(message, level)
def kill_flow(self, f):
"""
Kills a flow immediately. No further data will be sent to the client or the server.
"""
f.kill(self._master)
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.
"""
return self._master.replay_request(f, block=True, run_scripthooks=False)
@property
def app_registry(self):
return self._master.apps
def add_contentview(self, view_obj):
contentviews.add(view_obj)
def remove_contentview(self, view_obj):
contentviews.remove(view_obj)

2
test/test_examples.py
View File

@ -22,7 +22,7 @@ def test_load_scripts():
if "modify_response_body" in f:
f += " foo bar" # two arguments required
try:
s = script.Script(f, tmaster) # Loads the script file.
s = script.Script(f, script.ScriptContext(tmaster)) # Loads the script file.
except Exception as v:
if "ImportError" not in str(v):
raise

25
test/test_script.py
View File

@ -9,13 +9,13 @@ 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)
p = script.Script("%s --var 40" % sp, script.ScriptContext(fm))
assert "here" in p.ns
assert p.run("here") == 41
assert p.run("here") == 42
tutils.raises(script.ScriptError, p.run, "errargs")
tutils.raises(script.ScriptException, p.run, "errargs")
# Check reload
p.load()
@ -36,29 +36,30 @@ def test_duplicate_flow():
def test_err():
s = flow.State()
fm = flow.FlowMaster(None, s)
sc = script.ScriptContext(fm)
tutils.raises(
"not found",
script.Script, "nonexistent", fm
script.Script, "nonexistent", sc
)
tutils.raises(
"not a file",
script.Script, tutils.test_data.path("scripts"), fm
script.Script, tutils.test_data.path("scripts"), sc
)
tutils.raises(
script.ScriptError,
script.Script, tutils.test_data.path("scripts/syntaxerr.py"), fm
script.ScriptException,
script.Script, tutils.test_data.path("scripts/syntaxerr.py"), sc
)
tutils.raises(
script.ScriptError,
script.Script, tutils.test_data.path("scripts/loaderr.py"), fm
script.ScriptException,
script.Script, tutils.test_data.path("scripts/loaderr.py"), sc
)
scr = script.Script(tutils.test_data.path("scripts/unloaderr.py"), fm)
tutils.raises(script.ScriptError, scr.unload)
scr = script.Script(tutils.test_data.path("scripts/unloaderr.py"), sc)
tutils.raises(script.ScriptException, scr.unload)
def test_concurrent():
@ -84,7 +85,7 @@ def test_concurrent2():
fm = flow.FlowMaster(None, s)
s = script.Script(
tutils.test_data.path("scripts/concurrent_decorator.py"),
fm)
script.ScriptContext(fm))
s.load()
m = mock.Mock()
@ -125,6 +126,6 @@ def test_command_parsing():
s = flow.State()
fm = flow.FlowMaster(None, s)
absfilepath = os.path.normcase(tutils.test_data.path("scripts/a.py"))
s = script.Script(absfilepath, fm)
s = script.Script(absfilepath, script.ScriptContext(fm))
assert os.path.isfile(s.args[0])