Achievement Unlocked: Proxy Chain

3 files changed, 74 insertions, 8 deletions

diff --git a/test/test_fuzzing.py b/test/test_fuzzing.py
index ba7b751c..646ce5c1 100644
--- a/test/test_fuzzing.py
+++ b/test/test_fuzzing.py
@@ -32,8 +32,8 @@ class TestFuzzy(tservers.HTTPProxTest):
         assert p.request(req%self.server.port).status_code == 502
 
     def test_upstream_disconnect(self):
-        req = r'200:d0:h"Date"="Sun, 03 Mar 2013 04:00:00 GMT"'
+        req = r'200:d0'
         p = self.pathod(req)
-        assert p.status_code == 400
+        assert p.status_code == 502
 
 
diff --git a/test/test_protocol_http.py b/test/test_protocol_http.py
index bf569fb7..76d192de 100644
--- a/test/test_protocol_http.py
+++ b/test/test_protocol_http.py
@@ -1,7 +1,7 @@
 from libmproxy import proxy  # FIXME: Remove
 from libmproxy.protocol.http import *
 from cStringIO import StringIO
-import tutils
+import tutils, tservers
 
 
 def test_HttpAuthenticationError():
@@ -84,3 +84,30 @@ class TestHTTPResponse:
         assert r.code == 200
         assert r.content == ""
         tutils.raises("Invalid server response: 'content", HTTPResponse.from_stream, s, "GET")
+
+
+class TestProxyChaining(tservers.HTTPChainProxyTest):
+    def test_all(self):
+        self.chain[1].tmaster.replacehooks.add("~q", "foo", "bar") # replace in request
+        self.chain[0].tmaster.replacehooks.add("~q", "foo", "oh noes!")
+        self.proxy.tmaster.replacehooks.add("~q", "bar", "baz")
+        self.chain[0].tmaster.replacehooks.add("~s", "baz", "ORLY")  # replace in response
+
+        p = self.pathoc()
+        req = p.request("get:'%s/p/418:b\"foo\"'" % self.server.urlbase)
+        assert req.content == "ORLY"
+        assert req.status_code == 418
+
+        self.chain[0].tmaster.replacehooks.clear()
+        self.chain[1].tmaster.replacehooks.clear()
+        self.proxy.tmaster.replacehooks.clear()
+
+
+class TestProxyChainingSSL(tservers.HTTPChainProxyTest):
+    ssl = True
+
+    def test_full(self):
+        p = self.pathoc()
+        req = p.request("get:'/p/418:b@100'")
+        assert len(req.content) == 100
+        assert req.status_code == 418
\ No newline at end of file
diff --git a/test/tservers.py b/test/tservers.py
index f9008cd6..156796c4 100644
--- a/test/tservers.py
+++ b/test/tservers.py
@@ -21,13 +21,12 @@ def errapp(environ, start_response):
 
 
 class TestMaster(flow.FlowMaster):
-    def __init__(self, testq, config):
+    def __init__(self, config):
         s = proxy.ProxyServer(config, 0)
         state = flow.State()
         flow.FlowMaster.__init__(self, s, state)
         self.apps.add(testapp, "testapp", 80)
         self.apps.add(errapp, "errapp", 80)
-        self.testq = testq
         self.clear_log()
         self.start_app(APP_HOST, APP_PORT, False)
 
@@ -51,6 +50,8 @@ class ProxyThread(threading.Thread):
     def __init__(self, tmaster):
         threading.Thread.__init__(self)
         self.tmaster = tmaster
+        self.name = "ProxyThread (%s:%s)" % (tmaster.server.address.host, tmaster.server.address.port)
+
         controller.should_exit = False
 
     @property
@@ -68,7 +69,7 @@ class ProxyThread(threading.Thread):
         self.tmaster.shutdown()
 
 
-class ProxTestBase:
+class ProxTestBase(object):
     # Test Configuration
     ssl = None
     ssloptions = False
@@ -79,7 +80,6 @@ class ProxTestBase:
     masterclass = TestMaster
     @classmethod
     def setupAll(cls):
-        cls.tqueue = Queue.Queue()
         cls.server = libpathod.test.Daemon(ssl=cls.ssl, ssloptions=cls.ssloptions)
         cls.server2 = libpathod.test.Daemon(ssl=cls.ssl, ssloptions=cls.ssloptions)
         pconf = cls.get_proxy_config()
@@ -89,7 +89,7 @@ class ProxTestBase:
             authenticator = cls.authenticator,
             **pconf
         )
-        tmaster = cls.masterclass(cls.tqueue, config)
+        tmaster = cls.masterclass(config)
         cls.proxy = ProxyThread(tmaster)
         cls.proxy.start()
 
@@ -249,5 +249,44 @@ class ReverseProxTest(ProxTestBase):
         return p.request(q)
 
 
+class ChainProxTest(ProxTestBase):
+    """
+    Chain n instances of mitmproxy in a row - because we can.
+    """
+    n = 2
+    chain = []
+    chain_config = [lambda: proxy.ProxyConfig(
+        cacert = tutils.test_data.path("data/serverkey.pem")
+    )] * n
+
+    @classmethod
+    def setupAll(cls):
+        super(ChainProxTest, cls).setupAll()
+        for i in range(cls.n):
+            config = cls.chain_config[i]()
+            config.forward_proxy = ("http", "127.0.0.1",
+                                    cls.proxy.port if i == 0 else
+                                    cls.chain[-1].port
+            )
+            tmaster = cls.masterclass(config)
+            cls.chain.append(ProxyThread(tmaster))
+            cls.chain[-1].start()
+
+    @classmethod
+    def teardownAll(cls):
+        super(ChainProxTest, cls).teardownAll()
+        for p in cls.chain:
+            p.tmaster.server.shutdown()
 
 
+class HTTPChainProxyTest(ChainProxTest):
+    def pathoc(self, sni=None):
+        """
+            Returns a connected Pathoc instance.
+        """
+        p = libpathod.pathoc.Pathoc(("localhost", self.chain[-1].port), ssl=self.ssl, sni=sni)
+        if self.ssl:
+            p.connect(("127.0.0.1", self.server.port))
+        else:
+            p.connect()
+        return p
\ No newline at end of file