diff --git a/bfxapi/websocket/BfxWebsocketClient.py b/bfxapi/websocket/BfxWebsocketClient.py index afb24a3..1d32e53 100644 --- a/bfxapi/websocket/BfxWebsocketClient.py +++ b/bfxapi/websocket/BfxWebsocketClient.py @@ -4,10 +4,29 @@ from pyee.asyncio import AsyncIOEventEmitter from .handlers import Channels, PublicChannelsHandler, AuthenticatedChannelsHandler -from .errors import BfxWebsocketException, ConnectionNotOpen, InvalidAuthenticationCredentials, EventNotSupported, OutdatedClientVersion +from .errors import ConnectionNotOpen, WebsocketAuthenticationRequired, InvalidAuthenticationCredentials, EventNotSupported, OutdatedClientVersion HEARTBEAT = "hb" +def _require_websocket_connection(function): + async def wrapper(self, *args, **kwargs): + if self.websocket == None or self.websocket.open == False: + raise ConnectionNotOpen("No open connection with the server.") + + await function(self, *args, **kwargs) + + return wrapper + +def _require_websocket_authentication(function): + @_require_websocket_connection + async def wrapper(self, *args, **kwargs): + if self.authentication == False: + raise WebsocketAuthenticationRequired("To perform this action you need to authenticate using your API_KEY and API_SECRET.") + + await function(self, *args, **kwargs) + + return wrapper + class BfxWebsocketClient(object): VERSION = 2 @@ -22,6 +41,8 @@ class BfxWebsocketClient(object): self.websocket, self.API_KEY, self.API_SECRET = None, API_KEY, API_SECRET + self.authentication = False + self.handlers = { "public": PublicChannelsHandler(event_emitter=self.event_emitter), "authenticated": AuthenticatedChannelsHandler(event_emitter=self.event_emitter) @@ -53,26 +74,19 @@ class BfxWebsocketClient(object): elif isinstance(message, dict) and message["event"] == "auth": if message["status"] == "OK": self.event_emitter.emit("authenticated", message) + + self.authentication = True else: raise InvalidAuthenticationCredentials("Cannot authenticate with given API-KEY and API-SECRET.") elif isinstance(message, dict) and message["event"] == "error": self.event_emitter.emit("error", message["code"], message["msg"]) - elif isinstance(message, list) and ((chanId := message[0]) or True) and message[1] != HEARTBEAT: - if chanId == 0: + elif isinstance(message, list) and message[1] != HEARTBEAT: + if ((chanId := message[0]) or True) and chanId == 0: self.handlers["authenticated"].handle(message[1], message[2]) else: self.handlers["public"].handle(self.chanIds[chanId], *message[1:]) except websockets.ConnectionClosed: continue - def __require_websocket_connection(function): - async def wrapper(self, *args, **kwargs): - if self.websocket == None or self.websocket.open == False: - raise ConnectionNotOpen("No open connection with the server.") - - await function(self, *args, **kwargs) - - return wrapper - - @__require_websocket_connection + @_require_websocket_connection async def subscribe(self, channel, **kwargs): await self.websocket.send(json.dumps({ "event": "subscribe", @@ -80,14 +94,14 @@ class BfxWebsocketClient(object): **kwargs })) - @__require_websocket_connection + @_require_websocket_connection async def unsubscribe(self, chanId): await self.websocket.send(json.dumps({ "event": "unsubscribe", "chanId": chanId })) - @__require_websocket_connection + @_require_websocket_connection async def authenticate(self, API_KEY, API_SECRET, filter=None): data = { "event": "auth", "filter": filter, "apiKey": API_KEY } @@ -103,6 +117,10 @@ class BfxWebsocketClient(object): await self.websocket.send(json.dumps(data)) + @_require_websocket_authentication + async def notify(self, MESSAGE_ID, info): + await self.websocket.send(json.dumps([ 0, "n", MESSAGE_ID, { "type": "ucm-test", "info": info } ])) + async def clear(self): for chanId in self.chanIds.keys(): await self.unsubscribe(chanId) diff --git a/bfxapi/websocket/__init__.py b/bfxapi/websocket/__init__.py index 2a4bfa0..f36a1b0 100644 --- a/bfxapi/websocket/__init__.py +++ b/bfxapi/websocket/__init__.py @@ -1,3 +1,3 @@ from .BfxWebsocketClient import BfxWebsocketClient from .handlers import Channels -from .errors import BfxWebsocketException, ConnectionNotOpen, InvalidAuthenticationCredentials, EventNotSupported, OutdatedClientVersion \ No newline at end of file +from .errors import BfxWebsocketException, ConnectionNotOpen, WebsocketAuthenticationRequired, InvalidAuthenticationCredentials, EventNotSupported, OutdatedClientVersion \ No newline at end of file diff --git a/bfxapi/websocket/errors.py b/bfxapi/websocket/errors.py index 6a45a94..7565bc9 100644 --- a/bfxapi/websocket/errors.py +++ b/bfxapi/websocket/errors.py @@ -18,6 +18,13 @@ class ConnectionNotOpen(BfxWebsocketException): pass +class WebsocketAuthenticationRequired(BfxWebsocketException): + """ + This error indicates an attempt to access a protected resource without logging in first. + """ + + pass + class InvalidAuthenticationCredentials(BfxWebsocketException): """ This error indicates that the user has provided incorrect credentials (API-KEY and API-SECRET) for authentication.