mirror of
https://github.com/Grasscutters/mitmproxy.git
synced 2024-11-23 08:11:00 +00:00
Zap old scripts infrastructure, fix concurrency tests
This commit is contained in:
parent
a3a22fba33
commit
a6821aad8e
@ -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()
|
||||
|
@ -1,11 +1,5 @@
|
||||
from . import reloader
|
||||
from .concurrent import concurrent
|
||||
from .script import Script
|
||||
from ..exceptions import ScriptException
|
||||
|
||||
__all__ = [
|
||||
"Script",
|
||||
"concurrent",
|
||||
"ScriptException",
|
||||
"reloader"
|
||||
]
|
||||
|
@ -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"]
|
@ -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
|
@ -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")
|
||||
)
|
||||
|
@ -1,7 +1,6 @@
|
||||
import time
|
||||
from mitmproxy.script import concurrent
|
||||
|
||||
|
||||
@concurrent
|
||||
def request(flow):
|
||||
time.sleep(0.1)
|
@ -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))
|
||||
|
@ -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]
|
||||
|
Loading…
Reference in New Issue
Block a user