Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

added --target-api option #157

Closed
wants to merge 1 commit into from
Closed
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 42 additions & 7 deletions websockify/websocketproxy.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@

'''

import signal, socket, optparse, time, os, sys, subprocess, logging
import signal, socket, optparse, time, os, sys, subprocess, logging, urllib2, json
try: from socketserver import ForkingMixIn
except: from SocketServer import ForkingMixIn
try: from http.server import HTTPServer
Expand Down Expand Up @@ -46,6 +46,8 @@ def new_websocket_client(self):
# for a valid target for it then
if self.server.target_cfg:
(self.server.target_host, self.server.target_port) = self.get_target(self.server.target_cfg, self.path)
elif self.server.target_api:
(self.server.target_host, self.server.target_port) = self.get_target_via_api(self.server.target_api, self.path)

# Connect to the target
if self.server.wrap_cmd:
Expand Down Expand Up @@ -73,7 +75,7 @@ def new_websocket_client(self):
if tsock:
tsock.shutdown(socket.SHUT_RDWR)
tsock.close()
if self.verbose:
if self.verbose:
self.log_message("%s:%s: Closed target",
self.server.target_host, self.server.target_port)
raise
Expand Down Expand Up @@ -117,6 +119,28 @@ def get_target(self, target_cfg, path):
else:
raise self.EClose("Token '%s' not found" % token)

def get_target_via_api(self,target_api,path):
"""
Parses the path, extracts a token, and looks for a valid
target for that token in the Installer API. Sets
target_host and target_port if successful
"""

# Extract the token parameter from url
args = parse_qs(urlparse(path)[4]) # 4 is the query from url

if not 'token' in args or not len(args['token']):
raise self.EClose("Token not present")

# installId to be passed in in request/query string
token = args['token'][0].rstrip('\n')
try:
output = urllib2.urlopen(target_api+token).read()
decoded = json.loads(output)
return decoded['ip'],decoded['port']
except:
raise self.EClose("Token '%s' not found" % token)

def do_proxy(self, target):
"""
Proxy client WebSocket to normal target socket.
Expand Down Expand Up @@ -147,7 +171,7 @@ def do_proxy(self, target):

if closed:
# TODO: What about blocking on client socket?
if self.verbose:
if self.verbose:
self.log_message("%s:%s: Client closed connection",
self.server.target_host, self.server.target_port)
raise self.CClose(closed['code'], closed['reason'])
Expand Down Expand Up @@ -196,6 +220,7 @@ def __init__(self, RequestHandlerClass=ProxyRequestHandler, *args, **kwargs):
self.unix_target = kwargs.pop('unix_target', None)
self.ssl_target = kwargs.pop('ssl_target', None)
self.target_cfg = kwargs.pop('target_cfg', None)
self.target_api = kwargs.pop('target_api', None)
# Last 3 timestamps command was run
self.wrap_times = [0, 0, 0]

Expand Down Expand Up @@ -254,6 +279,9 @@ def started(self):
if self.target_cfg:
msg = " - proxying from %s:%s to targets in %s" % (
self.listen_host, self.listen_port, self.target_cfg)
elif self.target_api:
msg = " - proxying from %s:%s to targets (if found) in API" % (
self.listen_host, self.listen_port)
else:
msg = " - proxying from %s:%s to %s" % (
self.listen_host, self.listen_port, dst_string)
Expand Down Expand Up @@ -352,6 +380,11 @@ def websockify_init():
parser.add_option("--prefer-ipv6", "-6",
action="store_true", dest="source_is_ipv6",
help="prefer IPv6 when resolving source_addr")
parser.add_option("--target-api", metavar="FILE",
dest="target_api",
help="JSON Rest API address to check token, token is appended"
"e.g. http://localhost/api/[token]"
"API should return ip and port in JSON response")
parser.add_option("--target-config", metavar="FILE",
dest="target_cfg",
help="Configuration file containing valid targets "
Expand All @@ -365,7 +398,7 @@ def websockify_init():
logging.getLogger(WebSocketProxy.log_prefix).setLevel(logging.DEBUG)

# Sanity checks
if len(args) < 2 and not (opts.target_cfg or opts.unix_target):
if len(args) < 2 and not (opts.target_cfg or opts.unix_target or opts.target_api):
parser.error("Too few arguments")
if sys.argv.count('--'):
opts.wrap_cmd = args[1:]
Expand All @@ -390,7 +423,7 @@ def websockify_init():
try: opts.listen_port = int(opts.listen_port)
except: parser.error("Error parsing listen port")

if opts.wrap_cmd or opts.unix_target or opts.target_cfg:
if opts.wrap_cmd or opts.unix_target or opts.target_cfg or opts.target_api:
opts.target_host = None
opts.target_port = None
else:
Expand Down Expand Up @@ -434,7 +467,9 @@ def __init__(self, RequestHandlerClass=ProxyRequestHandler, **kwargs):
self.unix_target = kwargs.pop('unix_target', None)
self.ssl_target = kwargs.pop('ssl_target', None)
self.target_cfg = kwargs.pop('target_cfg', None)
self.target_api = kwargs.pop('target_api', None)
self.daemon = False
self.target_api = None
self.target_cfg = None

# Server configuration
Expand All @@ -456,8 +491,8 @@ def __init__(self, RequestHandlerClass=ProxyRequestHandler, **kwargs):

if web:
os.chdir(web)
HTTPServer.__init__(self, (listen_host, listen_port),

HTTPServer.__init__(self, (listen_host, listen_port),
RequestHandlerClass)


Expand Down