Merge branch 'token_redis_improvements' of https://github.com/javicacheiro/websockifyHEADmaster

4 files changed, 254 insertions, 23 deletions

diff --git a/setup.py b/setup.py
index ac3f51a..c5c349a 100644
--- a/setup.py
+++ b/setup.py
@@ -32,7 +32,7 @@ setup(name=name,
       install_requires=[
         'numpy', 'requests',
         'jwcrypto',
-        'redis', 'simplejson',
+        'redis',
         ],
       zip_safe=False,
       entry_points={
diff --git a/test-requirements.txt b/test-requirements.txt
index 77a3f5f..4eeff97 100644
--- a/test-requirements.txt
+++ b/test-requirements.txt
@@ -1,4 +1,5 @@
 mock
 nose2
 six
+redis
 wrapt<=1.12.1;python_version<="3.4"
diff --git a/tests/test_token_plugins.py b/tests/test_token_plugins.py
index 3e1fd19..d13d944 100644
--- a/tests/test_token_plugins.py
+++ b/tests/test_token_plugins.py
@@ -203,3 +203,153 @@ class TokenRedisTestCase(unittest.TestCase):
         self.assertIsNotNone(result)
         self.assertEqual(result[0], 'remote_host')
         self.assertEqual(result[1], 'remote_port')
+
+    @patch('redis.Redis')
+    def test_json_token_with_spaces(self, mock_redis):
+        plugin = TokenRedis('127.0.0.1:1234')
+
+        instance = mock_redis.return_value
+        instance.get.return_value = b' {"host": "remote_host:remote_port"} '
+
+        result = plugin.lookup('testhost')
+
+        instance.get.assert_called_once_with('testhost')
+        self.assertIsNotNone(result)
+        self.assertEqual(result[0], 'remote_host')
+        self.assertEqual(result[1], 'remote_port')
+
+    @patch('redis.Redis')
+    def test_text_token(self, mock_redis):
+        plugin = TokenRedis('127.0.0.1:1234')
+
+        instance = mock_redis.return_value
+        instance.get.return_value = b'remote_host:remote_port'
+
+        result = plugin.lookup('testhost')
+
+        instance.get.assert_called_once_with('testhost')
+        self.assertIsNotNone(result)
+        self.assertEqual(result[0], 'remote_host')
+        self.assertEqual(result[1], 'remote_port')
+
+    @patch('redis.Redis')
+    def test_text_token_with_spaces(self, mock_redis):
+        plugin = TokenRedis('127.0.0.1:1234')
+
+        instance = mock_redis.return_value
+        instance.get.return_value = b' remote_host:remote_port '
+
+        result = plugin.lookup('testhost')
+
+        instance.get.assert_called_once_with('testhost')
+        self.assertIsNotNone(result)
+        self.assertEqual(result[0], 'remote_host')
+        self.assertEqual(result[1], 'remote_port')
+
+    @patch('redis.Redis')
+    def test_invalid_token(self, mock_redis):
+        plugin = TokenRedis('127.0.0.1:1234')
+
+        instance = mock_redis.return_value
+        instance.get.return_value = b'{"host": "remote_host:remote_port"   '
+
+        result = plugin.lookup('testhost')
+
+        instance.get.assert_called_once_with('testhost')
+        self.assertIsNone(result)
+
+    def test_src_only_host(self):
+        plugin = TokenRedis('127.0.0.1')
+
+        self.assertEqual(plugin._server, '127.0.0.1')
+        self.assertEqual(plugin._port, 6379)
+        self.assertEqual(plugin._db, 0)
+        self.assertEqual(plugin._password, None)
+
+    def test_src_with_host_port(self):
+        plugin = TokenRedis('127.0.0.1:1234')
+
+        self.assertEqual(plugin._server, '127.0.0.1')
+        self.assertEqual(plugin._port, 1234)
+        self.assertEqual(plugin._db, 0)
+        self.assertEqual(plugin._password, None)
+
+    def test_src_with_host_port_db(self):
+        plugin = TokenRedis('127.0.0.1:1234:2')
+
+        self.assertEqual(plugin._server, '127.0.0.1')
+        self.assertEqual(plugin._port, 1234)
+        self.assertEqual(plugin._db, 2)
+        self.assertEqual(plugin._password, None)
+
+    def test_src_with_host_port_db_pass(self):
+        plugin = TokenRedis('127.0.0.1:1234:2:verysecret')
+
+        self.assertEqual(plugin._server, '127.0.0.1')
+        self.assertEqual(plugin._port, 1234)
+        self.assertEqual(plugin._db, 2)
+        self.assertEqual(plugin._password, 'verysecret')
+
+    def test_src_with_host_empty_port_empty_db_pass(self):
+        plugin = TokenRedis('127.0.0.1:::verysecret')
+
+        self.assertEqual(plugin._server, '127.0.0.1')
+        self.assertEqual(plugin._port, 6379)
+        self.assertEqual(plugin._db, 0)
+        self.assertEqual(plugin._password, 'verysecret')
+
+    def test_src_with_host_empty_port_empty_db_empty_pass(self):
+        plugin = TokenRedis('127.0.0.1:::')
+
+        self.assertEqual(plugin._server, '127.0.0.1')
+        self.assertEqual(plugin._port, 6379)
+        self.assertEqual(plugin._db, 0)
+        self.assertEqual(plugin._password, None)
+
+    def test_src_with_host_empty_port_empty_db_no_pass(self):
+        plugin = TokenRedis('127.0.0.1::')
+
+        self.assertEqual(plugin._server, '127.0.0.1')
+        self.assertEqual(plugin._port, 6379)
+        self.assertEqual(plugin._db, 0)
+        self.assertEqual(plugin._password, None)
+
+    def test_src_with_host_empty_port_no_db_no_pass(self):
+        plugin = TokenRedis('127.0.0.1:')
+
+        self.assertEqual(plugin._server, '127.0.0.1')
+        self.assertEqual(plugin._port, 6379)
+        self.assertEqual(plugin._db, 0)
+        self.assertEqual(plugin._password, None)
+
+    def test_src_with_host_empty_port_db_no_pass(self):
+        plugin = TokenRedis('127.0.0.1::2')
+
+        self.assertEqual(plugin._server, '127.0.0.1')
+        self.assertEqual(plugin._port, 6379)
+        self.assertEqual(plugin._db, 2)
+        self.assertEqual(plugin._password, None)
+
+    def test_src_with_host_port_empty_db_pass(self):
+        plugin = TokenRedis('127.0.0.1:1234::verysecret')
+
+        self.assertEqual(plugin._server, '127.0.0.1')
+        self.assertEqual(plugin._port, 1234)
+        self.assertEqual(plugin._db, 0)
+        self.assertEqual(plugin._password, 'verysecret')
+
+    def test_src_with_host_empty_port_db_pass(self):
+        plugin = TokenRedis('127.0.0.1::2:verysecret')
+
+        self.assertEqual(plugin._server, '127.0.0.1')
+        self.assertEqual(plugin._port, 6379)
+        self.assertEqual(plugin._db, 2)
+        self.assertEqual(plugin._password, 'verysecret')
+
+    def test_src_with_host_empty_port_db_empty_pass(self):
+        plugin = TokenRedis('127.0.0.1::2:')
+
+        self.assertEqual(plugin._server, '127.0.0.1')
+        self.assertEqual(plugin._port, 6379)
+        self.assertEqual(plugin._db, 2)
+        self.assertEqual(plugin._password, None)
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]