85 lines
2.5 KiB
Python
85 lines
2.5 KiB
Python
import websockets
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
|
|
logging.basicConfig(level=logging.INFO)
|
|
|
|
class JsonWebsocketConnection(object):
|
|
|
|
def __init__(self,
|
|
websocket: websockets.WebSocketServerProtocol):
|
|
|
|
logging.info("incoming connection")
|
|
self._websocket = websocket
|
|
self._is_closed = False;
|
|
|
|
def is_closed(self):
|
|
return self._is_closed
|
|
|
|
async def close(self):
|
|
logging.info("closing connection")
|
|
try:
|
|
if self._is_closed:
|
|
return
|
|
self._websocket.close()
|
|
self._is_closed = True
|
|
except Exception as e:
|
|
logging.warning("error closing connection: %s", str(e))
|
|
|
|
async def send(self, message: dict):
|
|
string_message = json.dumps(message)
|
|
logging.debug("sending message: %s", string_message)
|
|
try:
|
|
await self._websocket.send(string_message)
|
|
except Exception as e:
|
|
logging.warning("error sending message: %s", str(e))
|
|
|
|
async def handle_message(self, message: dict):
|
|
pass # override this function
|
|
|
|
async def run(self):
|
|
try:
|
|
async for message in self._websocket:
|
|
try:
|
|
json_message = json.loads(message)
|
|
await self.handle_message(json_message)
|
|
except ValueError as e:
|
|
logging.warning("received unprocessable message %s", str(e))
|
|
|
|
|
|
except Exception as e:
|
|
logging.warning("error in websocket connection: %s", str(e))
|
|
self._is_closed = True
|
|
finally:
|
|
self._is_closed = True
|
|
|
|
|
|
|
|
|
|
|
|
class JsonWebsocketServer(object):
|
|
def __init__(self, handler_class: JsonWebsocketConnection, host:str = 'localhost', port:int = 8765, ssl_context = None):
|
|
self._host = host
|
|
self._port = port
|
|
self._handler_class = handler_class
|
|
self._ssl_context = ssl_context
|
|
|
|
def run(self):
|
|
async def main(websocket: websockets.WebSocketServerProtocol,
|
|
path: str):
|
|
|
|
connection = self._handler_class(websocket)
|
|
|
|
await connection.run()
|
|
|
|
if (self._ssl_context is not None):
|
|
start_server = websockets.serve(main, self._host, self._port, ssl=self._ssl_context)
|
|
else:
|
|
start_server = websockets.serve(main, self._host, self._port)
|
|
|
|
asyncio.get_event_loop().run_until_complete(start_server)
|
|
asyncio.get_event_loop().run_forever()
|
|
|
|
|