crabapples
/
websockify
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

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

This commit is contained in:
Pierre Ossman 2023-01-20 17:01:32 +01:00
parent 7f53e9c22c 3d2e93aeb0
commit a1346552fb
4 changed files with 254 additions and 23 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
setup.py
View File

@ -32,7 +32,7 @@ setup(name=name,
install_requires=[
'numpy', 'requests',
'jwcrypto',
'redis', 'simplejson',
'redis',
],
zip_safe=False,
entry_points={

1
test-requirements.txt
View File

@ -1,4 +1,5 @@
mock
nose2
six
redis
wrapt<=1.12.1;python_version<="3.4"

150
tests/test_token_plugins.py
View File

@ -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)

124
websockify/token_plugins.py
View File

@ -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]