diff --git a/websocket.py b/websocket.py index 9af0424..8b5f28d 100644 --- a/websocket.py +++ b/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]) diff --git a/websockify b/websockify index ce12194..4dac6fc 100644 --- a/websockify +++ b/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" % ( - self.listen_host, self.listen_port, dst_string)) + msg = " - proxying from %s:%s to %s\n" % ( + 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" % ( - self.listen_host, self.listen_port, dst_string)) + msg = " - proxying from %s:%s to %s\n" % ( + 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" % ( - self.target_host, self.target_port)) + msg = "connecting to: %s:%s" % ( + 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)