diff --git a/dockercloud/api/base.py b/dockercloud/api/base.py index bc904a5..c6bf3d0 100644 --- a/dockercloud/api/base.py +++ b/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 diff --git a/dockercloud/api/events.py b/dockercloud/api/events.py index ea3bd83..f84e4aa 100644 --- a/dockercloud/api/events.py +++ b/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,