From e23d4e337cd3f4a83d160220af4b09e446eed61a Mon Sep 17 00:00:00 2001
From: Javier Cacheiro <javier.cacheiro.lopez@cesga.es>
Date: Fri, 22 Apr 2022 13:54:27 +0200
Subject: Token Redis: Support both json and plain text tokens

---
 websockify/token_plugins.py | 45 +++++++++++++++++++++++++++++++++++++--------
 1 file changed, 37 insertions(+), 8 deletions(-)

(limited to 'websockify/token_plugins.py')

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]
 
 
-- 
cgit v1.2.1