| # -*- coding: utf-8 -*- |
| import ssl |
| |
| from tornado import iostream, escape |
| from ambari_ws4py.client import WebSocketBaseClient |
| from ambari_ws4py.exc import HandshakeError |
| |
| __all__ = ['TornadoWebSocketClient'] |
| |
| class TornadoWebSocketClient(WebSocketBaseClient): |
| def __init__(self, url, protocols=None, extensions=None, |
| io_loop=None, ssl_options=None, headers=None, exclude_headers=None): |
| """ |
| .. code-block:: python |
| |
| from tornado import ioloop |
| |
| class MyClient(TornadoWebSocketClient): |
| def opened(self): |
| for i in range(0, 200, 25): |
| self.send("*" * i) |
| |
| def received_message(self, m): |
| print((m, len(str(m)))) |
| |
| def closed(self, code, reason=None): |
| ioloop.IOLoop.instance().stop() |
| |
| ws = MyClient('ws://localhost:9000/echo', protocols=['http-only', 'chat']) |
| ws.connect() |
| |
| ioloop.IOLoop.instance().start() |
| """ |
| WebSocketBaseClient.__init__(self, url, protocols, extensions, |
| ssl_options=ssl_options, headers=headers, exclude_headers=exclude_headers) |
| if self.scheme == "wss": |
| self.sock = ssl.wrap_socket(self.sock, do_handshake_on_connect=False, **self.ssl_options) |
| self._is_secure = True |
| self.io = iostream.SSLIOStream(self.sock, io_loop, ssl_options=self.ssl_options) |
| else: |
| self.io = iostream.IOStream(self.sock, io_loop) |
| self.io_loop = io_loop |
| |
| def connect(self): |
| """ |
| Connects the websocket and initiate the upgrade handshake. |
| """ |
| self.io.set_close_callback(self.__connection_refused) |
| self.io.connect((self.host, int(self.port)), self.__send_handshake) |
| |
| def _write(self, b): |
| """ |
| Trying to prevent a write operation |
| on an already closed websocket stream. |
| |
| This cannot be bullet proof but hopefully |
| will catch almost all use cases. |
| """ |
| if self.terminated: |
| raise RuntimeError("Cannot send on a terminated websocket") |
| |
| self.io.write(b) |
| |
| def __connection_refused(self, *args, **kwargs): |
| self.server_terminated = True |
| self.closed(1005, 'Connection refused') |
| |
| def __send_handshake(self): |
| self.io.set_close_callback(self.__connection_closed) |
| self.io.write(escape.utf8(self.handshake_request), |
| self.__handshake_sent) |
| |
| def __connection_closed(self, *args, **kwargs): |
| self.server_terminated = True |
| self.closed(1006, 'Connection closed during handshake') |
| |
| def __handshake_sent(self): |
| self.io.read_until(b"\r\n\r\n", self.__handshake_completed) |
| |
| def __handshake_completed(self, data): |
| self.io.set_close_callback(None) |
| try: |
| response_line, _, headers = data.partition(b'\r\n') |
| self.process_response_line(response_line) |
| protocols, extensions = self.process_handshake_header(headers) |
| except HandshakeError: |
| self.close_connection() |
| raise |
| |
| self.opened() |
| self.io.set_close_callback(self.__stream_closed) |
| self.io.read_bytes(self.reading_buffer_size, self.__fetch_more) |
| |
| def __fetch_more(self, bytes): |
| try: |
| should_continue = self.process(bytes) |
| except: |
| should_continue = False |
| |
| if should_continue: |
| self.io.read_bytes(self.reading_buffer_size, self.__fetch_more) |
| else: |
| self.__gracefully_terminate() |
| |
| def __gracefully_terminate(self): |
| self.client_terminated = self.server_terminated = True |
| |
| try: |
| if not self.stream.closing: |
| self.closed(1006) |
| finally: |
| self.close_connection() |
| |
| def __stream_closed(self, *args, **kwargs): |
| self.io.set_close_callback(None) |
| code = 1006 |
| reason = None |
| if self.stream.closing: |
| code, reason = self.stream.closing.code, self.stream.closing.reason |
| self.closed(code, reason) |
| self.stream._cleanup() |
| |
| def close_connection(self): |
| """ |
| Close the underlying connection |
| """ |
| self.io.close() |
| |
| if __name__ == '__main__': |
| from tornado import ioloop |
| |
| class MyClient(TornadoWebSocketClient): |
| def opened(self): |
| def data_provider(): |
| for i in range(0, 200, 25): |
| yield "#" * i |
| |
| self.send(data_provider()) |
| |
| for i in range(0, 200, 25): |
| self.send("*" * i) |
| |
| def received_message(self, m): |
| print("#%d" % len(m)) |
| if len(m) == 175: |
| self.close() |
| |
| def closed(self, code, reason=None): |
| ioloop.IOLoop.instance().stop() |
| print(("Closed down", code, reason)) |
| |
| ws = MyClient('ws://localhost:9000/ws', protocols=['http-only', 'chat']) |
| ws.connect() |
| |
| ioloop.IOLoop.instance().start() |
