diff options
-rw-r--r-- | tests/test_token_plugins.py | 54 | ||||
-rw-r--r-- | websockify/token_plugins.py | 45 |
2 files changed, 91 insertions, 8 deletions
diff --git a/tests/test_token_plugins.py b/tests/test_token_plugins.py index 0d3b578..b3a1847 100644 --- a/tests/test_token_plugins.py +++ b/tests/test_token_plugins.py @@ -204,6 +204,60 @@ class TokenRedisTestCase(unittest.TestCase): 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') diff --git a/websockify/token_plugins.py b/websockify/token_plugins.py index 644b6a3..4e92c56 100644 --- a/websockify/token_plugins.py +++ b/websockify/token_plugins.py @@ -174,18 +174,32 @@ class TokenRedis(BasePlugin): my-redis-host:6379:0:verysecretpass - The TokenRedis plugin expects the format of the data in a form of json. + 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' module + 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): @@ -234,11 +248,26 @@ class TokenRedis(BasePlugin): 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 = json.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] |