You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
78 lines
2.2 KiB
78 lines
2.2 KiB
4 years ago
|
import json
|
||
|
|
import sys
|
||
|
|
|
||
|
|
import gevent
|
||
|
|
|
||
|
|
if sys.version_info[0] < 3:
|
||
|
|
from urlparse import urlparse, urlunparse
|
||
|
|
else:
|
||
|
|
from urllib.parse import urlparse, urlunparse
|
||
|
|
|
||
|
|
from websocket import create_connection
|
||
|
|
from ._transport import Transport
|
||
|
|
|
||
|
|
|
||
|
|
class WebSocketsTransport(Transport):
|
||
|
|
def __init__(self, session, connection):
|
||
|
|
Transport.__init__(self, session, connection)
|
||
|
|
self.ws = None
|
||
|
|
self.__requests = {}
|
||
|
|
|
||
|
|
def _get_name(self):
|
||
|
|
return 'webSockets'
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def __get_ws_url_from(url):
|
||
|
|
parsed = urlparse(url)
|
||
|
|
scheme = 'wss' if parsed.scheme == 'https' else 'ws'
|
||
|
|
url_data = (scheme, parsed.netloc, parsed.path, parsed.params, parsed.query, parsed.fragment)
|
||
|
|
|
||
|
|
return urlunparse(url_data)
|
||
|
|
|
||
|
|
def start(self):
|
||
|
|
ws_url = self.__get_ws_url_from(self._get_url('connect'))
|
||
|
|
|
||
|
|
self.ws = create_connection(ws_url,
|
||
|
|
header=self.__get_headers(),
|
||
|
|
cookie=self.__get_cookie_str(),
|
||
|
|
enable_multithread=True)
|
||
|
|
self._session.get(self._get_url('start'))
|
||
|
|
|
||
|
|
def _receive():
|
||
|
|
try:
|
||
|
|
for notification in self.ws:
|
||
|
|
self._handle_notification(notification)
|
||
|
|
except ConnectionError:
|
||
|
|
raise ConnectionError
|
||
|
|
|
||
|
|
return _receive
|
||
|
|
|
||
|
|
def send(self, data):
|
||
|
|
self.ws.send(json.dumps(data))
|
||
|
|
gevent.sleep()
|
||
|
|
|
||
|
|
def close(self):
|
||
|
|
self.ws.close()
|
||
|
|
|
||
|
|
def accept(self, negotiate_data):
|
||
|
|
return bool(negotiate_data['TryWebSockets'])
|
||
|
|
|
||
|
|
class HeadersLoader(object):
|
||
|
|
def __init__(self, headers):
|
||
|
|
self.headers = headers
|
||
|
|
|
||
|
|
def __get_headers(self):
|
||
|
|
headers = self._session.headers
|
||
|
|
loader = WebSocketsTransport.HeadersLoader(headers)
|
||
|
|
|
||
|
|
if self._session.auth:
|
||
|
|
self._session.auth(loader)
|
||
|
|
|
||
|
|
return ['%s: %s' % (name, headers[name]) for name in headers]
|
||
|
|
|
||
|
|
def __get_cookie_str(self):
|
||
|
|
return '; '.join([
|
||
|
|
'%s=%s' % (name, value)
|
||
|
|
for name, value in self._session.cookies.items()
|
||
|
|
])
|