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

1 files changed, 78 insertions, 0 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'])