fix the auth error handler error in dockercloud events

dockercloud/api/base.py

@@ -264,7 +264,7 @@ class StreamingAPI(BasicObject):
         header = {'User-Agent': user_agent}
         header.update(dockercloud.auth.get_auth_header())
         self.header = [": ".join([key, value]) for key, value in header.items()]
-        logger.info("websocket: %s %s" % (self.url, self.header))
+        logger.info("Websocket: %s %s" % (self.url, self.header))
         self.open_handler = None
         self.message_handler = None
         self.error_handler = None

dockercloud/api/events.py

@@ -1,6 +1,7 @@
 from __future__ import absolute_import
 
 import json
+import logging
 
 import websocket
 
@@ -8,32 +9,40 @@ import dockercloud
 from .base import StreamingAPI
 from .exceptions import AuthError
 
+logger = logging.getLogger("python-dockercloud")
+
 
 class Events(StreamingAPI):
     def __init__(self):
         endpoint = "events"
         url = "/".join([dockercloud.stream_host.rstrip("/"), "api", "audit", self._api_version, endpoint.lstrip("/")])
+        self.invaid_auth_headers = set()
+        self.auth_error = ""
         super(self.__class__, self).__init__(url)
 
     def _on_message(self, ws, message):
+        logger.info("Websocket Message: %s" % message)
         try:
             event = json.loads(message)
         except ValueError:
             return
-
-        if event.get("type") == "error" and event.get("data", {}).get("errorMessage") == "UNAUTHORIZED":
-            self.auth_error = True
-            raise AuthError("Not authorized")
         if event.get("type") == "auth":
             return
 
         if self.message_handler:
             self.message_handler(message)
 
+    def _on_error(self, ws, e):
+        if isinstance(e, websocket._exceptions.WebSocketBadStatusException) and getattr(e, "status_code") == 401:
+            self.auth_error = "Not Authorized"
+            self.invaid_auth_headers.add(str(dockercloud.auth.get_auth_header()))
+
+        super(self.__class__, self)._on_error(ws, e)
+
     def run_forever(self, *args, **kwargs):
         while True:
-            if self.auth_error:
-                raise AuthError("Not authorized")
+            if str(dockercloud.auth.get_auth_header()) in self.invaid_auth_headers:
+                raise AuthError(self.auth_error)
             ws = websocket.WebSocketApp(self.url, header=self.header,
                                         on_open=self._on_open,
                                         on_message=self._on_message,