fix the auth error handler error in dockercloud events

This commit is contained in:
tifayuki 2016-04-19 18:57:59 +02:00
commit 94a632aa4c
2 changed files with 16 additions and 7 deletions

View file

@ -264,7 +264,7 @@ class StreamingAPI(BasicObject):
header = {'User-Agent': user_agent} header = {'User-Agent': user_agent}
header.update(dockercloud.auth.get_auth_header()) header.update(dockercloud.auth.get_auth_header())
self.header = [": ".join([key, value]) for key, value in header.items()] 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.open_handler = None
self.message_handler = None self.message_handler = None
self.error_handler = None self.error_handler = None

View file

@ -1,6 +1,7 @@
from __future__ import absolute_import from __future__ import absolute_import
import json import json
import logging
import websocket import websocket
@ -8,32 +9,40 @@ import dockercloud
from .base import StreamingAPI from .base import StreamingAPI
from .exceptions import AuthError from .exceptions import AuthError
logger = logging.getLogger("python-dockercloud")
class Events(StreamingAPI): class Events(StreamingAPI):
def __init__(self): def __init__(self):
endpoint = "events" endpoint = "events"
url = "/".join([dockercloud.stream_host.rstrip("/"), "api", "audit", self._api_version, endpoint.lstrip("/")]) 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) super(self.__class__, self).__init__(url)
def _on_message(self, ws, message): def _on_message(self, ws, message):
logger.info("Websocket Message: %s" % message)
try: try:
event = json.loads(message) event = json.loads(message)
except ValueError: except ValueError:
return 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": if event.get("type") == "auth":
return return
if self.message_handler: if self.message_handler:
self.message_handler(message) 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): def run_forever(self, *args, **kwargs):
while True: while True:
if self.auth_error: if str(dockercloud.auth.get_auth_header()) in self.invaid_auth_headers:
raise AuthError("Not authorized") raise AuthError(self.auth_error)
ws = websocket.WebSocketApp(self.url, header=self.header, ws = websocket.WebSocketApp(self.url, header=self.header,
on_open=self._on_open, on_open=self._on_open,
on_message=self._on_message, on_message=self._on_message,