From cd4eea39348fc9d59d4f0cc3f71384f7c1c2b2e2 Mon Sep 17 00:00:00 2001
From: Aldo Cortesi <aldo@nullcube.com>
Date: Fri, 18 Feb 2011 12:40:45 +1300
Subject: [PATCH] First pass of script hooks for mitmdump.

Also stub out docs, improve mitmdump error handling.
---
 doc-src/02-docstyle.css |  2 +-
 doc-src/index.py        | 23 ++++++++++---
 doc-src/library.html    |  5 +--
 doc-src/mitmdump.html   |  0
 doc-src/mitmproxy.html  |  0
 doc-src/scripts.html    | 11 +++++++
 examples/simple_script  |  6 ++++
 libmproxy/console.py    |  3 +-
 libmproxy/dump.py       | 30 +++++++++++++----
 libmproxy/flow.py       |  3 +-
 mitmdump                | 24 +++++++++++---
 test/test_dump.py       | 72 ++++++++++++++++++++++++-----------------
 test/test_flow.py       |  2 +-
 13 files changed, 127 insertions(+), 54 deletions(-)
 create mode 100644 doc-src/mitmdump.html
 create mode 100644 doc-src/mitmproxy.html
 create mode 100644 doc-src/scripts.html
 create mode 100755 examples/simple_script

