diff options
Diffstat (limited to 'websockify')
| -rw-r--r-- | websockify/token_plugins.py | 124 |
1 files changed, 102 insertions, 22 deletions
diff --git a/websockify/token_plugins.py b/websockify/token_plugins.py index 19005d3..d42414e 100644 --- a/websockify/token_plugins.py +++ b/websockify/token_plugins.py @@ -3,6 +3,7 @@ import os import sys import time import re +import json logger = logging.getLogger(__name__) @@ -154,57 +155,136 @@ class JWTTokenApi(BasePlugin): logger.error("package jwcrypto not found, are you sure you've installed it correctly?") return None -class TokenRedis(): - """ - The TokenRedis plugin expects the format of the data in a form of json. + +class TokenRedis(BasePlugin): + """Token plugin based on the Redis in-memory data store. + + The token source is in the format: + + host[:port[:db[:password]]] + + where port, db and password are optional. If port or db are left empty + they will take its default value, ie. 6379 and 0 respectively. + + If your redis server is using the default port (6379) then you can use: + + my-redis-host + + In case you need to authenticate with the redis server and you are using + the default database and port you can use: + + my-redis-host:::verysecretpass + + In the more general case you will use: + + my-redis-host:6380:1:verysecretpass + + The TokenRedis plugin expects the format of the target in one of these two + formats: + + - JSON + + {"host": "target-host:target-port"} + + - Plain text + + target-host:target-port Prepare data with: - redis-cli set hello '{"host":"127.0.0.1:5000"}' + + redis-cli set my-token '{"host": "127.0.0.1:5000"}' Verify with: - redis-cli --raw get hello + + redis-cli --raw get my-token Spawn a test "server" using netcat + nc -l 5000 -v - Note: you have to install also the 'redis' and 'simplejson' modules - pip install redis simplejson + Note: This Token Plugin depends on the 'redis' module, so you have + to install it before using this plugin: + + pip install redis """ def __init__(self, src): try: - # import those ahead of time so we provide error earlier import redis - import simplejson - self._server, self._port = src.split(":") + except ImportError: + logger.error("Unable to load redis module") + sys.exit() + # Default values + self._port = 6379 + self._db = 0 + self._password = None + try: + fields = src.split(":") + if len(fields) == 1: + self._server = fields[0] + elif len(fields) == 2: + self._server, self._port = fields + if not self._port: + self._port = 6379 + elif len(fields) == 3: + self._server, self._port, self._db = fields + if not self._port: + self._port = 6379 + if not self._db: + self._db = 0 + elif len(fields) == 4: + self._server, self._port, self._db, self._password = fields + if not self._port: + self._port = 6379 + if not self._db: + self._db = 0 + if not self._password: + self._password = None + else: + raise ValueError + self._port = int(self._port) + self._db = int(self._db) logger.info("TokenRedis backend initilized (%s:%s)" % (self._server, self._port)) except ValueError: - logger.error("The provided --token-source='%s' is not in an expected format <host>:<port>" % - src) - sys.exit() - except ImportError: - logger.error("package redis or simplejson not found, are you sure you've installed them correctly?") + logger.error("The provided --token-source='%s' is not in the " + "expected format <host>[:<port>[:<db>[:<password>]]]" % + src) sys.exit() def lookup(self, token): try: import redis - import simplejson except ImportError: - logger.error("package redis or simplejson not found, are you sure you've installed them correctly?") + logger.error("package redis not found, are you sure you've installed them correctly?") sys.exit() logger.info("resolving token '%s'" % token) - client = redis.Redis(host=self._server, port=self._port) + client = redis.Redis(host=self._server, port=self._port, + db=self._db, password=self._password) stuff = client.get(token) if stuff is None: return None else: - responseStr = stuff.decode("utf-8") + responseStr = stuff.decode("utf-8").strip() logger.debug("response from redis : %s" % responseStr) - combo = simplejson.loads(responseStr) - (host, port) = combo["host"].split(':') - logger.debug("host: %s, port: %s" % (host,port)) + if responseStr.startswith("{"): + try: + combo = json.loads(responseStr) + host, port = combo["host"].split(":") + except ValueError: + logger.error("Unable to decode JSON token: %s" % + responseStr) + return None + except KeyError: + logger.error("Unable to find 'host' key in JSON token: %s" % + responseStr) + return None + elif re.match(r'\S+:\S+', responseStr): + host, port = responseStr.split(":") + else: + logger.error("Unable to parse token: %s" % responseStr) + return None + logger.debug("host: %s, port: %s" % (host, port)) return [host, port] |