Zap old scripts infrastructure, fix concurrency tests

This commit is contained in:
Aldo Cortesi 2016-07-14 14:39:07 +12:00
parent a3a22fba33
commit a6821aad8e
9 changed files with 43 additions and 216 deletions

View File

@ -390,6 +390,3 @@ class FlowMaster(controller.Master):
@controller.handler
def tcp_close(self, flow):
self.active_flows.discard(flow)
def shutdown(self):
super(FlowMaster, self).shutdown()

View File

@ -1,11 +1,5 @@
from . import reloader
from .concurrent import concurrent
from .script import Script
from ..exceptions import ScriptException
__all__ = [
"Script",
"concurrent",
"ScriptException",
"reloader"
]

View File

@ -1,47 +0,0 @@
from __future__ import absolute_import, print_function, division
import os
from watchdog.events import RegexMatchingEventHandler
from watchdog.observers.polling import PollingObserver as Observer
# We occasionally have watchdog errors on Windows, Linux and Mac when using the native observers.
# After reading through the watchdog source code and issue tracker,
# we may want to replace this with a very simple implementation of our own.
_observers = {}
def watch(script, callback):
if script in _observers:
raise RuntimeError("Script already observed")
script_dir = os.path.dirname(os.path.abspath(script.path))
script_name = os.path.basename(script.path)
event_handler = _ScriptModificationHandler(callback, filename=script_name)
observer = Observer()
observer.schedule(event_handler, script_dir)
observer.start()
_observers[script] = observer
def unwatch(script):
observer = _observers.pop(script, None)
if observer:
observer.stop()
observer.join()
class _ScriptModificationHandler(RegexMatchingEventHandler):
def __init__(self, callback, filename='.*'):
super(_ScriptModificationHandler, self).__init__(
ignore_directories=True,
regexes=['.*' + filename]
)
self.callback = callback
def on_modified(self, event):
self.callback()
__all__ = ["watch", "unwatch"]

View File

