Adding support for proxying from a unix socket
This commit is contained in:
Karim Allah Ahmed
parent
e1c206b315
commit
c8018f29c9
|
|
@ -4,3 +4,5 @@
|
|||
other/.lein-deps-sum
|
||||
other/classes
|
||||
other/lib
|
||||
.project
|
||||
.pydevproject
|
||||
|
|
|
|||
54
websocket.py
54
websocket.py
|
|
@ -75,6 +75,7 @@ class WebSocketServer(object):
|
|||
|
||||
buffer_size = 65536
|
||||
|
||||
|
||||
server_handshake_hixie = """HTTP/1.1 101 Web Socket Protocol Handshake\r
|
||||
Upgrade: WebSocket\r
|
||||
Connection: Upgrade\r
|
||||
|
|
@ -103,7 +104,7 @@ Sec-WebSocket-Accept: %s\r
|
|||
def __init__(self, listen_host='', listen_port=None, source_is_ipv6=False,
|
||||
verbose=False, cert='', key='', ssl_only=None,
|
||||
daemon=False, record='', web='',
|
||||
run_once=False, timeout=0):
|
||||
run_once=False, timeout=0, unix=None):
|
||||
|
||||
# settings
|
||||
self.verbose = verbose
|
||||
|
|
@ -113,6 +114,8 @@ Sec-WebSocket-Accept: %s\r
|
|||
self.daemon = daemon
|
||||
self.run_once = run_once
|
||||
self.timeout = timeout
|
||||
|
||||
self.unix_socket = unix
|
||||
|
||||
self.launch_time = time.time()
|
||||
self.ws_connection = False
|
||||
|
|
@ -163,7 +166,7 @@ Sec-WebSocket-Accept: %s\r
|
|||
#
|
||||
|
||||
@staticmethod
|
||||
def socket(host, port=None, connect=False, prefer_ipv6=False):
|
||||
def socket(host, port=None, connect=False, prefer_ipv6=False, unix_socket=None):
|
||||
""" 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.
|
||||
|
|
@ -171,24 +174,30 @@ Sec-WebSocket-Accept: %s\r
|
|||
flags = 0
|
||||
if host == '':
|
||||
host = None
|
||||
if connect and not port:
|
||||
if connect and not (port or unix_socket):
|
||||
raise Exception("Connect mode requires a port")
|
||||
if not connect:
|
||||
flags = flags | socket.AI_PASSIVE
|
||||
addrs = socket.getaddrinfo(host, port, 0, socket.SOCK_STREAM,
|
||||
socket.IPPROTO_TCP, flags)
|
||||
if not addrs:
|
||||
raise Exception("Could 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])
|
||||
else:
|
||||
sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
||||
sock.bind(addrs[0][4])
|
||||
sock.listen(100)
|
||||
|
||||
if not unix_socket:
|
||||
addrs = socket.getaddrinfo(host, port, 0, socket.SOCK_STREAM,
|
||||
socket.IPPROTO_TCP, flags)
|
||||
if not addrs:
|
||||
raise Exception("Could 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])
|
||||
else:
|
||||
sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
||||
sock.bind(addrs[0][4])
|
||||
sock.listen(100)
|
||||
else:
|
||||
sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
|
||||
sock.connect(unix_socket)
|
||||
|
||||
return sock
|
||||
|
||||
@staticmethod
|
||||
|
|
@ -569,10 +578,10 @@ Sec-WebSocket-Accept: %s\r
|
|||
- Send a WebSockets handshake server response.
|
||||
- Return the socket for this WebSocket client.
|
||||
"""
|
||||
|
||||
stype = ""
|
||||
|
||||
ready = select.select([sock], [], [], 3)[0]
|
||||
|
||||
|
||||
if not ready:
|
||||
raise self.EClose("ignoring socket not ready")
|
||||
# Peek, but do not read the data so that we have a opportunity
|
||||
|
|
@ -751,6 +760,11 @@ Sec-WebSocket-Accept: %s\r
|
|||
self.start_time = int(time.time()*1000)
|
||||
|
||||
# handler process
|
||||
|
||||
dst_string = self.unix_socket or "%s:%s" % (self.target_host, self.target_port)
|
||||
dst_string = self.unix_socket or "'%s' (port %s)" % (" ".join(self.wrap_cmd), self.target_port)
|
||||
|
||||
|
||||
try:
|
||||
try:
|
||||
self.client = self.do_handshake(startsock, address)
|
||||
|
|
@ -848,7 +862,7 @@ Sec-WebSocket-Accept: %s\r
|
|||
continue
|
||||
else:
|
||||
raise
|
||||
|
||||
|
||||
if self.run_once:
|
||||
# Run in same process if run_once
|
||||
self.top_new_client(startsock, address)
|
||||
|
|
|
|||
|
|
@ -88,14 +88,14 @@ Traffic Legend:
|
|||
# Need to call wrapped command after daemonization so we can
|
||||
# know when the wrapped command exits
|
||||
if self.wrap_cmd:
|
||||
print(" - proxying from %s:%s to '%s' (port %s)\n" % (
|
||||
self.listen_host, self.listen_port,
|
||||
" ".join(self.wrap_cmd), self.target_port))
|
||||
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))
|
||||
self.run_wrap_cmd()
|
||||
else:
|
||||
print(" - proxying from %s:%s to %s:%s\n" % (
|
||||
self.listen_host, self.listen_port,
|
||||
self.target_host, self.target_port))
|
||||
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))
|
||||
|
||||
def poll(self):
|
||||
# If we are wrapping a command, check it's status
|
||||
|
|
@ -137,12 +137,15 @@ Traffic Legend:
|
|||
"""
|
||||
Called after a new WebSocket connection has been established.
|
||||
"""
|
||||
|
||||
# Connect to the target
|
||||
self.msg("connecting to: %s:%s" % (
|
||||
self.target_host, self.target_port))
|
||||
if self.unix_socket:
|
||||
self.msg("connecting to unix socket : %s" % self.unix_socket)
|
||||
else:
|
||||
self.msg("connecting to: %s:%s" % (
|
||||
self.target_host, self.target_port))
|
||||
|
||||
tsock = self.socket(self.target_host, self.target_port,
|
||||
connect=True)
|
||||
connect=True, unix_socket=self.unix_socket)
|
||||
|
||||
if self.verbose and not self.daemon:
|
||||
print(self.traffic_legend)
|
||||
|
|
@ -223,6 +226,8 @@ def websockify_init():
|
|||
help="verbose messages and per frame traffic")
|
||||
parser.add_option("--record",
|
||||
help="record sessions to FILE.[session_number]", metavar="FILE")
|
||||
parser.add_option("--unix",
|
||||
help="unix socket to proxy network from", metavar="FILE")
|
||||
parser.add_option("--daemon", "-D",
|
||||
dest="daemon", action="store_true",
|
||||
help="become a daemon (background process)")
|
||||
|
|
@ -245,7 +250,7 @@ def websockify_init():
|
|||
(opts, args) = parser.parse_args()
|
||||
|
||||
# Sanity checks
|
||||
if len(args) < 2:
|
||||
if len(args) < 1:
|
||||
parser.error("Too few arguments")
|
||||
if sys.argv.count('--'):
|
||||
opts.wrap_cmd = args[1:]
|
||||
|
|
@ -270,12 +275,15 @@ def websockify_init():
|
|||
opts.target_host = None
|
||||
opts.target_port = None
|
||||
else:
|
||||
if args[1].count(':') > 0:
|
||||
opts.target_host, opts.target_port = args[1].rsplit(':', 1)
|
||||
if hasattr(opts, 'unix'):
|
||||
opts.target_host = opts.target_port = None
|
||||
else:
|
||||
parser.error("Error parsing target")
|
||||
try: opts.target_port = int(opts.target_port)
|
||||
except: parser.error("Error parsing target port")
|
||||
if args[1].count(':') > 0:
|
||||
opts.target_host, opts.target_port = args[1].rsplit(':', 1)
|
||||
else:
|
||||
parser.error("Error parsing target")
|
||||
try: opts.target_port = int(opts.target_port)
|
||||
except: parser.error("Error parsing target port")
|
||||
|
||||
# Create and start the WebSockets proxy
|
||||
server = WebSocketProxy(**opts.__dict__)
|
||||
|
|
|
|||
Loading…
Reference in New Issue