diff options
Diffstat (limited to 'websockify/websocketproxy.py')
| -rwxr-xr-x | websockify/websocketproxy.py | 102 |
1 files changed, 60 insertions, 42 deletions
diff --git a/websockify/websocketproxy.py b/websockify/websocketproxy.py index ac9ed63..fdad0da 100755 --- a/websockify/websocketproxy.py +++ b/websockify/websocketproxy.py @@ -44,8 +44,8 @@ Traffic Legend: """ # Checks if we receive a token, and look # 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) + if self.server.token_plugin: + (self.server.target_host, self.server.target_port) = self.get_target(self.server.token_plugin, self.path) # Connect to the target if self.server.wrap_cmd: @@ -73,15 +73,15 @@ Traffic Legend: 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 - def get_target(self, target_cfg, path): + def get_target(self, target_plugin, path): """ - Parses the path, extracts a token, and looks for a valid - target for that token in the configuration file(s). Sets + Parses the path, extracts a token, and looks up a target + for that token using the token plugin. Sets target_host and target_port if successful """ # The files in targets contain the lines @@ -91,31 +91,16 @@ Traffic Legend: 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") + raise self.server.EClose("Token not present") token = args['token'][0].rstrip('\n') - # target_cfg can be a single config file or directory of - # config files - if os.path.isdir(target_cfg): - cfg_files = [os.path.join(target_cfg, f) - for f in os.listdir(target_cfg)] - else: - cfg_files = [target_cfg] - - targets = {} - for f in cfg_files: - for line in [l.strip() for l in open(f).readlines()]: - if line and not line.startswith('#'): - ttoken, target = line.split(': ') - targets[ttoken] = target.strip() + result_pair = target_plugin.lookup(token) - self.vmsg("Target config: %s" % repr(targets)) - - if token in targets: - return targets[token].split(':') + if result_pair is not None: + return result_pair else: - raise self.EClose("Token '%s' not found" % token) + raise self.server.EClose("Token '%s' not found" % token) def do_proxy(self, target): """ @@ -147,7 +132,7 @@ Traffic Legend: 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']) @@ -195,7 +180,23 @@ class WebSocketProxy(websocket.WebSocketServer): self.wrap_mode = kwargs.pop('wrap_mode', None) self.unix_target = kwargs.pop('unix_target', None) self.ssl_target = kwargs.pop('ssl_target', None) - self.target_cfg = kwargs.pop('target_cfg', None) + + token_plugin = kwargs.pop('token_plugin', None) + token_source = kwargs.pop('token_source', None) + + if token_plugin is not None: + if '.' not in token_plugin: + token_plugin = 'websockify.token_plugins.%s' % token_plugin + + token_plugin_module, token_plugin_cls = token_plugin.rsplit('.', 1) + + __import__(token_plugin_module) + token_plugin_cls = getattr(sys.modules[token_plugin_module], token_plugin_cls) + + self.token_plugin = token_plugin_cls(token_source) + else: + self.token_plugin = None + # Last 3 timestamps command was run self.wrap_times = [0, 0, 0] @@ -251,9 +252,9 @@ class WebSocketProxy(websocket.WebSocketServer): else: dst_string = "%s:%s" % (self.target_host, self.target_port) - if self.target_cfg: - msg = " - proxying from %s:%s to targets in %s" % ( - self.listen_host, self.listen_port, self.target_cfg) + if self.token_plugin: + msg = " - proxying from %s:%s to targets generated by %s" % ( + self.listen_host, self.listen_port, type(self.token_plugin).__name__) else: msg = " - proxying from %s:%s to %s" % ( self.listen_host, self.listen_port, dst_string) @@ -352,20 +353,41 @@ 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("--libserver", action="store_true", + help="use Python library SocketServer engine") parser.add_option("--target-config", metavar="FILE", dest="target_cfg", help="Configuration file containing valid targets " "in the form 'token: host:port' or, alternatively, a " - "directory containing configuration files of this form") - parser.add_option("--libserver", action="store_true", - help="use Python library SocketServer engine") + "directory containing configuration files of this form " + "(DEPRECATED: use `--token-plugin TokenFile --token-source " + " path/to/token/file` instead)") + parser.add_option("--token-plugin", default=None, metavar="PLUGIN", + help="use the given Python class to process tokens " + "into host:port pairs") + parser.add_option("--token-source", default=None, metavar="ARG", + help="an argument to be passed to the token plugin" + "on instantiation") + (opts, args) = parser.parse_args() if opts.verbose: logging.getLogger(WebSocketProxy.log_prefix).setLevel(logging.DEBUG) + if opts.token_source and not opts.token_plugin: + parser.error("You must use --token-plugin to use --token-source") + + # Transform to absolute path as daemon may chdir + if opts.target_cfg: + opts.target_cfg = os.path.abspath(opts.target_cfg) + + if opts.target_cfg: + opts.token_plugin = 'TokenFile' + opts.token_source = opts.target_cfg + del opts.target_cfg + # Sanity checks - if len(args) < 2 and not (opts.target_cfg or opts.unix_target): + if len(args) < 2 and not (opts.token_plugin or opts.unix_target): parser.error("Too few arguments") if sys.argv.count('--'): opts.wrap_cmd = args[1:] @@ -390,7 +412,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.token_plugin: opts.target_host = None opts.target_port = None else: @@ -402,10 +424,6 @@ def websockify_init(): try: opts.target_port = int(opts.target_port) except: parser.error("Error parsing target port") - # Transform to absolute path as daemon may chdir - if opts.target_cfg: - opts.target_cfg = os.path.abspath(opts.target_cfg) - # Create and start the WebSockets proxy libserver = opts.libserver del opts.libserver @@ -456,8 +474,8 @@ class LibProxyServer(ForkingMixIn, HTTPServer): if web: os.chdir(web) - - HTTPServer.__init__(self, (listen_host, listen_port), + + HTTPServer.__init__(self, (listen_host, listen_port), RequestHandlerClass) |