diff --git a/doc-src/02-docstyle.css b/doc-src/02-docstyle.css
index 236bfc529..dd61d2e82 100644
--- a/doc-src/02-docstyle.css
+++ b/doc-src/02-docstyle.css
@@ -2,7 +2,7 @@ body {
     -x-system-font:none;
     font-family: Helvetica,Arial,Tahoma,Verdana,Sans-Serif;
     color: #555555;
-    font-size: 1.3em;
+    font-size: 1.1em;
 }
 
 a {
diff --git a/doc-src/index.py b/doc-src/index.py
index 1ff22108d..c34fcb78b 100644
--- a/doc-src/index.py
+++ b/doc-src/index.py
@@ -1,7 +1,7 @@
+import os
 import countershape
 from countershape import Page, Directory, PythonModule, markup
-import countershape.grok
-
+import countershape.grok, countershape.template    
 
 this.layout = countershape.Layout("_layout.html")
 ns.docTitle = "mitmproxy"
@@ -16,12 +16,25 @@ ns.sidebar = countershape.widgets.SiblingPageIndex(
           )
 
 ns.license = file("../LICENSE").read()
-ns.index_contents = file("../README").read()
-ns.example = file("../examples/stickycookies.py").read()
+ns.index_contents = file("../README.mkd").read()
+
+
+top = os.path.abspath(os.getcwd())
+def example(s):
+    d = file(os.path.join(top, s)).read()
+    return countershape.template.pySyntax(d)
+
+
+ns.example = example
+
+
 
 pages = [
     Page("index.html", "introduction"),
-    Page("library.html", "library"),
+    Page("mitmproxy.html", "mitmproxy"),
+    Page("mitmdump.html", "mitmdump"),
+    Page("scripts.html", "scripts"),
+    Page("library.html", "libmproxy"),
     Page("faq.html", "faq"),
     Page("admin.html", "administrivia")
 ]
diff --git a/doc-src/library.html b/doc-src/library.html
index e8533731e..2266c0777 100644
--- a/doc-src/library.html
+++ b/doc-src/library.html
@@ -8,8 +8,5 @@ this lets you log in to a site using your browser, and then make subsequent
 requests using a tool like __curl__, which will then seem to be part of the
 authenticated session.
 
-
-<!--(block |pySyntax)-->
-$!example!$
-<!--(end)-->
+$!example("../examples/stickycookies.py")!$
 
diff --git a/doc-src/mitmdump.html b/doc-src/mitmdump.html
new file mode 100644
index 000000000..e69de29bb
diff --git a/doc-src/mitmproxy.html b/doc-src/mitmproxy.html
new file mode 100644
index 000000000..e69de29bb
diff --git a/doc-src/scripts.html b/doc-src/scripts.html
new file mode 100644
index 000000000..b7085c79e
--- /dev/null
+++ b/doc-src/scripts.html
@@ -0,0 +1,11 @@
+
+Both __mitmproxy__ and __mitmdump__ allow you to modify requests and responses
+with external scripts. The script interface is simple - scripts simply read,
+modify and return a single __libmproxy.flow.Flow__ object, using the methods
+defined in the __libmproxy.script__ module. Scripts must be executable.
+
+$!example("../examples/simple_script")!$
+
+
+
+
diff --git a/examples/simple_script b/examples/simple_script
new file mode 100755
index 000000000..aed937ec7
--- /dev/null
+++ b/examples/simple_script
@@ -0,0 +1,6 @@
+#!/usr/bin/env python
+from libmproxy import script
+
+f = script.load_flow()
+f.request.headers["newheader"] = ["foo"]
+script.return_flow(f)
diff --git a/libmproxy/console.py b/libmproxy/console.py
index c14f1aefb..895974d2e 100644
--- a/libmproxy/console.py
+++ b/libmproxy/console.py
@@ -549,7 +549,7 @@ class ConnectionView(WWrap):
         path = os.path.expanduser(path)
         self.state.last_script = path
         try:
-            newflow, serr = self.flow.run_script(path)
+            serr = self.flow.run_script(path)
         except flow.RunException, e:
             if e.errout:
                 serr = "Script error code: %s\n\n"%e.returncode + e.errout
@@ -559,7 +559,6 @@ class ConnectionView(WWrap):
         if serr:
             serr = "Script output:\n\n" + serr
             self.master.spawn_external_viewer(serr, None)
-        self.flow.load_state(newflow.get_state())
         self.master.refresh_connection(self.flow)
 
 
diff --git a/libmproxy/dump.py b/libmproxy/dump.py
index 6b8c886de..372e6ef60 100644
--- a/libmproxy/dump.py
+++ b/libmproxy/dump.py
@@ -8,6 +8,8 @@ class Options(object):
     __slots__ = [
         "verbosity",
         "wfile",
+        "request_script",
+        "response_script",
     ]
     def __init__(self, **kwargs):
         for k, v in kwargs.items():
@@ -44,8 +46,24 @@ class DumpMaster(flow.FlowMaster):
         flow.FlowMaster.handle_error(self, r)
         r.ack()
 
+    def _runscript(self, f, script):
+        try:
+            ret = f.run_script(script)
+            if self.o.verbosity > 0:
+                print >> self.outfile, ret
+        except flow.RunException, e:
+            if e.errout:
+                eout = "Script output:\n" + self.indent(4, e.errout) + "\n"
+            else:
+                eout = ""
+            raise DumpError(
+                    "%s: %s\n%s"%(script, e.args[0], eout)
+                )
+
     def handle_request(self, r):
-        flow.FlowMaster.handle_request(self, r)
+        f = flow.FlowMaster.handle_request(self, r)
+        if self.o.request_script:
+            self._runscript(f, self.o.request_script)
         r.ack()
 
     def indent(self, n, t):
@@ -55,6 +73,8 @@ class DumpMaster(flow.FlowMaster):
     def handle_response(self, msg):
         f = flow.FlowMaster.handle_response(self, msg)
         if f:
+            if self.o.response_script:
+                self._runscript(f, self.o.response_script)
             msg.ack()
             if self.filt and not f.match(self.filt):
                     return
@@ -96,8 +116,6 @@ class DumpMaster(flow.FlowMaster):
     def run(self):
         try:
             return flow.FlowMaster.run(self)
-        except KeyboardInterrupt:
-            pass
-        except Exception, v:
-            traceback.print_exc()
-        self.shutdown()
+        except BaseException, v:
+            self.shutdown()
+            raise
diff --git a/libmproxy/flow.py b/libmproxy/flow.py
index 63da0230a..cea0ca1cc 100644
--- a/libmproxy/flow.py
+++ b/libmproxy/flow.py
@@ -82,7 +82,8 @@ class Flow:
                     p.returncode,
                     se
                 )
