Allow websocket redirects

resources/lib/websocket.py

@@ -105,6 +105,19 @@ class WebSocketTimeoutException(WebSocketException):
     pass
 
 
+class WebsocketRedirect(WebSocketException):
+    """
+    WebsocketRedirect will be raised if a status code 301 is returned
+    The Exception will be instantiated with a dict containing all response
+    headers; which should contain the redirect address under the key 'location'
+
+    Access the headers via the attribute headers
+    """
+    def __init__(self, headers):
+        self.headers = headers
+        super(WebsocketRedirect, self).__init__()
+
+
 DEFAULT_TIMEOUT = None
 TRACE_ENABLED = False
 
@@ -162,10 +175,10 @@ def _parse_url(url):
         port = parsed.port
 
     is_secure = False
-    if scheme == "ws":
+    if scheme == "ws" or scheme == 'http':
         if not port:
             port = 80
-    elif scheme == "wss":
+    elif scheme == "wss" or scheme == 'https':
         is_secure = True
         if not port:
             port = 443
@@ -500,6 +513,9 @@ class WebSocket(object):
             LOG.debug("-----------------------")
 
         status, resp_headers = self._read_headers()
+        if status == 301:
+            # Redirect
+            raise WebsocketRedirect(resp_headers)
         if status != 101:
             self.close()
             raise WebSocketException("Handshake Status %d" % status)

resources/lib/websocket_client.py

@@ -18,6 +18,7 @@ class WebSocket(backgroundthread.KillableThread):
 
     def __init__(self):
         self.ws = None
+        self.redirect_uri = None
         super(WebSocket, self).__init__()
 
     def process(self, opcode, message):
@@ -91,6 +92,16 @@ class WebSocket(backgroundthread.KillableThread):
                              self.__class__.__name__)
                     self.ws = None
                     app.APP.monitor.waitForAbort(1)
+                except websocket.WebsocketRedirect as e:
+                    LOG.info('301 redirect detected')
+                    self.redirect_uri = e.headers.get('location', e.headers.get('Location'))
+                    if self.redirect_uri:
+                        self.redirect_uri.decode('utf-8')
+                    counter += 1
+                    if counter >= 10:
+                        LOG.info('%s: Repeated WebsocketRedirect detected. Stopping now',
+                                 self.__class__.__name__)
+                        break
                 except websocket.WebSocketException as e:
                     LOG.info('%s: WebSocketException: %s',
                              self.__class__.__name__, e)
@@ -141,6 +152,10 @@ class PMS_Websocket(WebSocket):
                 app.SYNC.background_sync_disabled)
 
     def getUri(self):
+        if self.redirect_uri:
+            uri = self.redirect_uri
+            self.redirect_uri = None
+        else:
             server = app.CONN.server
             # Get the appropriate prefix for the websocket
             if server.startswith('https'):
@@ -209,6 +224,10 @@ class Alexa_Websocket(WebSocket):
                 app.ACCOUNT.restricted_user)
 
     def getUri(self):
+        if self.redirect_uri:
+            uri = self.redirect_uri
+            self.redirect_uri = None
+        else:
             uri = ('wss://pubsub.plex.tv/sub/websockets/%s/%s?X-Plex-Token=%s'
                    % (app.ACCOUNT.plex_user_id,
                       v.PKC_MACHINE_IDENTIFIER,