def get_case_count(server): uri = urlparse(server + '/getCaseCount') connection = WSConnection(CLIENT, uri.netloc, uri.path) reader, writer = yield from asyncio.open_connection( uri.hostname, uri.port or 80) writer.write(connection.bytes_to_send()) case_count = None while case_count is None: data = yield from reader.read(65535) connection.receive_bytes(data) data = "" for event in connection.events(): if isinstance(event, TextReceived): data += event.data if event.message_finished: case_count = json.loads(data) connection.close() try: writer.write(connection.bytes_to_send()) yield from writer.drain() except (ConnectionError, OSError): break return case_count
def get_case_count(server): uri = urlparse(server + '/getCaseCount') connection = WSConnection(CLIENT, uri.netloc, uri.path) sock = socket.socket() sock.connect((uri.hostname, uri.port or 80)) sock.sendall(connection.bytes_to_send()) case_count = None while case_count is None: data = sock.recv(65535) connection.receive_bytes(data) data = "" for event in connection.events(): if isinstance(event, TextReceived): data += event.data if event.message_finished: case_count = json.loads(data) connection.close() try: sock.sendall(connection.bytes_to_send()) except CONNECTION_EXCEPTIONS: break sock.close() return case_count
async def ws_adapter(in_q, out_q, client, _): """A simple, queue-based Curio-Sans-IO websocket bridge.""" client.setsockopt(IPPROTO_TCP, TCP_NODELAY, 1) wsconn = WSConnection(SERVER) closed = False while not closed: wstask = await spawn(client.recv, 65535) outqtask = await spawn(out_q.get) async with TaskGroup([wstask, outqtask]) as g: task = await g.next_done() result = await task.join() await g.cancel_remaining() if task is wstask: wsconn.receive_bytes(result) for event in wsconn.events(): cl = event.__class__ if cl in DATA_TYPES: await in_q.put(event.data) elif cl is ConnectionRequested: # Auto accept. Maybe consult the handler? wsconn.accept(event) elif cl is ConnectionClosed: # The client has closed the connection. await in_q.put(None) closed = True else: print(event) await client.sendall(wsconn.bytes_to_send()) else: # We got something from the out queue. if result is None: # Terminate the connection. print("Closing the connection.") wsconn.close() closed = True else: wsconn.send_data(result) payload = wsconn.bytes_to_send() await client.sendall(payload) print("Bridge done.")
def update_reports(server, agent): uri = urlparse(server + '/updateReports?agent=%s' % agent) connection = WSConnection(CLIENT, uri.netloc, '%s?%s' % (uri.path, uri.query)) sock = socket.socket() sock.connect((uri.hostname, uri.port or 80)) sock.sendall(connection.bytes_to_send()) closed = False while not closed: data = sock.recv(65535) connection.receive_bytes(data) for event in connection.events(): if isinstance(event, ConnectionEstablished): connection.close() sock.sendall(connection.bytes_to_send()) try: sock.close() except CONNECTION_EXCEPTIONS: pass finally: closed = True
def update_reports(server, agent): uri = urlparse(server + '/updateReports?agent=%s' % agent) connection = WSConnection(CLIENT, uri.netloc, '%s?%s' % (uri.path, uri.query)) reader, writer = yield from asyncio.open_connection( uri.hostname, uri.port or 80) writer.write(connection.bytes_to_send()) closed = False while not closed: data = yield from reader.read(65535) connection.receive_bytes(data) for event in connection.events(): if isinstance(event, ConnectionEstablished): connection.close() writer.write(connection.bytes_to_send()) try: yield from writer.drain() writer.close() except (ConnectionError, OSError): pass finally: closed = True
class WebSocketEndpoint(AbstractEndpoint): """ Implements websocket endpoints. Subprotocol negotiation is currently not supported. """ __slots__ = ('ctx', '_client', '_ws') def __init__(self, ctx: Context, client: BaseHTTPClientConnection): self.ctx = ctx self._client = client self._ws = WSConnection(ConnectionType.SERVER) def _process_ws_events(self): for event in self._ws.events(): if isinstance(event, ConnectionRequested): self._ws.accept(event) self.on_connect() elif isinstance(event, DataReceived): self.on_data(event.data) elif isinstance(event, ConnectionClosed): self.on_close() bytes_to_send = self._ws.bytes_to_send() if bytes_to_send: self._client.write(bytes_to_send) def begin_request(self, request: HTTPRequest): trailing_data = self._client.upgrade() self._ws.receive_bytes(trailing_data) self._process_ws_events() def receive_body_data(self, data: bytes) -> None: self._ws.receive_bytes(data) self._process_ws_events() def send_message(self, payload: Union[str, bytes]) -> None: """ Send a message to the client. :param payload: either a unicode string or a bytestring """ self._ws.send_data(payload) bytes_to_send = self._ws.bytes_to_send() self._client.write(bytes_to_send) def close(self) -> None: """Close the websocket.""" self._ws.close() self._process_ws_events() def on_connect(self) -> None: """Called when the websocket handshake has been done.""" def on_close(self) -> None: """Called when the connection has been closed.""" def on_data(self, data: bytes) -> None: """Called when there is new data from the peer."""
def wsproto_demo(host, port): ''' Demonstrate wsproto: 0) Open TCP connection 1) Negotiate WebSocket opening handshake 2) Send a message and display response 3) Send ping and display pong 4) Negotiate WebSocket closing handshake :param stream: a socket stream ''' # 0) Open TCP connection print('Connecting to {}:{}'.format(host, port)) conn = socket.socket(socket.AF_INET, socket.SOCK_STREAM) conn.connect((host, port)) # 1) Negotiate WebSocket opening handshake print('Opening WebSocket') ws = WSConnection(ConnectionType.CLIENT, host=host, resource='server') # events is a generator that yields websocket event objects. Usually you # would say `for event in ws.events()`, but the synchronous nature of this # client requires us to use next(event) instead so that we can interleave # the network I/O. It will raise StopIteration when it runs out of events # (i.e. needs more network data), but since this script is synchronous, we # will explicitly resume the generator whenever we have new network data. events = ws.events() # Because this is a client WebSocket, wsproto has automatically queued up # a handshake, and we need to send it and wait for a response. net_send_recv(ws, conn) event = next(events) if isinstance(event, ConnectionEstablished): print('WebSocket negotiation complete') else: raise Exception('Expected ConnectionEstablished event!') # 2) Send a message and display response message = "wsproto is great" print('Sending message: {}'.format(message)) ws.send_data(message) net_send_recv(ws, conn) event = next(events) if isinstance(event, TextReceived): print('Received message: {}'.format(event.data)) else: raise Exception('Expected TextReceived event!') # 3) Send ping and display pong payload = b"table tennis" print('Sending ping: {}'.format(payload)) ws.ping(payload) net_send_recv(ws, conn) event = next(events) if isinstance(event, PongReceived): print('Received pong: {}'.format(event.payload)) else: raise Exception('Expected PongReceived event!') # 4) Negotiate WebSocket closing handshake print('Closing WebSocket') ws.close(code=1000, reason='sample reason') # After sending the closing frame, we won't get any more events. The server # should send a reply and then close the connection, so we need to receive # twice: net_send_recv(ws, conn) conn.shutdown(socket.SHUT_WR) net_recv(ws, conn)
class ClientWebsocket(object): """ Represents a ClientWebsocket. """ def __init__(self, address: Tuple[str, str, bool, str], *, reconnecting: bool = True): """ :param type_: The :class:`wsproto.connection.ConnectionType` for this websocket. :param address: A 4-tuple of (host, port, ssl, endpoint). :param endpoint: The endpoint to open the connection to. :param reconnecting: If this websocket reconnects automatically. Only useful on the client. """ # these are all used to construct the state object self._address = address self.state = None # type: WSConnection self._ready = False self._reconnecting = reconnecting self.sock = None # type: multio.SocketWrapper @property def closed(self) -> bool: """ :return: If this websocket is closed. """ return self.state.closed async def open_connection(self): """ Opens a connection, and performs the initial handshake. """ _sock = await multio.asynclib.open_connection(self._address[0], self._address[1], ssl=self._address[2]) self.sock = multio.SocketWrapper(_sock) self.state = WSConnection(ConnectionType.CLIENT, host=self._address[0], resource=self._address[3]) res = self.state.bytes_to_send() await self.sock.sendall(res) async def __aiter__(self): # initiate the websocket await self.open_connection() # setup buffers buf_bytes = BytesIO() buf_text = StringIO() while True: data = await self.sock.recv(4096) self.state.receive_bytes(data) # do ping/pongs if needed to_send = self.state.bytes_to_send() if to_send: await self.sock.sendall(to_send) for event in self.state.events(): if isinstance(event, events.ConnectionEstablished): self._ready = True yield WebsocketConnectionEstablished(event) elif isinstance(event, events.ConnectionClosed): self._ready = False yield WebsocketClosed(event) if self._reconnecting: await self.open_connection() else: return elif isinstance(event, events.DataReceived): buf = buf_bytes if isinstance( event, events.BytesReceived) else buf_text buf.write(event.data) # yield events as appropriate if event.message_finished: buf.seek(0) read = buf.read() # empty buffer buf.truncate(0) buf.seek(0) typ = WebsocketBytesMessage if isinstance(event, events.BytesReceived) \ else WebsocketTextMessage yield typ(read) elif isinstance(event, events.ConnectionFailed): self._ready = False yield WebsocketConnectionFailed(event) if self._reconnecting: await self.open_connection() else: return async def send_message(self, data: Union[str, bytes]): """ Sends a message on the websocket. :param data: The data to send. Either str or bytes. """ self.state.send_data(data, final=True) return await self.sock.sendall(self.state.bytes_to_send()) async def close(self, *, code: int = 1000, reason: str = "No reason", allow_reconnects: bool = False): """ Closes the websocket. :param code: The close code to use. :param reason: The close reason to use. If the websocket is marked as reconnecting: :param allow_reconnects: If the websocket can reconnect after this close. """ # do NOT reconnect if we close explicitly and don't allow reconnects if not allow_reconnects: self._reconnecting = False self.state.close(code=code, reason=reason) to_send = self.state.bytes_to_send() await self.sock.sendall(to_send) await self.sock.close() self.state.receive_bytes(None)
async def wsproto_client_demo(host, port, use_ssl): ''' Demonstrate wsproto: 0) Open TCP connection 1) Negotiate WebSocket opening handshake 2) Send a message and display response 3) Send ping and display pong 4) Negotiate WebSocket closing handshake :param stream: a socket stream ''' # 0) Open TCP connection print('[C] Connecting to {}:{}'.format(host, port)) conn = await trio.open_tcp_stream(host, port) if use_ssl: conn = upgrade_stream_to_ssl(conn, host) # 1) Negotiate WebSocket opening handshake print('[C] Opening WebSocket') ws = WSConnection(ConnectionType.CLIENT, host=host, resource='server') events = ws.events() # Because this is a client WebSocket, wsproto has automatically queued up # a handshake, and we need to send it and wait for a response. await net_send_recv(ws, conn) event = next(events) if isinstance(event, ConnectionEstablished): print('[C] WebSocket negotiation complete') else: raise Exception(f'Expected ConnectionEstablished event! Got: {event}') # 2) Send a message and display response message = "wsproto is great" * 10 print('[C] Sending message: {}'.format(message)) ws.send_data(message) await net_send_recv(ws, conn) event = next(events) if isinstance(event, TextReceived): print('[C] Received message: {}'.format(event.data)) else: raise Exception(f'Expected TextReceived event! Got: {event}') # 3) Send ping and display pong payload = b"table tennis" print('[C] Sending ping: {}'.format(payload)) ws.ping(payload) await net_send_recv(ws, conn) event = next(events) if isinstance(event, PongReceived): print('[C] Received pong: {}'.format(event.payload)) else: raise Exception(f'Expected PongReceived event! Got: {event}') # 4) Negotiate WebSocket closing handshake print('[C] Closing WebSocket') ws.close(code=1000, reason='sample reason') # After sending the closing frame, we won't get any more events. The server # should send a reply and then close the connection, so we need to receive # twice: await net_send_recv(ws, conn) await conn.aclose()
class ClientWebsocket(object): """ Represents a ClientWebsocket. """ def __init__(self, address: Tuple[str, str, bool, str], *, reconnecting: bool = True): """ :param type_: The :class:`wsproto.connection.ConnectionType` for this websocket. :param address: A 4-tuple of (host, port, ssl, endpoint). :param endpoint: The endpoint to open the connection to. :param reconnecting: If this websocket reconnects automatically. Only useful on the client. """ # these are all used to construct the state object self._address = address self.state = None # type: WSConnection self._ready = False self._reconnecting = reconnecting self.sock = None # type: trio.socket.Socket @property def closed(self) -> bool: """ :return: If this websocket is closed. """ return self.state.closed def _create_ssl_ctx(self, sslp): if isinstance(sslp, ssl.SSLContext): return sslp ca = sslp.get('ca') capath = sslp.get('capath') hasnoca = ca is None and capath is None ctx = ssl.create_default_context(cafile=ca, capath=capath) ctx.check_hostname = not hasnoca and sslp.get('check_hostname', True) ctx.verify_mode = ssl.CERT_NONE if hasnoca else ssl.CERT_REQUIRED if 'cert' in sslp: ctx.load_cert_chain(sslp['cert'], keyfile=sslp.get('key')) if 'cipher' in sslp: ctx.set_ciphers(sslp['cipher']) ctx.options |= ssl.OP_NO_SSLv2 ctx.options |= ssl.OP_NO_SSLv3 return ctx async def open_connection(self): """ Opens a connection, and performs the initial handshake. """ _ssl = self._address[2] _sock = await trio.open_tcp_stream(self._address[0], self._address[1]) if _ssl: if _ssl is True: _ssl = {} _ssl = self._create_ssl_ctx(_ssl) _sock = trio.ssl.SSLStream(_sock, _ssl, server_hostname=self._address[0], https_compatible=True) await _sock.do_handshake() self.sock = _sock self.state = WSConnection(ConnectionType.CLIENT, host=self._address[0], resource=self._address[3]) res = self.state.bytes_to_send() await self.sock.send_all(res) async def __aiter__(self): # initiate the websocket await self.open_connection() # setup buffers buf_bytes = BytesIO() buf_text = StringIO() while self.sock is not None: data = await self.sock.receive_some(4096) self.state.receive_bytes(data) # do ping/pongs if needed to_send = self.state.bytes_to_send() if to_send: await self.sock.send_all(to_send) for event in self.state.events(): if isinstance(event, events.ConnectionEstablished): self._ready = True yield WebsocketConnectionEstablished(event) elif isinstance(event, events.ConnectionClosed): self._ready = False yield WebsocketClosed(event) if self._reconnecting: await self.open_connection() else: return elif isinstance(event, events.DataReceived): buf = buf_bytes if isinstance(event, events.BytesReceived) else buf_text buf.write(event.data) # yield events as appropriate if event.message_finished: buf.seek(0) read = buf.read() # empty buffer buf.truncate(0) buf.seek(0) typ = WebsocketBytesMessage if isinstance(event, events.BytesReceived) \ else WebsocketTextMessage yield typ(read) elif isinstance(event, events.ConnectionFailed): self._ready = False yield WebsocketConnectionFailed(event) if self._reconnecting: await self.open_connection() else: return else: raise RuntimeError("I don't understand this message", event) async def send_message(self, data: Union[str, bytes]): """ Sends a message on the websocket. :param data: The data to send. Either str or bytes. """ self.state.send_data(data, final=True) return await self.sock.send_all(self.state.bytes_to_send()) async def aclose(self, *, code: int = 1000, reason: str = "No reason", allow_reconnects: bool = False): """ Closes the websocket. :param code: The close code to use. :param reason: The close reason to use. If the websocket is marked as reconnecting: :param allow_reconnects: If the websocket can reconnect after this close. """ # do NOT reconnect if we close explicitly and don't allow reconnects if not allow_reconnects: self._reconnecting = False self.state.close(code=code, reason=reason) to_send = self.state.bytes_to_send() await self.sock.send_all(to_send) await self.sock.aclose() self.sock = None self.state.receive_bytes(None)