-        return f, se
+        self.load_state(f.get_state())
+        return se
 
     def get_state(self, nobackup=False):
         d = dict(
diff --git a/mitmdump b/mitmdump
index 886bc5cc6..8a8e45b79 100755
--- a/mitmdump
+++ b/mitmdump
@@ -41,8 +41,16 @@ if __name__ == '__main__':
     parser.add_option("-w", "--writefile",
                       action="store", dest="wfile", default=None,
                       help="Write flows to file.")
+    parser.add_option("", "--reqscript",
+                      action="store", dest="request_script", default=None,
+                      help="Script to run when a request is recieved.")
+    parser.add_option("", "--respscript",
+                      action="store", dest="response_script", default=None,
+                      help="Script to run when a response is recieved.")
+
     options, args = parser.parse_args()
 
+
     if options.quiet:
         options.verbose = 0
 
@@ -58,15 +66,23 @@ if __name__ == '__main__':
 
     dumpopts = dump.Options(
         verbosity = options.verbose,
-        wfile = options.wfile
+        wfile = options.wfile,
+        request_script = options.request_script,
+        response_script = options.response_script,
     )
     if args:
         filt = " ".join(args)
     else:
         filt = None
 
-    m = dump.DumpMaster(server, dumpopts, filt)
-
     if options.verbose > 0:
         print >> sys.stderr, "Running on port %s"%options.port
-    m.run()
+
+    try:
+        m = dump.DumpMaster(server, dumpopts, filt)
+        m.run()
+    except dump.DumpError, e:
+        print >> sys.stderr, "mitmdump:", e
+        sys.exit(1)
+    except KeyboardInterrupt:
+        pass
diff --git a/test/test_dump.py b/test/test_dump.py
index 0d0a62199..978bf1384 100644
--- a/test/test_dump.py
+++ b/test/test_dump.py
@@ -6,7 +6,7 @@ import utils
 
 
 class uDumpMaster(libpry.AutoTree):
-    def _dummy_cycle(self, m):
+    def _cycle(self, m):
         req = utils.treq()
         cc = req.client_conn
         resp = utils.tresp(req)
@@ -14,50 +14,64 @@ class uDumpMaster(libpry.AutoTree):
         m.handle_request(req)
         m.handle_response(resp)
 
+    def _dummy_cycle(self, filt, **options):
+        cs = StringIO()
+        o = dump.Options(**options)
+        m = dump.DumpMaster(None, o, filt, outfile=cs)
+        self._cycle(m)
+        return cs.getvalue()
+
     def test_options(self):
         o = dump.Options(verbosity = 2)
         assert o.verbosity == 2
         libpry.raises(AttributeError, dump.Options, nonexistent = 2)
 
     def test_filter(self):
-            cs = StringIO()
-            o = dump.Options(
-                verbosity = 1
-            )
-            m = dump.DumpMaster(None, o, "~u foo", outfile=cs)
-            self._dummy_cycle(m)
-            assert not "GET" in cs.getvalue()
+        assert not "GET" in self._dummy_cycle("~u foo", verbosity=1)
 
     def test_basic(self):
         for i in (1, 2, 3):
-            cs = StringIO()
-            o = dump.Options(
-                verbosity = i
-            )
-            m = dump.DumpMaster(None, o, "~s", outfile=cs)
-            self._dummy_cycle(m)
-            assert "GET" in cs.getvalue()
+            assert "GET" in self._dummy_cycle("~s", verbosity=i)
 
     def test_write(self):
         d = self.tmpdir()
         p = os.path.join(d, "a")
-        o = dump.Options(
-            wfile = p,
-            verbosity = 0
-        )
-        cs = StringIO()
-        m = dump.DumpMaster(None, o, None, outfile=cs)
-        self._dummy_cycle(m)
-        del m
+        self._dummy_cycle(None, wfile=p, verbosity=0)
         assert len(list(flow.FlowReader(open(p)).stream())) == 1
 
     def test_write_err(self):
-        o = dump.Options(
-            wfile = "nonexistentdir/foo",
-            verbosity = 0
+        libpry.raises(
+            dump.DumpError,
+            self._dummy_cycle,
+            None,
+            wfile = "nonexistentdir/foo"
+        )
+
+    def test_request_script(self):
+        ret = self._dummy_cycle(None, request_script="scripts/a", verbosity=1)
+        assert "TESTOK" in ret
+        assert "DEBUG" in ret
+        libpry.raises(
+            dump.DumpError,
+            self._dummy_cycle, None, request_script="nonexistent"
+        )
+        libpry.raises(
+            dump.DumpError,
+            self._dummy_cycle, None, request_script="scripts/err_data"
+        )
+
+    def test_response_script(self):
+        ret = self._dummy_cycle(None, response_script="scripts/a", verbosity=1)
+        assert "TESTOK" in ret
+        assert "DEBUG" in ret
+        libpry.raises(
+            dump.DumpError,
+            self._dummy_cycle, None, response_script="nonexistent"
+        )
+        libpry.raises(
+            dump.DumpError,
+            self._dummy_cycle, None, response_script="scripts/err_data"
         )
-        cs = StringIO()
-        libpry.raises(dump.DumpError, dump.DumpMaster, None, o, None)
 
 
 
@@ -67,5 +81,3 @@ class uDumpMaster(libpry.AutoTree):
 tests = [
     uDumpMaster()
 ]
-
-
diff --git a/test/test_flow.py b/test/test_flow.py
index 9871b090e..35d336e8c 100644
--- a/test/test_flow.py
+++ b/test/test_flow.py
@@ -8,7 +8,7 @@ class uFlow(libpry.AutoTree):
         f = utils.tflow()
         f.response = utils.tresp()
         f.request = f.response.request
-        f, se = f.run_script("scripts/a")
+        se = f.run_script("scripts/a")
         assert "DEBUG" == se.strip()
         assert f.request.host == "TESTOK"