Merge with kanaka's mainstream master branch

websocket.py

@@ -166,7 +166,7 @@ Sec-WebSocket-Accept: %s\r
     #
 
     @staticmethod
-    def socket(host, port=None, connect=False, prefer_ipv6=False, unix_socket=None):
+    def socket(host, port=None, connect=False, prefer_ipv6=False, unix_socket=None, use_ssl=False):
         """ Resolve a host (and optional port) to an IPv4 or IPv6
         address. Create a socket. Bind to it if listen is set,
         otherwise connect to it. Return the socket.
@@ -176,6 +176,10 @@ Sec-WebSocket-Accept: %s\r
             host = None
         if connect and not (port or unix_socket):
             raise Exception("Connect mode requires a port")
+        if use_ssl and not ssl:
+            raise Exception("SSL socket requested but Python SSL module not loaded.");
+        if not connect and use_ssl:
+            raise Exception("SSL only supported in connect mode (for now)")
         if not connect:
             flags = flags | socket.AI_PASSIVE
             
@@ -183,13 +187,15 @@ Sec-WebSocket-Accept: %s\r
             addrs = socket.getaddrinfo(host, port, 0, socket.SOCK_STREAM,
                     socket.IPPROTO_TCP, flags)
             if not addrs:
-                raise Exception("Could resolve host '%s'" % host)
+                raise Exception("Could not resolve host '%s'" % host)
             addrs.sort(key=lambda x: x[0])
             if prefer_ipv6:
                 addrs.reverse()
             sock = socket.socket(addrs[0][0], addrs[0][1])
             if connect:
                 sock.connect(addrs[0][4])
+                if use_ssl:
+                    sock = ssl.wrap_socket(sock)
             else:
                 sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
                 sock.bind(addrs[0][4])

websockify

@@ -13,9 +13,9 @@ as taken from http://docs.python.org/dev/library/ssl.html#certificates
 
 import socket, optparse, time, os, sys, subprocess
 from select import select
-from websocket import WebSocketServer
+import websocket
 
-class WebSocketProxy(WebSocketServer):
+class WebSocketProxy(websocket.WebSocketServer):
     """
     Proxy traffic to and from a WebSockets client to a normal TCP
     socket server target. All traffic to/from the client is base64
@@ -43,6 +43,7 @@ Traffic Legend:
         self.target_port    = kwargs.pop('target_port')
         self.wrap_cmd       = kwargs.pop('wrap_cmd')
         self.wrap_mode      = kwargs.pop('wrap_mode')
+        self.ssl_target     = kwargs.pop('ssl_target')
         # Last 3 timestamps command was run
         self.wrap_times    = [0, 0, 0]
 
@@ -71,7 +72,7 @@ Traffic Legend:
                 "REBIND_OLD_PORT": str(kwargs['listen_port']),
                 "REBIND_NEW_PORT": str(self.target_port)})
 
-        WebSocketServer.__init__(self, *args, **kwargs)
+        websocket.WebSocketServer.__init__(self, *args, **kwargs)
 
     def run_wrap_cmd(self):
         print("Starting '%s'" % " ".join(self.wrap_cmd))
@@ -89,13 +90,21 @@ Traffic Legend:
         # know when the wrapped command exits
         if self.wrap_cmd:
             dst_string = self.unix_socket or "'%s' (port %s)" % (" ".join(self.wrap_cmd), self.target_port)
-            print("  - proxying from %s:%s to %s\n" % (
+            msg = "  - proxying from %s:%s to %s\n" % (
-                    self.listen_host, self.listen_port, dst_string))
+                    self.listen_host, self.listen_port, dst_string)
             self.run_wrap_cmd()
         else:
             dst_string = self.unix_socket or "%s:%s" % (self.target_host, self.target_port)
-            print("  - proxying from %s:%s to %s\n" % (
+            msg = "  - proxying from %s:%s to %s\n" % (
-                    self.listen_host, self.listen_port, dst_string))
+                    self.listen_host, self.listen_port, dst_string)
+
+        if self.ssl_target:
+            msg += " (using SSL)"
+
+        print(msg + "\n")
+
+        if self.wrap_cmd:
+            self.run_wrap_cmd()
 
     def poll(self):
         # If we are wrapping a command, check it's status
@@ -139,13 +148,20 @@ Traffic Legend:
         """
         # Connect to the target
         if self.unix_socket:
-            self.msg("connecting to unix socket : %s" % self.unix_socket)
+            msg = "connecting to unix socket : %s" % self.unix_socket
         else:
-            self.msg("connecting to: %s:%s" % (
+            msg = "connecting to: %s:%s" % (
-                                    self.target_host, self.target_port))
+                                    self.target_host, self.target_port)
         
         tsock = self.socket(self.target_host, self.target_port,
-                            connect=True, unix_socket=self.unix_socket)
+                            connect=True, use_ssl=self.ssl_target, unix_socket=self.unix_socket)
+         
+        if self.ssl_target:
+            msg += " (using SSL)"
+        self.msg(msg)
+
+        tsock = self.socket(self.target_host, self.target_port,
+                connect=True, use_ssl=self.ssl_target)
 
         if self.verbose and not self.daemon:
             print(self.traffic_legend)
@@ -240,7 +256,9 @@ def websockify_init():
     parser.add_option("--key", default=None,
             help="SSL key file (if separate from cert)")
     parser.add_option("--ssl-only", action="store_true",
-            help="disallow non-encrypted connections")
+            help="disallow non-encrypted client connections")
+    parser.add_option("--ssl-target", action="store_true",
+            help="connect to SSL target as SSL client")
     parser.add_option("--web", default=None, metavar="DIR",
             help="run webserver on same port. Serve files from DIR.")
     parser.add_option("--wrap-mode", default="exit", metavar="MODE",
@@ -259,6 +277,9 @@ def websockify_init():
         if len(args) > 2:
             parser.error("Too many arguments")
 
+    if not websocket.ssl and opts.ssl_target:
+        parser.error("SSL target requested and Python SSL module not loaded.");
+        
     if opts.ssl_only and not os.path.exists(opts.cert):
         parser.error("SSL only and %s not found" % opts.cert)