Introduce Token Pluginsfeature/token-plugins

Token plugins provide a generic interface for transforming a token
into a `(host, port)` tuple.

The plugin name is specified using the '--token-plugin' option,
and may either be the name of a class from `websockify.token_plugins`,
or a fully qualified python path to the token plugin class (see below).

An optional plugin parameter can be specified using the '--token-source'
option (a value of `None` will be used if no '--token-source' option is
passed).

Token plugins should inherit from `websockify.token_plugins.BasePlugin`,
and should implement the `lookup(token)` method.  The value of the
'--token-source' option is available as `self.source`.

Several plugins are included by default.  The `ReadOnlyTokenFile`
and `TokenFile` plugins implement functionality from '--target-config'
(with the former only reading the file(s) once, and the latter reading
them every time).  The 'BaseTokenAPI' plugin fetches the value from
an API, returning the result of `process_result(response_object)`.
By default, `process_result` simply returns the text of the response,
but may be overriden.  The `JSONTokenAPI` does just this, returning
the 'host' and 'port' values from the response JSON object.

The old '--target-config' option is now deprecated, and maps to the
`TokenFile` plugin under the hood.

Also-Authored-By: James Portman (@james-portman)

Closes #157

2 files changed, 138 insertions, 42 deletions

diff --git a/websockify/token_plugins.py b/websockify/token_plugins.py
new file mode 100644
index 0000000..fcb6080
--- /dev/null
+++ b/websockify/token_plugins.py
@@ -0,0 +1,78 @@
+import os
+
+class BasePlugin(object):
+    def __init__(self, src):
+        self.source = src
+
+    def lookup(self, token):
+        return None
+
+
+class ReadOnlyTokenFile(BasePlugin):
+    # source is a token file with lines like
+    #   token: host:port
+    # or a directory of such files
+    def _load_targets(self):
+        if os.path.isdir(self.source):
+            cfg_files = [os.path.join(self.source, f) for
+                         f in os.listdir(self.source)]
+        else:
+            cfg_files = [self.source]
+
+        self._targets = {}
+        for f in cfg_files:
+            for line in [l.strip() for l in open(f).readlines()]:
+                if line and not line.startswith('#'):
+                    tok, target = line.split(': ')
+                    self._targets[tok] = target.strip().split(':')
+
+    def lookup(self, token):
+        if self._targets is None:
+            self._load_targets()
+
+        if token in self._targets:
+            return self._targets[token]
+        else:
+            return None
+
+
+# the above one is probably more efficient, but this one is
+# more backwards compatible (although in most cases
+# ReadOnlyTokenFile should suffice)
+class TokenFile(ReadOnlyTokenFile):
+    # source is a token file with lines like
+    #   token: host:port
+    # or a directory of such files
+    def lookup(self, token):
+        self._load_targets()
+
+        return super(TokenFile, self).lookup(token)
+
+
+class BaseTokenAPI(BasePlugin):
+    # source is a url with a '%s' in it where the token
+    # should go
+
+    # we import things on demand so that other plugins
+    # in this file can be used w/o unecessary dependencies
+
+    def process_result(self, resp):
+        return resp.text.split(':')
+
+    def lookup(self, token):
+        import requests
+
+        resp = requests.get(self.source % token)
+
+        if resp.ok:
+            return self.process_result(resp)
+        else:
+            return None
+
+
+class JSONTokenApi(BaseTokenAPI):
+    # source is a url with a '%s' in it where the token
+    # should go
+
+    def process_result(self, resp):
+        return (resp.json['host'], resp.json['port'])
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)