py-as-zhong/exchange/gate_usdt_swap_ws.py

from re import sub
import aiohttp
import time
import asyncio
import zlib
import json, ujson
import zlib
import hmac, sys
import base64, csv, random
import traceback, hashlib
import logging, logging.handlers
import utils
import model
from decimal import Decimal


def inflate(data):
    '''
        解压缩数据
    '''
    decompress = zlib.decompressobj(-zlib.MAX_WBITS)
    inflated = decompress.decompress(data)
    inflated += decompress.flush()
    return inflated

def empty_call(msg):
    pass

class GateUsdtSwapWs:

    def __init__(self, params:model.ClientParams, colo=0, is_print=0):
        if colo:
            print('使用colo高速线路')
            self.URL = 'wss://fxws-private.gateapi.io/v4/ws/usdt'
        else:
            self.URL = 'wss://fx-ws.gateio.ws/v4/ws/usdt'
        self.params = params
        self.name = self.params.name
        self.base = self.params.pair.split('_')[0].upper()
        self.quote = self.params.pair.split('_')[1].upper()
        self.symbol = self.base + '_' + self.quote
        self.callback = {
            "onMarket":self.save_market,
            "onPosition":empty_call,
            "onTicker":empty_call,
            "onDepth":empty_call,
            "onEquity":empty_call,
            "onOrder":empty_call,
            "onTrade":empty_call,
            "onExit":empty_call,
            }
        self.is_print = is_print
        self.proxy = None
        if 'win' in sys.platform:
            self.proxy = self.params.proxy
        self.logger = self.get_logger()
        self.ticker_info = {"name":self.name,'bp':0.0,'ap':0.0}
        self.stop_flag = 0
        self.update_t = 0
        self.max_buy = 0.0
        self.min_sell = 0.0
        self.buy_v = 0.0
        self.buy_q = 0.0
        self.sell_v = 0.0
        self.sell_q = 0.0
        # 过期检查
        self.public_update_time = time.time()
        self.private_update_time = time.time()
        self.expired_time = 300
        self.multiplier = None
        self.depth = []
        #### 指定发包ip
        iplist = utils.get_local_ip_list()
        self.ip = iplist[int(self.params.ip)]
    
    def get_logger(self):
        logger = logging.getLogger(__name__)
        logger.setLevel(logging.DEBUG)
        # log to txt
        formatter = logging.Formatter('[%(asctime)s] - %(levelname)s - %(message)s')
        handler = logging.handlers.RotatingFileHandler(f"log.log",maxBytes=1024*1024)
        handler.setLevel(logging.DEBUG)
        handler.setFormatter(formatter)
        logger.addHandler(handler)
        return logger

    def gen_signed(self, channel, event, timestamp):
        # 为消息签名
        api_key = self.params.access_key
        api_secret = self.params.secret_key

        s = 'channel=%s&event=%s&time=%d' % (channel, event, timestamp)
        sign = hmac.new(api_secret.encode('utf-8'), s.encode('utf-8'), hashlib.sha512).hexdigest()
        return {'method': 'api_key', 'KEY': api_key, 'SIGN': sign}

    def save_market(self, msg):
        date = time.strftime('%Y-%m-%d',time.localtime())
        interval = self.params.interval
        if msg:
            exchange = msg['name']
            if len(msg['data']) > 1:
                with open(f'./history/{exchange}_{self.symbol}_{interval}_{date}.csv',
                            'a',
                            newline='',
                            encoding='utf-8') as f:
                    writer = csv.writer(f, delimiter=',')
                    writer.writerow(msg['data'])
                if self.is_print:print(f'写入行情 {self.symbol}')

    def _update_depth(self, msg):
        self.public_update_time = time.time()
        if msg['t'] > self.update_t:
            self.update_t = msg['t']
            self.ticker_info["bp"] = float(msg['bids'][0]['p'])
            self.ticker_info["ap"] = float(msg['asks'][0]['p'])
            self.callback['onTicker'](self.ticker_info)
            ##### 标准化深度
            mp = (self.ticker_info["bp"] + self.ticker_info["ap"])*0.5
            step = mp * utils.EFF_RANGE / utils.LEVEL
            bp = []
            ap = []
            bv = [0 for _ in range(utils.LEVEL)]
            av = [0 for _ in range(utils.LEVEL)]
            for i in range(utils.LEVEL):
                bp.append(self.ticker_info["bp"]-step*i)
            for i in range(utils.LEVEL):
                ap.append(self.ticker_info["ap"]+step*i)
            # 
            price_thre = self.ticker_info["bp"] - step
            index = 0
            for bid in msg['bids']:
                price = float(bid['p'])
                amount = float(bid['s'])
                if price > price_thre:
                    bv[index] += amount
                else:
                    price_thre -= step
                    index += 1
                    if index == utils.LEVEL:
                        break
                    bv[index] += amount
            price_thre = self.ticker_info["ap"] + step
            index = 0
            for ask in msg['asks']:
                price = float(ask['p'])
                amount = float(ask['s'])
                if price < price_thre:
                    av[index] += amount
                else:
                    price_thre += step
                    index += 1
                    if index == utils.LEVEL:
                        break
                    av[index] += amount
            self.depth = bp + bv + ap + av
            self.callback['onDepth']({'name':self.name,'data':self.depth})
        else:   
            self.logger.error(f"收到过时的depth推送 {self.update_t}")
    
    def _update_trade(self, msg):
        self.public_update_time = time.time()
        for i in msg:
            amount = float(i['size'])*self.multiplier
            price = float(i['price'])
            side = "buy" if amount > 0.0 else "sell"
            if price > self.max_buy or self.max_buy == 0.0:
                self.max_buy = price
            if price < self.min_sell or self.min_sell == 0.0:
                self.min_sell = price
            if side == 'buy':
                self.buy_q += amount
                self.buy_v += amount*price
            elif side == 'sell':
                self.sell_q += amount
                self.sell_v += amount*price

    def _update_account(self, msg):
        self.private_update_time = time.time()
        for i in msg:
            if self.symbol in i['text']:
                cash = float(i['balance'])
                self.callback['onEquity'] = {
                    self.quote:cash
                }
                self.logger.debug(f"ws cash {cash}")

    def _update_order(self, msg):
        self.private_update_time = time.time()
        self.logger.debug(f"ws订单推送 {msg}")
        for i in msg:
            if i['status'] in ['open']:
                order_event = dict()
                order_event['filled'] = 0
                order_event['filled_price'] = 0
                order_event['client_id'] = i["text"]
                order_event['order_id'] = i['id']
                order_event['status'] = "NEW"
                self.callback['onOrder'](order_event)
            elif i['status'] in ['finished']:
                order_event = dict()
                filled_paper = Decimal(abs(float(i["size"]))) - Decimal(abs(float(i["left"])))
                filled_amount = filled_paper*Decimal(str(self.multiplier))
                order_event['filled'] = float(filled_amount)
                order_event['filled_price'] = float(i["fill_price"])
                order_event['client_id'] = i["text"]
                order_event['order_id'] = i['id']
                order_event['fee'] = 0.0
                order_event['status'] = "REMOVE"
                self.callback['onOrder'](order_event)
                # 根据成交信息更新仓位信息 因为账户信息推送有延迟
                # 但订单信息和账户信息到达先后时间可能有前有后 可能平仓 账户先置零仓位 然后sell成交达到 导致仓位变成负数

    def _update_usertrade(self, msg):
        '''暂时不用'''
        return

    def _update_position(self, msg):
        self.private_update_time = time.time()
        long_pos, short_pos = 0, 0
        long_avg, short_avg = 0, 0 
        for i in msg:
            if i['contract'] == self.symbol:
                size = float(i['size'])*self.multiplier
                if size > 0:
                    long_pos = abs(size)
                    long_avg = float(i['entry_price'])
                if size < 0:
                    short_pos = abs(size)
                    short_avg = float(i['entry_price'])
        pos = model.Position()
        pos.longPos = long_pos
        pos.longAvg = long_avg
        pos.shortPos = short_pos
        pos.shortAvg = short_avg
        self.callback['onPosition'](pos)

    def _get_data(self):
        market_data = self.depth + [self.max_buy, self.min_sell]
        self.max_buy = 0.0
        self.min_sell = 0.0
        self.buy_v = 0.0
        self.buy_q = 0.0
        self.sell_v = 0.0
        self.sell_q = 0.0
        return {'name': self.name,'data':market_data}

    async def go(self):
        interval = float(self.params.interval)
        if self.is_print:print(f'Ws循环器启动 interval {interval}')
        ### onTrade
        while 1:
            try:
                # 更新市场信息
                market_data = self._get_data()
                self.callback['onMarket'](market_data)
            except:
                traceback.print_exc()
            await asyncio.sleep(interval)

    def get_sign(self, message):
        h = hmac.new(self.params.secret_key.encode("utf8"), message.encode("utf8"), hashlib.sha512)
        return h.hexdigest()

    def _get_uid(self):
        pass

    def generate_signature(self, method, uri, query_param=None, body=None):
        t = time.time()
        m = hashlib.sha512()
        m.update((body or "").encode('utf-8'))
        hashed_payload = m.hexdigest()
        s = '%s\n%s\n%s\n%s\n%s' % (method, uri, query_param or "", hashed_payload, t)
        sign = hmac.new(self.params.secret_key.encode('utf-8'), s.encode('utf-8'), hashlib.sha512).hexdigest()
        return {'KEY': self.params.access_key, 'Timestamp': str(t), 'SIGN': sign}

    async def run(self, is_auth=0, sub_trade=0, sub_fast=0):
        while True:
            try:
                # 重置更新时间
                self.public_update_time = time.time()
                self.private_update_time = time.time()
                ping_time = time.time()
                # 获取uid
                headers = {
                    "Accept": "application/json",
                    "Content-type": "application/json"
                }
                if is_auth:
                    user_id = ""
                    uri = "/api/v4/wallet/fee"
                    query_param = ''
                    sign_headers = self.generate_signature('GET', uri, query_param)
                    headers.update(sign_headers)
                    async with aiohttp.ClientSession(connector = aiohttp.TCPConnector(
                                limit=50,
                                keepalive_timeout=120,
                                verify_ssl=False,
                                local_addr=(self.ip,0)
                            )) as session:
                        response = await session.get(
                            "https://api.gateio.ws" + uri,  
                            headers=headers,
                            proxy=self.proxy
                        )
                        res = await response.json()
                        user_id = str(res['user_id'])
                        print(f"uid {user_id}")
                # 获取合约乘数
                async with aiohttp.ClientSession(connector = aiohttp.TCPConnector(
                                limit=50,
                                keepalive_timeout=120,
                                verify_ssl=False,
                                local_addr=(self.ip,0)
                            )) as session:
                    uri = "/api/v4/futures/usdt/contracts"
                    response = await session.get(
                        "https://api.gateio.ws" + uri,  
                        headers=headers,
                        proxy=self.proxy
                    )
                    res = await response.json()
                    if res:
                        for i in res:
                            if self.symbol == i['name']:
                                self.multiplier = float(i['quanto_multiplier'])
                                print(f"contract multiplier {self.multiplier}")
                    # 尝试连接
                print(f'{self.name} 尝试连接ws')
                ws_url = self.URL
                async with aiohttp.ClientSession(
                        connector = aiohttp.TCPConnector(
                                limit=50,
                                keepalive_timeout=120,
                                verify_ssl=False,
                                local_addr=(self.ip,0)
                            )
                        ).ws_connect(
                            ws_url,
                            proxy=self.proxy,
                            timeout=30,
                            receive_timeout=30,
                            ) as _ws:
                    print(f'{self.name} ws连接成功')
                    # 登陆
                    if is_auth:
                        # userorders
                        current_time = int(time.time())
                        channel = "futures.orders"
                        sub_str = {
                            "time": current_time,
                            "channel": channel,
                            "event": "subscribe", 
                            "payload": [user_id,self.symbol]
                            }
                        sub_str["auth"] = self.gen_signed(sub_str['channel'], sub_str['event'], sub_str['time'])
                        await _ws.send_str(ujson.dumps(sub_str))
                        # positions
                        current_time = int(time.time())
                        channel = "futures.positions"
                        sub_str = {
                            "time": current_time,
                            "channel": channel,
                            "event": "subscribe", 
                            "payload": [user_id,self.symbol]
                            }
                        sub_str["auth"] = self.gen_signed(sub_str['channel'], sub_str['event'], sub_str['time'])
                        await _ws.send_str(ujson.dumps(sub_str))
                        # usertrades
                        # current_time = int(time.time())
                        # channel = "futures.usertrades"
                        # sub_str = {
                        #     "time": current_time,
                        #     "channel": channel,
                        #     "event": "subscribe", 
                        #     "payload": [self.symbol]
                        # }
                        # message = 'channel=%s&event=%s&time=%d' % (channel, "subscribe", current_time)
                        # sub_str["auth"] = {
                        #     "method": "api_key",
                        #     "KEY": self.params.access_key,
                        #     "SIGN": self.get_sign(message)}
                        # await _ws.send_str(ujson.dumps(sub_str))
                        # balance
                        current_time = int(time.time())
                        channel = "futures.balances"
                        sub_str = {
                            "time": current_time,
                            "channel": channel,
                            "event": "subscribe", 
                            "payload": [user_id]
                        }
                        sub_str["auth"] = self.gen_signed(sub_str['channel'], sub_str['event'], sub_str['time'])
                        await _ws.send_str(ujson.dumps(sub_str))
                    if sub_trade:
                        # public trade
                        current_time = int(time.time())
                        channel = "futures.trades"
                        sub_str = {
                            "time": current_time,
                            "channel": channel,
                            "event": "subscribe", 
                            "payload": [self.symbol]
                        }
                        await _ws.send_str(ujson.dumps(sub_str))
                    # 订阅
                    # tickers 速度慢
                    # current_time = int(time.time())
                    # channel = "futures.tickers"
                    # sub_str = {
                    #     "time": current_time,
                    #     "channel": channel,
                    #     "event": "subscribe", 
                    #     "payload": [self.symbol]
                    # }
                    # await _ws.send_str(ujson.dumps(sub_str))
                    # depth
                    current_time = int(time.time())
                    channel = "futures.order_book"
                    sub_str = {
                        "time": current_time,
                        "channel": channel,
                        "event": "subscribe", 
                        "payload": [self.symbol,"20","0"]
                    }
                    await _ws.send_str(ujson.dumps(sub_str))
                    while True:
                        # 停机信号
                        if self.stop_flag:
                            await _ws.close()
                            return
                        # 接受消息
                        try:
                            msg = await _ws.receive(timeout=10)
                        except:
                            print(f'{self.name} ws长时间没有收到消息 准备重连...')
                            self.logger.error(f'{self.name} ws长时间没有收到消息 准备重连...')
                            break
                        msg = ujson.loads(msg.data)
                        # 处理消息
                        if msg['event'] in ['update', 'all']:
                            if msg['channel'] == 'futures.order_book':self._update_depth(msg['result'])
                            elif msg['channel'] == 'futures.balances':self._update_account(msg['result'])
                            elif msg['channel'] == 'futures.orders':self._update_order(msg['result'])
                            # elif msg['channel'] == 'futures.usertrades':self._update_usertrade(msg['result'])
                            elif msg['channel'] == 'futures.positions':self._update_position(msg['result'])
                            elif msg['channel'] == 'futures.trades':self._update_trade(msg['result'])
                        else:
                            pass
                        # pong
                        if time.time() - ping_time > 5:
                            await _ws.send_str('{"time": %d, "channel" : "futures.ping"}' % int(time.time()))
                            ping_time = time.time()
                        if is_auth:                            
                            if time.time() - self.private_update_time > self.expired_time*5:
                                raise Exception('长期未更新私有信息重连')
                        if time.time() - self.public_update_time > self.expired_time:
                            raise Exception('长期未更新公有信息重连')
            except:
                traceback.print_exc()
                print(f'{self.name} ws连接失败 开始重连...')
                self.logger.error(f'{self.name} ws连接失败 开始重连...')
                self.logger.error(traceback.format_exc())
                await asyncio.sleep(1)