OKX python webSocket

import hmac
import os
import socket
import time

import socks
import websocket
import base64
import json
import okx.consts as c
import queue as Queue

try:
	import thread
except ImportError:
	import _thread as thread

class TestOkxWebSocket ( object ):
	def __init__( self, apiKey, secretKey, passphrase, sandbox=False ):
		self.sandbox = sandbox
		self.apiKey = apiKey
		self.secretKey = secretKey
		self.passphrase = passphrase
		self.msgQueue = Queue.Queue ()
		self.ws = None
		self.channels = []

	def private_subscribe( self ):
		# 余额与订单频道
		'''
		实盘交易
		WS_PUBLIC_URL='wss://ws.okx.com:8443/ws/v5/public'
		WS_PRIVATE_URL='wss://ws.okx.com:8443/ws/v5/private'
	
		模拟盘交易
		WS_SANBOX_PUBLIC_URL='wss://wspap.okx.com:8443/ws/v5/public?brokerId=9999'
		WS_SANBOX_PRIVATE_URL='wss://wspap.okx.com:8443/ws/v5/private?brokerId=9999'	
		'''
		
		self.url = c.WS_SANBOX_PRIVATE_URL if self.sandbox else c.WS_PRIVATE_URL
		self.login = True
		self.channels.append ( {"channel": "orders", "instType": "SWAP"} )
		# self.channels.append ( {"channel": "balance_and_position", 'ccy': 'USDT'} )
		# self.channels.append ( {"channel": "account", 'ccy': 'USDT'} )
		# self.channels.append ( {"channel": "orders-algo", "instType": "SWAP"} )
		self.listenOnThread ()

	def private_unsubscribe( self ):
		sub_param = {'op': 'unsubscribe', 'args': self.channels}
		msg = json.dumps ( sub_param )
		self.ws.send ( msg )
		self.close ()

	def public_subscribe( self, bassAsset ):
		if self.ws == None:
			self._add_pubilc_subscribe ( bassAsset )
		else:
			self.msgQueue.put ( {'event': 'add', 'asset': bassAsset} )

	def _add_pubilc_subscribe( self, bassAsset ):

		self.url = c.WS_SANBOX_PUBLIC_URL if self.sandbox else c.WS_PUBLIC_URL
		self.login = False

		instId = str ( bassAsset.upper () + "-USDT-SWAP" ).strip ()
		self.channels.append (
			{"channel": "mark-price", "instType": "SWAP", "instId": instId} )
		self.channels.append (
			{"channel": "tickers", "instType": "SWAP", "instId": instId} )

		if self.ws is None:
			self.listenOnThread ()
		else:  # 实时订阅
			addchannels = []
			addchannels.append (
				{"channel": "mark-price", "instType": "SWAP", "instId": instId} )
			addchannels.append (
				{"channel": "tickers", "instType": "SWAP", "instId": instId} )
			sub_param = {'op': 'subscribe', 'args': addchannels}
			msg = json.dumps ( sub_param )
			self.ws.send ( msg )
			time.sleep ( 1 )
			print ( 'TestOkxWebSocket 添加交易对', self.channels )

	def remove_public_subscribe( self, bassAsset ):
		self.msgQueue.put ( {'event': 'remove', 'asset': bassAsset} )

	def _remove_pubilc_subscribe( self, bassAsset=None ):
		if bassAsset == None:
			raise Exception ( "TestOkxWebSocket 策略不能为None" )
		instId = str ( bassAsset.upper () + "-USDT-SWAP" ).strip ()
		temp = list ( filter ( lambda x: x['instId'] != instId, self.channels ) )
		self.channels = None
		self.channels = temp
		delchannels = []
		delchannels.append (
			{"channel": "mark-price", "instId": instId} )
		delchannels.append (
			{"channel": "tickers", "instId": instId} )
		sub_param = {'op': 'unsubscribe', 'args': delchannels}
		msg = json.dumps ( sub_param )
		self.ws.send ( msg )
		if len ( self.channels ) == 0:
			self.close ()

	def messagePublish( self, message ):

		print ( "messagePublish:", message )
		# 订单频道
		if message['arg']['channel'] == 'orders':
			msg = message['data'][0]
			print ( "TestOkxWebSocket 订单频道 result-->", msg )
		# 行情频道
		elif message['arg']['channel'] == 'tickers':
			print ( "TestOkxWebSocket 行情频道 result-->", message )

		# 标记价格频道
		elif message['arg']['channel'] == 'mark-price':
			print ( "TestOkxWebSocket 标记价格频道 result-->", message )
		else:
			print ( "TestOkxWebSocket 频道 result-->", message )

	def on_message( self, msg ):
		message = json.loads ( msg )
		# 订阅事件
		if 'event' in message:
			# 登录
			if message['event'] == 'login' and message['code'] == '0':
				params = self.get_subscribe ()
				print ( "TestOkxWebSocket 登录中 -->", params )
				self.ws.send ( params )
			# 订阅成功
			elif message['event'] == 'subscribe':
				# 订阅成功
				print ( "TestOkxWebSocket 频道订阅成功 -->", message )
			elif message['event'] == 'unsubscribe':
				print ( "TestOkxWebSocket 频道取消订阅成功 -->", message )

			else:
				print ( "TestOkxWebSocket 频道error --> OKX WebSocket 错误码", message )
		# 数据类型
		elif 'event' not in message:
			self.messagePublish ( message )

	def on_error( self, msg ):
		print ( "on_error===", msg )

	def on_pong( self, msg ):
		print ( "on_pong--->", msg )
		# 动态订阅
		if not self.msgQueue is None:
			if not self.msgQueue.empty ():
				qmsg = self.msgQueue.get ()
				if qmsg['event'] == 'add':
					self._add_pubilc_subscribe ( qmsg['asset'] )
				elif qmsg['event'] == 'remove':
					self._remove_pubilc_subscribe ( qmsg['asset'] )

	def on_close( self ):
		print ( "### TestOkxWebSocket 关闭WebSocket ###" )

	def on_open( self ):
		print ( "### TestOkxWebSocket 开启WebSocket ###" )
		if self.login:
			req_param = self.login_params ()
		else:
			req_param = self.get_subscribe ()
		self.ws.send ( req_param )

	def listenStreams( self ):
		print ( 'TestOkxWebSocket before 监听WebSocket' )
		websocket.enableTrace ( True )
		# 期待一个文字字符串'pong'作为回应。如果在N秒内未收到，请发出错误或重新连接。
		self.ws = websocket.WebSocketApp ( self.url,
										   on_message=self.on_message,
										   on_error=self.on_error,
										   on_open=self.on_open,
										   on_pong=self.on_pong,
										   on_close=self.on_close )

		print ( 'TestOkxWebSocket after 监听WebSocket' )
		self.ws.run_forever ( ping_interval=15 )

	def listenOnThread( self ):
		def run():
			self.listenStreams ()

		thread.start_new_thread ( run, () )

	def close( self ):
		self.ws.close ()
		self.ws = None

	# 登录参数
	def login_params( self ):
		timestamp = self.get_local_timestamp ()
		message = str ( timestamp ) + 'GET' + '/users/self/verify'
		mac = hmac.new ( bytes ( self.secretKey, encoding='utf8' ), bytes ( message, encoding='utf-8' ),
						 digestmod='sha256' )
		d = mac.digest ()
		sign = base64.b64encode ( d )
		login_param = {"op": "login", "args": [{"apiKey": self.apiKey,
												"passphrase": self.passphrase,
												"timestamp": timestamp,
												"sign": sign.decode ( "utf-8" )}]}
		login_str = json.dumps ( login_param )
		return login_str

	# 生成订阅请求头
	def get_subscribe( self ):
		sub_param = {'op': 'subscribe', 'args': self.channels}
		return json.dumps ( sub_param )

	def get_local_timestamp( self ):
		return int ( time.time () )


if __name__ == '__main__':
	os.environ["http_proxy"] = "http://127.0.0.1:10809"
	os.environ["https_proxy"] = "https://127.0.0.1:10809"
	socks.set_default_proxy ( socks.SOCKS5, "127.0.0.1", 10808 )
	socket.socket = socks.socksocket

	api_key = "xxxxx"
	secret_key = "xxxx"
	passphrase = "xxxxx"
	t = TestOkxWebSocket ( api_key, secret_key, passphrase, True )
	t.private_subscribe ()
	i = 0
	while (True):
		if i == 2:
			t.public_subscribe ( 'ETH' )
		elif i == 3:
			t.public_subscribe ( 'TRX' )
		elif i == 4:
			t.remove_public_subscribe ( 'BTC' )
		elif i == 5:
			t.remove_public_subscribe ( 'ETH' )
		elif i == 6:
			t.remove_public_subscribe ( 'TRX' )
		elif i == 7:
			t.public_subscribe ( 'ETC' )
		i = i + 1
		time.sleep ( 10 )