@ -1,136 +0,0 @@
"""
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.
"""
# Do not import __future__ here, this would apply transitively to the inline scripts.
from __future__ import absolute_import, print_function, division
import os
import shlex
import sys
import contextlib
import six
from typing import List # noqa
from mitmproxy import exceptions
@contextlib.contextmanager
def scriptenv(path, args):
# type: (str, List[str]) -> None
oldargs = sys.argv
script_dir = os.path.dirname(os.path.abspath(path))
sys.argv = [path] + args
sys.path.append(script_dir)
try:
yield
finally:
sys.argv = oldargs
sys.path.pop()
class Script(object):
"""
Script object representing an inline script.
"""
def __init__(self, command):
self.command = command
self.path, self.args = self.parse_command(command)
self.ns = None
def __enter__(self):
self.load()
return self
def __exit__(self, exc_type, exc_val, exc_tb):
if exc_val:
return False # re-raise the exception
self.unload()
@staticmethod
def parse_command(command):
# type: (str) -> Tuple[str,List[str]]
"""
Returns a (path, args) tuple.
"""
if not command or not command.strip():
raise exceptions.ScriptException("Empty script command.")
# Windows: escape all backslashes in the path.
if os.name == "nt": # pragma: no cover
backslashes = shlex.split(command, posix=False)[0].count("\\")
command = command.replace("\\", "\\\\", backslashes)
args = shlex.split(command) # pragma: no cover
args[0] = os.path.expanduser(args[0])
if not os.path.exists(args[0]):
raise exceptions.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 exceptions.ScriptException("Not a file: %s" % args[0])
return args[0], args[1:]
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:
raise exceptions.ScriptException("Script is already loaded")
self.ns = {'__file__': os.path.abspath(self.path)}
with scriptenv(self.path, self.args):
try:
with open(self.path) as f:
code = compile(f.read(), self.path, 'exec')
exec(code, self.ns, self.ns)
except Exception:
six.reraise(
exceptions.ScriptException,
exceptions.ScriptException.from_exception_context(),
sys.exc_info()[2]
)
return self.run("start")
def unload(self):
try:
return self.run("done")
finally:
self.ns = None
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.
"""
if self.ns is None:
raise exceptions.ScriptException("Script not loaded.")
f = self.ns.get(name)
if f:
try:
with scriptenv(self.path, self.args):
return f(*args, **kwargs)
except Exception:
six.reraise(
exceptions.ScriptException,
exceptions.ScriptException.from_exception_context(),
sys.exc_info()[2]
)
else:
return None

View File

@ -47,15 +47,6 @@ def test_load_script():
assert ns["configure"]
class RecordingMaster(master.FlowMaster):
def __init__(self, *args, **kwargs):
master.FlowMaster.__init__(self, *args, **kwargs)
self.event_log = []
def add_event(self, e, level):
self.event_log.append((level, e))
class TestScript(mastertest.MasterTest):
def test_simple(self):
s = state.State()
@ -77,7 +68,7 @@ class TestScript(mastertest.MasterTest):
def test_reload(self):
s = state.State()
m = RecordingMaster(options.Options(), None, s)
m = mastertest.RecordingMaster(options.Options(), None, s)
with tutils.tmpdir():
with open("foo.py", "w"):
pass
@ -94,7 +85,7 @@ class TestScript(mastertest.MasterTest):
def test_exception(self):
s = state.State()
m = RecordingMaster(options.Options(), None, s)
m = mastertest.RecordingMaster(options.Options(), None, s)
sc = script.Script(
tutils.test_data.path("data/addonscripts/error.py")
)

View File

@ -1,7 +1,6 @@
import time
from mitmproxy.script import concurrent
@concurrent
def request(flow):
time.sleep(0.1)

View File

@ -3,6 +3,7 @@ import mock
from . import tutils
import netlib.tutils
from mitmproxy.flow import master
from mitmproxy import flow, proxy, models, controller
@ -39,3 +40,12 @@ class MasterTest:
t = tutils.tflow(resp=True)
fw.add(t)
f.close()
class RecordingMaster(master.FlowMaster):
def __init__(self, *args, **kwargs):
master.FlowMaster.__init__(self, *args, **kwargs)
self.event_log = []
def add_event(self, e, level):
self.event_log.append((level, e))

View File

@ -1,28 +1,47 @@
from mitmproxy.script import Script
from test.mitmproxy import tutils
from mitmproxy import controller
from mitmproxy.builtins import script
from mitmproxy import options
from mitmproxy.flow import master
from mitmproxy.flow import state
import time
from .. import mastertest, tutils
class Thing:
def __init__(self):
self.reply = controller.DummyReply()
self.live = True
@tutils.skip_appveyor
def test_concurrent():
with Script(tutils.test_data.path("data/scripts/concurrent_decorator.py")) as s:
f1, f2 = Thing(), Thing()
s.run("request", f1)
s.run("request", f2)
class TestConcurrent(mastertest.MasterTest):
@tutils.skip_appveyor
def test_concurrent(self):
s = state.State()
m = master.FlowMaster(options.Options(), None, s)
sc = script.Script(
tutils.test_data.path(
"data/addonscripts/concurrent_decorator.py"
)
)
m.addons.add(sc)
f1, f2 = tutils.tflow(), tutils.tflow()
self.invoke(m, "request", f1)
self.invoke(m, "request", f2)
start = time.time()
while time.time() - start < 5:
if f1.reply.acked and f2.reply.acked:
return
raise ValueError("Script never acked")
def test_concurrent_err():
s = Script(tutils.test_data.path("data/scripts/concurrent_decorator_err.py"))
with tutils.raises("Concurrent decorator not supported for 'start' method"):
s.load()
def test_concurrent_err(self):
s = state.State()
m = mastertest.RecordingMaster(options.Options(), None, s)
sc = script.Script(
tutils.test_data.path(
"data/addonscripts/concurrent_decorator_err.py"
)
)
with m.handlecontext():
sc.configure(options.Options())
assert "decorator not supported" in m.event_log[0][1]