diff options
| -rw-r--r-- | examples/filt.py | 20 | ||||
| -rw-r--r-- | examples/flowwriter.py | 25 | ||||
| -rw-r--r-- | examples/iframe_injector.py | 37 | ||||
| -rw-r--r-- | mitmproxy/builtins/script.py | 45 | ||||
| -rw-r--r-- | test/mitmproxy/builtins/test_script.py | 21 | ||||
| -rw-r--r-- | test/mitmproxy/data/addonscripts/addon.py | 22 | ||||
| -rw-r--r-- | test/mitmproxy/data/addonscripts/recorder.py | 36 |
7 files changed, 136 insertions, 70 deletions
diff --git a/examples/filt.py b/examples/filt.py index 21744edd..9ccf9fa1 100644 --- a/examples/filt.py +++ b/examples/filt.py @@ -1,18 +1,20 @@ -# This scripts demonstrates how to use mitmproxy's filter pattern in inline scripts. +# This scripts demonstrates how to use mitmproxy's filter pattern in scripts. # Usage: mitmdump -s "filt.py FILTER" import sys from mitmproxy import filt -state = {} + +class Filter: + def __init__(self, spec): + self.filter = filt.parse(spec) + + def response(self, flow): + if flow.match(self.filter): + print("Flow matches filter:") + print(flow) def start(): if len(sys.argv) != 2: raise ValueError("Usage: -s 'filt.py FILTER'") - state["filter"] = filt.parse(sys.argv[1]) - - -def response(flow): - if flow.match(state["filter"]): - print("Flow matches filter:") - print(flow) + return Filter(sys.argv[1]) diff --git a/examples/flowwriter.py b/examples/flowwriter.py index 07c7ca20..df2e5a40 100644 --- a/examples/flowwriter.py +++ b/examples/flowwriter.py @@ -3,20 +3,21 @@ import sys from mitmproxy.flow import FlowWriter -state = {} + +class Writer: + def __init__(self, path): + if path == "-": + f = sys.stdout + else: + f = open(path, "wb") + self.w = FlowWriter(f) + + def response(self, flow): + if random.choice([True, False]): + self.w.add(flow) def start(): if len(sys.argv) != 2: raise ValueError('Usage: -s "flowriter.py filename"') - - if sys.argv[1] == "-": - f = sys.stdout - else: - f = open(sys.argv[1], "wb") - state["flow_writer"] = FlowWriter(f) - - -def response(flow): - if random.choice([True, False]): - state["flow_writer"].add(flow) + return Writer(sys.argv[1]) diff --git a/examples/iframe_injector.py b/examples/iframe_injector.py index 352c3c24..33d18bbd 100644 --- a/examples/iframe_injector.py +++ b/examples/iframe_injector.py @@ -3,26 +3,27 @@ import sys from bs4 import BeautifulSoup -iframe_url = None + +class Injector: + def __init__(self, iframe_url): + self.iframe_url = iframe_url + + def response(self, flow): + if flow.request.host in self.iframe_url: + return + html = BeautifulSoup(flow.response.content, "lxml") + if html.body: + iframe = html.new_tag( + "iframe", + src=self.iframe_url, + frameborder=0, + height=0, + width=0) + html.body.insert(0, iframe) + flow.response.content = str(html).encode("utf8") def start(): if len(sys.argv) != 2: raise ValueError('Usage: -s "iframe_injector.py url"') - global iframe_url - iframe_url = sys.argv[1] - - -def response(flow): - if flow.request.host in iframe_url: - return - html = BeautifulSoup(flow.response.content, "lxml") - if html.body: - iframe = html.new_tag( - "iframe", - src=iframe_url, - frameborder=0, - height=0, - width=0) - html.body.insert(0, iframe) - flow.response.content = str(html).encode("utf8") + return Injector(sys.argv[1]) diff --git a/mitmproxy/builtins/script.py b/mitmproxy/builtins/script.py index ab068e47..ed5f2ecd 100644 --- a/mitmproxy/builtins/script.py +++ b/mitmproxy/builtins/script.py @@ -16,6 +16,19 @@ import watchdog.events from watchdog.observers import polling +class NS: + def __init__(self, ns): + self.__dict__["ns"] = ns + + def __getattr__(self, key): + if key not in self.ns: + raise AttributeError("No such element: %s", key) + return self.ns[key] + + def __setattr__(self, key, value): + self.__dict__["ns"][key] = value + + def parse_command(command): """ Returns a (path, args) tuple. @@ -74,18 +87,27 @@ def load_script(path, args): ns = {'__file__': os.path.abspath(path)} with scriptenv(path, args): exec(code, ns, ns) - return ns + return NS(ns) class ReloadHandler(watchdog.events.FileSystemEventHandler): def __init__(self, callback): self.callback = callback + def filter(self, event): + if event.is_directory: + return False + if os.path.basename(event.src_path).startswith("."): + return False + return True + def on_modified(self, event): - self.callback() + if self.filter(event): + self.callback() def on_created(self, event): - self.callback() + if self.filter(event): + self.callback() class Script: @@ -118,27 +140,32 @@ class Script: # It's possible for ns to be un-initialised if we failed during # configure if self.ns is not None and not self.dead: - func = self.ns.get(name) + func = getattr(self.ns, name, None) if func: with scriptenv(self.path, self.args): - func(*args, **kwargs) + return func(*args, **kwargs) def reload(self): self.should_reload.set() + def load_script(self): + self.ns = load_script(self.path, self.args) + ret = self.run("start") + if ret: + self.ns = ret + self.run("start") + def tick(self): if self.should_reload.is_set(): self.should_reload.clear() ctx.log.info("Reloading script: %s" % self.name) - self.ns = load_script(self.path, self.args) - self.start() + self.load_script() self.configure(self.last_options) else: self.run("tick") def start(self): - self.ns = load_script(self.path, self.args) - self.run("start") + self.load_script() def configure(self, options): self.last_options = options diff --git a/test/mitmproxy/builtins/test_script.py b/test/mitmproxy/builtins/test_script.py index f37c7f94..826d2a91 100644 --- a/test/mitmproxy/builtins/test_script.py +++ b/test/mitmproxy/builtins/test_script.py @@ -48,7 +48,7 @@ def test_load_script(): "data/addonscripts/recorder.py" ), [] ) - assert ns["configure"] + assert ns.start class TestScript(mastertest.MasterTest): @@ -61,16 +61,16 @@ class TestScript(mastertest.MasterTest): ) ) m.addons.add(sc) - assert sc.ns["call_log"] == [ + assert sc.ns.call_log == [ ("solo", "start", (), {}), ("solo", "configure", (options.Options(),), {}) ] - sc.ns["call_log"] = [] + sc.ns.call_log = [] f = tutils.tflow(resp=True) self.invoke(m, "request", f) - recf = sc.ns["call_log"][0] + recf = sc.ns.call_log[0] assert recf[1] == "request" def test_reload(self): @@ -116,6 +116,19 @@ class TestScript(mastertest.MasterTest): assert not fm.state.view[0].request.is_replay assert fm.state.view[1].request.is_replay + def test_addon(self): + s = state.State() + m = master.FlowMaster(options.Options(), None, s) + sc = script.Script( + tutils.test_data.path( + "data/addonscripts/addon.py" + ) + ) + m.addons.add(sc) + assert sc.ns.event_log == [ + 'scriptstart', 'addonstart', 'addonconfigure' + ] + class TestScriptLoader(mastertest.MasterTest): def test_simple(self): diff --git a/test/mitmproxy/data/addonscripts/addon.py b/test/mitmproxy/data/addonscripts/addon.py new file mode 100644 index 00000000..3b09d231 --- /dev/null +++ b/test/mitmproxy/data/addonscripts/addon.py @@ -0,0 +1,22 @@ +event_log = [] + + +class Addon: + @property + def event_log(self): + return event_log + + def start(self): + event_log.append("addonstart") + + def configure(self, options): + event_log.append("addonconfigure") + + +def configure(options): + event_log.append("addonconfigure") + + +def start(): + event_log.append("scriptstart") + return Addon() diff --git a/test/mitmproxy/data/addonscripts/recorder.py b/test/mitmproxy/data/addonscripts/recorder.py index b6ac8d89..890e6f4e 100644 --- a/test/mitmproxy/data/addonscripts/recorder.py +++ b/test/mitmproxy/data/addonscripts/recorder.py @@ -2,24 +2,24 @@ from mitmproxy import controller from mitmproxy import ctx import sys -call_log = [] -if len(sys.argv) > 1: - name = sys.argv[1] -else: - name = "solo" +class CallLogger: + call_log = [] -# Keep a log of all possible event calls -evts = list(controller.Events) + ["configure"] -for i in evts: - def mkprox(): - evt = i + def __init__(self, name = "solo"): + self.name = name - def prox(*args, **kwargs): - lg = (name, evt, args, kwargs) - if evt != "log": - ctx.log.info(str(lg)) - call_log.append(lg) - ctx.log.debug("%s %s" % (name, evt)) - return prox - globals()[i] = mkprox() + def __getattr__(self, attr): + if attr in controller.Events: + def prox(*args, **kwargs): + lg = (self.name, attr, args, kwargs) + if attr != "log": + ctx.log.info(str(lg)) + self.call_log.append(lg) + ctx.log.debug("%s %s" % (self.name, attr)) + return prox + raise AttributeError + + +def start(): + return CallLogger(*sys.argv[1:]) |