510 lines
17 KiB
Python
510 lines
17 KiB
Python
import websocket
|
|
import threading
|
|
import logging
|
|
import time
|
|
import json
|
|
import hmac
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
SPOT = "wss://wbs.mexc.com/ws"
|
|
FUTURES = "wss://contract.mexc.com/ws"
|
|
|
|
class _WebSocketManager:
|
|
def __init__(self, callback_function, ws_name, api_key=None, api_secret=None,
|
|
ping_interval=20, ping_timeout=10, retries=10,
|
|
restart_on_error=True, trace_logging=False,
|
|
http_proxy_host = None,
|
|
http_proxy_port = None,
|
|
http_no_proxy = None,
|
|
http_proxy_auth = None,
|
|
http_proxy_timeout = None):
|
|
|
|
self.proxy_settings = dict(
|
|
http_proxy_host = http_proxy_host,
|
|
http_proxy_port = http_proxy_port,
|
|
http_no_proxy = http_no_proxy,
|
|
http_proxy_auth = http_proxy_auth,
|
|
http_proxy_timeout = http_proxy_timeout
|
|
)
|
|
|
|
# Set API keys.
|
|
self.api_key = api_key
|
|
self.api_secret = api_secret
|
|
|
|
self.callback = callback_function
|
|
self.ws_name = ws_name
|
|
if api_key:
|
|
self.ws_name += " (Auth)"
|
|
|
|
# Setup the callback directory following the format:
|
|
# {
|
|
# "topic_name": function
|
|
# }
|
|
self.callback_directory = {}
|
|
|
|
# Record the subscriptions made so that we can resubscribe if the WSS
|
|
# connection is broken.
|
|
self.subscriptions = []
|
|
|
|
# Set ping settings.
|
|
self.ping_interval = ping_interval
|
|
self.ping_timeout = ping_timeout
|
|
self.retries = retries
|
|
|
|
# Other optional data handling settings.
|
|
self.handle_error = restart_on_error
|
|
|
|
# Enable websocket-client's trace logging for extra debug information
|
|
# on the websocket connection, including the raw sent & recv messages
|
|
websocket.enableTrace(trace_logging)
|
|
|
|
# Set initial state, initialize dictionary and connect.
|
|
self._reset()
|
|
self.attempting_connection = False
|
|
|
|
def _on_open(self):
|
|
"""
|
|
Log WS open.
|
|
"""
|
|
logger.debug(f"WebSocket {self.ws_name} opened.")
|
|
|
|
def _on_message(self, message):
|
|
"""
|
|
Parse incoming messages.
|
|
"""
|
|
self.callback(json.loads(message))
|
|
|
|
def is_connected(self):
|
|
try:
|
|
if self.ws.sock or not self.ws.sock.is_connected:
|
|
return True
|
|
else:
|
|
return False
|
|
except AttributeError:
|
|
return False
|
|
|
|
@staticmethod
|
|
def _are_connections_connected(active_connections):
|
|
for connection in active_connections:
|
|
if not connection.is_connected():
|
|
return False
|
|
return True
|
|
|
|
def _ping_loop(self, ping_payload: str, ping_interval: int, ping_timeout: int):
|
|
"""
|
|
Ping the websocket.
|
|
"""
|
|
time.sleep(ping_timeout)
|
|
while True:
|
|
logger.info(f"WebSocket {self.ws_name} send ping...")
|
|
self.ws.send(ping_payload)
|
|
time.sleep(ping_interval)
|
|
|
|
def _connect(self, url):
|
|
"""
|
|
Open websocket in a thread.
|
|
"""
|
|
def resubscribe_to_topics():
|
|
if not self.subscriptions:
|
|
# There are no subscriptions to resubscribe to, probably
|
|
# because this is a brand new WSS initialisation so there was
|
|
# no previous WSS connection.
|
|
return
|
|
for subscription_message in self.subscriptions:
|
|
self.ws.send(subscription_message)
|
|
|
|
self.attempting_connection = True
|
|
|
|
# Set endpoint.
|
|
self.endpoint = url
|
|
|
|
|
|
# Attempt to connect for X seconds.
|
|
retries = self.retries
|
|
if retries == 0:
|
|
infinitely_reconnect = True
|
|
else:
|
|
infinitely_reconnect = False
|
|
while (infinitely_reconnect or retries > 0) and not self.is_connected():
|
|
logger.info(f"WebSocket {self.ws_name} attempting connection...")
|
|
self.ws = websocket.WebSocketApp(
|
|
url=url,
|
|
on_message=lambda ws, msg: self._on_message(msg),
|
|
on_close=self._on_close(),
|
|
on_open=self._on_open(),
|
|
on_error=lambda ws, err: self._on_error(err)
|
|
)
|
|
|
|
# Setup the thread running WebSocketApp.
|
|
self.wst = threading.Thread(target=lambda: self.ws.run_forever(
|
|
ping_interval=self.ping_interval,
|
|
ping_timeout=self.ping_timeout,
|
|
**self.proxy_settings
|
|
))
|
|
|
|
# Configure as daemon; start.
|
|
self.wst.daemon = True
|
|
self.wst.start()
|
|
|
|
# setup ping loop
|
|
self.wsl = threading.Thread(target=lambda: self._ping_loop(
|
|
ping_payload='{"method":"ping"}',
|
|
ping_interval=self.ping_interval,
|
|
ping_timeout=self.ping_timeout
|
|
))
|
|
|
|
self.wsl.daemon = True
|
|
self.wsl.start()
|
|
|
|
retries -= 1
|
|
time.sleep(1)
|
|
|
|
# If connection was not successful, raise error.
|
|
if not infinitely_reconnect and retries <= 0:
|
|
self.exit()
|
|
raise websocket.WebSocketTimeoutException(
|
|
f"WebSocket {self.ws_name} connection failed. Too many "
|
|
f"connection attempts. pybit will "
|
|
f"no longer try to reconnect.")
|
|
|
|
logger.info(f"WebSocket {self.ws_name} connected")
|
|
|
|
# If given an api_key, authenticate.
|
|
if self.api_key and self.api_secret:
|
|
self._auth()
|
|
|
|
resubscribe_to_topics()
|
|
|
|
self.attempting_connection = False
|
|
|
|
def _auth(self):
|
|
# Generate signature
|
|
|
|
# make auth if futures. spot has a different auth system.
|
|
|
|
isspot = self.endpoint.startswith(SPOT)
|
|
if isspot:
|
|
return
|
|
|
|
timestamp = str(int(time.time() * 1000))
|
|
_val = self.api_key + timestamp
|
|
signature = str(hmac.new(
|
|
bytes(self.api_secret, "utf-8"),
|
|
bytes(_val, "utf-8"), digestmod="sha256"
|
|
).hexdigest())
|
|
|
|
# Authenticate with API.
|
|
self.ws.send(
|
|
json.dumps({
|
|
"subscribe": False,
|
|
"method": "login",
|
|
"param": {
|
|
"apiKey": self.api_key,
|
|
"reqTime": timestamp,
|
|
"signature": signature
|
|
}
|
|
})
|
|
)
|
|
|
|
def _on_error(self, error):
|
|
"""
|
|
Exit on errors and raise exception, or attempt reconnect.
|
|
"""
|
|
if type(error).__name__ not in ["WebSocketConnectionClosedException",
|
|
"ConnectionResetError",
|
|
"WebSocketTimeoutException"]:
|
|
# Raises errors not related to websocket disconnection.
|
|
self.exit()
|
|
raise error
|
|
|
|
if not self.exited:
|
|
logger.error(f"WebSocket {self.ws_name} encountered error: {error}.")
|
|
self.exit()
|
|
|
|
# Reconnect.
|
|
if self.handle_error and not self.attempting_connection:
|
|
self._reset()
|
|
self._connect(self.endpoint)
|
|
|
|
def _on_close(self):
|
|
"""
|
|
Log WS close.
|
|
"""
|
|
logger.debug(f"WebSocket {self.ws_name} closed.")
|
|
|
|
def _reset(self):
|
|
"""
|
|
Set state booleans and initialize dictionary.
|
|
"""
|
|
self.exited = False
|
|
self.auth = False
|
|
self.data = {}
|
|
|
|
def exit(self):
|
|
"""
|
|
Closes the websocket connection.
|
|
"""
|
|
|
|
self.ws.close()
|
|
while self.ws.sock:
|
|
continue
|
|
self.exited = True
|
|
|
|
class _FuturesWebSocketManager(_WebSocketManager):
|
|
def __init__(self, ws_name, **kwargs):
|
|
callback_function = kwargs.pop("callback_function") if \
|
|
kwargs.get("callback_function") else self._handle_incoming_message
|
|
|
|
super().__init__(callback_function, ws_name, **kwargs)
|
|
|
|
self.private_topics = ["personal.order", "personal.asset",
|
|
"personal.position", "personal.risk.limit",
|
|
"personal.adl.level", "personal.position.mode"]
|
|
|
|
self.symbol_wildcard = "*"
|
|
self.symbol_separator = "|"
|
|
self.last_subsctiption = None
|
|
|
|
def subscribe(self, topic, callback, params: dict = {}):
|
|
subscription_args = {
|
|
"method": topic,
|
|
"param": params
|
|
}
|
|
self._check_callback_directory(subscription_args)
|
|
|
|
while not self.is_connected():
|
|
# Wait until the connection is open before subscribing.
|
|
time.sleep(0.1)
|
|
|
|
subscription_message = json.dumps(subscription_args)
|
|
self.ws.send(subscription_message)
|
|
self.subscriptions.append(subscription_message)
|
|
self._set_callback(topic.replace("sub.", ""), callback)
|
|
self.last_subsctiption = topic.replace("sub.", "")
|
|
|
|
def _initialise_local_data(self, topic):
|
|
# Create self.data
|
|
try:
|
|
self.data[topic]
|
|
except KeyError:
|
|
self.data[topic] = []
|
|
|
|
def _process_auth_message(self, message):
|
|
# If we get successful futures auth, notify user
|
|
if message.get("data") == "success":
|
|
logger.debug(f"Authorization for {self.ws_name} successful.")
|
|
self.auth = True
|
|
# If we get unsuccessful auth, notify user.
|
|
elif message.get("data") != "success": # !!!!
|
|
logger.debug(f"Authorization for {self.ws_name} failed. Please "
|
|
f"check your API keys and restart.")
|
|
|
|
def _process_subscription_message(self, message):
|
|
#try:
|
|
sub = message["channel"]
|
|
#except KeyError:
|
|
#sub = message["c"] # SPOT PUBLIC & PRIVATE format
|
|
|
|
# If we get successful futures subscription, notify user
|
|
if (
|
|
message.get("channel", "").startswith("rs.") or
|
|
message.get("channel", "").startswith("push.")
|
|
) and message.get("channel", "") != "rs.error":
|
|
|
|
logger.debug(f"Subscription to {sub} successful.")
|
|
# Futures subscription fail
|
|
else:
|
|
response = message["data"]
|
|
logger.error("Couldn't subscribe to topic. "
|
|
f"Error: {response}.")
|
|
if self.last_subsctiption:
|
|
self._pop_callback(self.last_subsctiption)
|
|
|
|
def _process_normal_message(self, message):
|
|
topic = message["channel"].replace("push.", "").replace("rs.sub.", "")
|
|
callback_data = message
|
|
callback_function = self._get_callback(topic)
|
|
callback_function(callback_data)
|
|
|
|
def _handle_incoming_message(self, message):
|
|
def is_auth_message():
|
|
if message.get("channel", "") == "rs.login":
|
|
return True
|
|
else:
|
|
return False
|
|
|
|
def is_subscription_message():
|
|
if str(message).startswith("{'channel': 'push."):
|
|
return True
|
|
else:
|
|
return False
|
|
|
|
def is_pong_message():
|
|
if message.get("channel", "") in ("pong", "clientId"):
|
|
return True
|
|
else:
|
|
return False
|
|
|
|
if is_auth_message():
|
|
self._process_auth_message(message)
|
|
elif is_subscription_message():
|
|
self._process_subscription_message(message)
|
|
elif is_pong_message():
|
|
pass
|
|
else:
|
|
self._process_normal_message(message)
|
|
|
|
def custom_topic_stream(self, topic, callback):
|
|
return self.subscribe(topic=topic, callback=callback)
|
|
|
|
def _check_callback_directory(self, topics):
|
|
for topic in topics:
|
|
if topic in self.callback_directory:
|
|
raise Exception(f"You have already subscribed to this topic: "
|
|
f"{topic}")
|
|
|
|
def _set_callback(self, topic, callback_function):
|
|
self.callback_directory[topic] = callback_function
|
|
|
|
def _get_callback(self, topic):
|
|
return self.callback_directory[topic]
|
|
|
|
def _pop_callback(self, topic):
|
|
self.callback_directory.pop(topic)
|
|
|
|
class _FuturesWebSocket(_FuturesWebSocketManager):
|
|
def __init__(self, **kwargs):
|
|
self.ws_name = "FuturesV1"
|
|
self.endpoint = "wss://contract.mexc.com/ws"
|
|
|
|
super().__init__(self.ws_name, **kwargs)
|
|
self.ws = None
|
|
|
|
self.active_connections = []
|
|
self.kwargs = kwargs
|
|
|
|
def is_connected(self):
|
|
return self._are_connections_connected(self.active_connections)
|
|
|
|
def _ws_subscribe(self, topic, callback, params: list = []):
|
|
if not self.ws:
|
|
self.ws = _FuturesWebSocketManager(
|
|
self.ws_name, **self.kwargs)
|
|
self.ws._connect(self.endpoint)
|
|
self.active_connections.append(self.ws)
|
|
self.ws.subscribe(topic, callback, params)
|
|
|
|
class _SpotWebSocketManager(_WebSocketManager):
|
|
def __init__(self, ws_name, **kwargs):
|
|
callback_function = kwargs.pop("callback_function") if \
|
|
kwargs.get("callback_function") else self._handle_incoming_message
|
|
super().__init__(callback_function, ws_name, **kwargs)
|
|
|
|
self.private_topics = ["account", "deals", "orders"]
|
|
|
|
self.last_subsctiption = None
|
|
|
|
def subscribe(self, topic: str, callback, params_list: list):
|
|
subscription_args = {
|
|
"method": "SUBSCRIPTION",
|
|
"params": [
|
|
'@'.join([f"spot@{topic}.v3.api"] + list(map(str, params.values())))
|
|
for params in params_list
|
|
]
|
|
}
|
|
self._check_callback_directory(subscription_args)
|
|
|
|
while not self.is_connected():
|
|
# Wait until the connection is open before subscribing.
|
|
time.sleep(0.1)
|
|
|
|
subscription_message = json.dumps(subscription_args)
|
|
self.ws.send(subscription_message)
|
|
self.subscriptions.append(subscription_message)
|
|
self._set_callback(topic, callback)
|
|
self.last_subsctiption = topic
|
|
|
|
def _initialise_local_data(self, topic):
|
|
# Create self.data
|
|
try:
|
|
self.data[topic]
|
|
except KeyError:
|
|
self.data[topic] = []
|
|
|
|
|
|
def _process_subscription_message(self, message):
|
|
sub = message["msg"].replace("spot@", "").split(".v3.api")[0]
|
|
|
|
# If we get successful futures subscription, notify user
|
|
if message.get("id") == 0 and message.get("code") == 0:
|
|
logger.debug(f"Subscription to {sub} successful.")
|
|
# Futures subscription fail
|
|
else:
|
|
response = message["msg"]
|
|
logger.error("Couldn't subscribe to topic. "
|
|
f"Error: {response}.")
|
|
if self.last_subsctiption:
|
|
self._pop_callback(self.last_subsctiption)
|
|
|
|
def _process_normal_message(self, message):
|
|
topic = message["c"].replace("spot@", "").split(".v3.api")[0]
|
|
callback_data = message
|
|
callback_function = self._get_callback(topic)
|
|
callback_function(callback_data)
|
|
|
|
def _handle_incoming_message(self, message):
|
|
def is_subscription_message():
|
|
if (message.get("id") == 0 and
|
|
message.get("code") == 0 and
|
|
message.get("msg")):
|
|
return True
|
|
else:
|
|
return False
|
|
|
|
if is_subscription_message():
|
|
self._process_subscription_message(message)
|
|
else:
|
|
self._process_normal_message(message)
|
|
|
|
def custom_topic_stream(self, topic, callback):
|
|
return self.subscribe(topic=topic, callback=callback)
|
|
|
|
def _check_callback_directory(self, topics):
|
|
for topic in topics:
|
|
if topic in self.callback_directory:
|
|
raise Exception(f"You have already subscribed to this topic: "
|
|
f"{topic}")
|
|
|
|
def _set_callback(self, topic, callback_function):
|
|
self.callback_directory[topic] = callback_function
|
|
|
|
def _get_callback(self, topic):
|
|
return self.callback_directory[topic]
|
|
|
|
def _pop_callback(self, topic):
|
|
self.callback_directory.pop(topic)
|
|
|
|
class _SpotWebSocket(_SpotWebSocketManager):
|
|
def __init__(self, endpoint: str = "wss://wbs.mexc.com/ws", **kwargs):
|
|
self.ws_name = "SpotV3"
|
|
self.endpoint = endpoint
|
|
|
|
|
|
super().__init__(self.ws_name, **kwargs)
|
|
self.ws = None
|
|
|
|
self.active_connections = []
|
|
self.kwargs = kwargs
|
|
|
|
def is_connected(self):
|
|
return self._are_connections_connected(self.active_connections)
|
|
|
|
def _ws_subscribe(self, topic, callback, params: list = []):
|
|
if not self.ws:
|
|
self.ws = _SpotWebSocketManager(
|
|
self.ws_name, **self.kwargs)
|
|
self.ws._connect(self.endpoint)
|
|
self.active_connections.append(self.ws)
|
|
self.ws.subscribe(topic, callback, params)
|