combine projects

5 files changed, 282 insertions, 0 deletions

diff --git a/mitmproxy/script/__init__.py b/mitmproxy/script/__init__.py
new file mode 100644
index 00000000..3ee19b04
--- /dev/null
+++ b/mitmproxy/script/__init__.py
@@ -0,0 +1,13 @@
+from .script import Script
+from .script_context import ScriptContext
+from .concurrent import concurrent
+from ..exceptions import ScriptException
+from . import reloader
+
+__all__ = [
+    "Script",
+    "ScriptContext",
+    "concurrent",
+    "ScriptException",
+    "reloader"
+]
diff --git a/mitmproxy/script/concurrent.py b/mitmproxy/script/concurrent.py
new file mode 100644
index 00000000..f0f5e3cd
--- /dev/null
+++ b/mitmproxy/script/concurrent.py
@@ -0,0 +1,63 @@
+"""
+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)
diff --git a/mitmproxy/script/reloader.py b/mitmproxy/script/reloader.py
new file mode 100644
index 00000000..b4acf51b
--- /dev/null
+++ b/mitmproxy/script/reloader.py
@@ -0,0 +1,46 @@
+import os
+import sys
+from watchdog.events import RegexMatchingEventHandler
+if sys.platform == 'darwin':
+    from watchdog.observers.polling import PollingObserver as Observer
+else:
+    from watchdog.observers import Observer
+# The OSX reloader in watchdog 0.8.3 breaks when unobserving paths.
+# We use the PollingObserver instead.
+
+_observers = {}
+
+
+def watch(script, callback):
+    if script in _observers:
+        raise RuntimeError("Script already observed")
+    script_dir = os.path.dirname(os.path.abspath(script.args[0]))
+    script_name = os.path.basename(script.args[0])
+    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"]
diff --git a/mitmproxy/script/script.py b/mitmproxy/script/script.py
new file mode 100644
index 00000000..55778851
--- /dev/null
+++ b/mitmproxy/script/script.py
@@ -0,0 +1,100 @@
+"""
+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.
+import os
+import shlex
+import traceback
+import sys
+from ..exceptions import ScriptException
+
+
+class Script(object):
+
+    """
+    Script object representing an inline script.
+    """
+
+    def __init__(self, command, context):
+        self.command = command
+        self.args = self.parse_command(command)
+        self.ctx = context
+        self.ns = None
+        self.load()
+
+    @property
+    def filename(self):
+        return self.args[0]
+
+    @staticmethod
+    def parse_command(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]))
+        self.ns = {'__file__': os.path.abspath(self.args[0])}
+        sys.path.append(script_dir)
+        try:
+            execfile(self.args[0], self.ns, self.ns)
+        except Exception as e:
+            # Python 3: use exception chaining, https://www.python.org/dev/peps/pep-3134/
+            raise ScriptException(traceback.format_exc(e))
+        finally:
+            sys.path.pop()
+        return self.run("start", self.args)
+
+    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 ScriptException("Script not loaded.")
+        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
diff --git a/mitmproxy/script/script_context.py b/mitmproxy/script/script_context.py
new file mode 100644
index 00000000..cd5d4b61
--- /dev/null
+++ b/mitmproxy/script/script_context.py
@@ -0,0 +1,60 @@
+"""
+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)