CONNECT through proxy mode for pathoc

Also, use this to unit test the pathod CONNECT implementation.

4 files changed, 45 insertions, 7 deletions

diff --git a/libpathod/pathoc.py b/libpathod/pathoc.py
index 3384dd11..a2d89aaf 100644
--- a/libpathod/pathoc.py
+++ b/libpathod/pathoc.py
@@ -15,8 +15,23 @@ class Pathoc(tcp.TCPClient):
         )
         self.ssl, self.sni = ssl, sni
 
-    def connect(self):
+    def http_connect(self, connect_to, wfile, rfile):
+        wfile.write(
+                    'CONNECT %s:%s HTTP/1.1\r\n'%tuple(connect_to) +
+                    '\r\n'
+                    )
+        wfile.flush()
+        rfile.readline()
+        headers = http.read_headers(self.rfile)
+
+    def connect(self, connect_to=None):
+        """
+            connect_to: A (host, port) tuple, which will be connected to with an
+            HTTP CONNECT request.
+        """
         tcp.TCPClient.connect(self)
+        if connect_to:
+            self.http_connect(connect_to, self.wfile, self.rfile)
         if self.ssl:
             try:
                 self.convert_to_ssl(sni=self.sni)
diff --git a/pathoc b/pathoc
index 42369b37..527b9a22 100755
--- a/pathoc
+++ b/pathoc
@@ -18,6 +18,11 @@ if __name__ == "__main__":
 
     parser = argparse.ArgumentParser(description='A perverse HTTP client.', parents=[preparser])
     parser.add_argument(
+        "-c", dest="connect_to", type=str, default=False,
+        metavar = "HOST:PORT",
+        help="Issue an HTTP CONNECT to connect to the specified host."
+    )
+    parser.add_argument(
         "-i", dest="sni", type=str, default=False,
         help="SSL Server Name Indication"
     )
@@ -94,11 +99,23 @@ if __name__ == "__main__":
     except ValueError:
         parser.error("Invalid return code specification: %s"%args.ignorecodes)
 
+    if args.connect_to:
+        parts = args.connect_to.split(":")
+        if len(parts) != 2:
+            parser.error("Invalid CONNECT specification: %s"%args.connect_to)
+        try:
+            parts[1] = int(parts[1])
+        except ValueError:
+            parser.error("Invalid CONNECT specification: %s"%args.connect_to)
+        connect_to = parts
+    else:
+        connect_to = None
+
     try:
         for i in range(args.repeat):
             p = pathoc.Pathoc(args.host, port, args.ssl, args.sni)
             try:
-                p.connect()
+                p.connect(connect_to)
             except tcp.NetLibError, v:
                 print >> sys.stderr, str(v)
                 sys.exit(1)
diff --git a/test/test_pathod.py b/test/test_pathod.py
index 0936fce3..0463fc55 100644
--- a/test/test_pathod.py
+++ b/test/test_pathod.py
@@ -150,6 +150,12 @@ class CommonTests(tutils.DaemonTests):
 
 class TestDaemon(CommonTests):
     ssl = False
+    def test_connect(self):
+        v = self.pathoc(r"get:'http://foo.com/p/202':da", connect_to=("localhost", self.d.port), ssl=True)
+        assert v[1] == 202
+
+    def test_connect_err(self):
+        tutils.raises(http.HttpError, self.pathoc, r"get:'http://foo.com/p/202':da", connect_to=("localhost", self.d.port))
 
 
 class TestDaemonSSL(CommonTests):
diff --git a/test/tutils.py b/test/tutils.py
index 0a5455a3..580174d9 100644
--- a/test/tutils.py
+++ b/test/tutils.py
@@ -44,11 +44,11 @@ class DaemonTests:
     def get(self, spec):
         return requests.get(self.d.p(spec), verify=False)
 
-    def pathoc(self, spec, timeout=None):
-        c = pathoc.Pathoc("localhost", self.d.port)
-        c.connect()
-        if self.ssl:
-            c.convert_to_ssl()
+    def pathoc(self, spec, timeout=None, connect_to=None, ssl=None):
+        if ssl is None:
+            ssl = self.ssl
+        c = pathoc.Pathoc("localhost", self.d.port, ssl=ssl)
+        c.connect(connect_to)
         if timeout:
             c.settimeout(timeout)
         return c.request(spec)