Shift SSL parameters into Pathoc class

3 files changed, 41 insertions, 10 deletions

diff --git a/libpathod/pathoc.py b/libpathod/pathoc.py
index dcee353f..3384dd11 100644
--- a/libpathod/pathoc.py
+++ b/libpathod/pathoc.py
@@ -7,12 +7,21 @@ class PathocError(Exception): pass
 
 
 class Pathoc(tcp.TCPClient):
-    def __init__(self, host, port):
+    def __init__(self, host, port, ssl=None, sni=None):
         tcp.TCPClient.__init__(self, host, port)
         self.settings = dict(
             staticdir = os.getcwd(),
             unconstrained_file_access = True,
         )
+        self.ssl, self.sni = ssl, sni
+
+    def connect(self):
+        tcp.TCPClient.connect(self)
+        if self.ssl:
+            try:
+                self.convert_to_ssl(sni=self.sni)
+            except tcp.NetLibError, v:
+                raise PathocError(str(v))
 
     def request(self, spec):
         """
diff --git a/pathoc b/pathoc
index 1d3f6004..42369b37 100755
--- a/pathoc
+++ b/pathoc
@@ -96,18 +96,12 @@ if __name__ == "__main__":
 
     try:
         for i in range(args.repeat):
-            p = pathoc.Pathoc(args.host, port)
+            p = pathoc.Pathoc(args.host, port, args.ssl, args.sni)
             try:
                 p.connect()
             except tcp.NetLibError, v:
                 print >> sys.stderr, str(v)
                 sys.exit(1)
-            if args.ssl:
-                try:
-                    p.convert_to_ssl(sni=args.sni)
-                except tcp.NetLibError, v:
-                    print "\n>> %s"%v
-                    continue
             if args.timeout:
                 p.settimeout(args.timeout)
             for spec in args.request:
diff --git a/test/test_pathoc.py b/test/test_pathoc.py
index ca25be58..38d3754a 100644
--- a/test/test_pathoc.py
+++ b/test/test_pathoc.py
@@ -3,10 +3,11 @@ from libpathod import pathoc, test, version
 import tutils
 
 
-class TestDaemon:
+class _TestDaemon:
     @classmethod
     def setUpAll(self):
         self.d = test.Daemon(
+            ssl=self.ssl,
             staticdir=tutils.test_data.path("data"),
             anchors=[("/anchor/.*", "202")]
         )
@@ -19,11 +20,34 @@ class TestDaemon:
         self.d.clear_log()
 
     def test_info(self):
-        c = pathoc.Pathoc("127.0.0.1", self.d.port)
+        c = pathoc.Pathoc(
+            "127.0.0.1",
+            self.d.port,
+            ssl = self.ssl
+        )
         c.connect()
         _, _, _, _, content = c.request("get:/api/info")
         assert tuple(json.loads(content)["version"]) == version.IVERSION
 
+
+class TestDaemonSSL(_TestDaemon):
+    ssl = True
+    def test_sni(self):
+        c = pathoc.Pathoc(
+            "127.0.0.1",
+            self.d.port,
+            ssl = True,
+            sni = "foobar.com"
+        )
+        c.connect()
+        c.request("get:/p/200")
+        _, _, _, _, content = c.request("get:/api/log")
+        d = json.loads(content)
+        assert d["log"][0]["request"]["sni"] == "foobar.com"
+
+
+class TestDaemon(_TestDaemon):
+    ssl = False
     def tval(self, requests, showreq=False, showresp=False, explain=False, hexdump=False, timeout=None, ignorecodes=None, ignoretimeout=None):
         c = pathoc.Pathoc("127.0.0.1", self.d.port)
         c.connect()
@@ -43,6 +67,10 @@ class TestDaemon:
             )
         return s.getvalue()
 
+    def test_ssl_error(self):
+        c = pathoc.Pathoc("127.0.0.1", self.d.port, ssl = True)
+        tutils.raises("ssl handshake", c.connect)
+
     def test_ignorecodes(self):
         assert "200" in self.tval(["get:'/p/200:b@1'"])
         assert "200" not in self.tval(["get:'/p/200:b@1'"], ignorecodes=[200])