Updated websocketproxy to support Cookies

websockify/websocketproxy.py

@@ -140,20 +140,39 @@ Traffic Legend:
         # The files in targets contain the lines
         # in the form of token: host:port
 
+        token = ''
+
+		# Extract the token parameter from cookie
+        if 'Cookie' in self.headers:
+            cookiestr = self.headers['Cookie']
+            cookiestr.replace( 'Cookie: ', '', 1 )
+            cookies = cookiestr.split('; ')
+
+            for cookie in cookies:
+                cookie = cookie.split('=')
+                if 'token' in cookie[0]:
+                    token = cookie[1]
+                    break
+
         # Extract the token parameter from url
-        args = parse_qs(urlparse(path)[4]) # 4 is the query from url
+        if token == '':
+            args = parse_qs(urlparse(path)[4]) # 4 is the query from url
 
-        if not 'token' in args or not len(args['token']):
-            raise self.server.EClose("Token not present")
+            if not 'token' in args or not len(args['token']):
+                raise self.server.EClose("Token not present")
 
-        token = args['token'][0].rstrip('\n')
+            token = args['token'][0].rstrip('\n')
 
-        result_pair = target_plugin.lookup(token)
+        if token == '':
+            raise self.server.EClose( "Empty Token defined" )
+
+        # Search for the token
+        result_pair = target_plugin.lookup( token )
 
         if result_pair is not None:
             return result_pair
         else:
-            raise self.server.EClose("Token '%s' not found" % token)
+            raise self.server.EClose( "Token '%s' not found" % token )
 
     def do_proxy(self, target):
         """