Вы не можете выбрать более 25 тем Темы должны начинаться с буквы или цифры, могут содержать дефисы(-) и должны содержать не более 35 символов.

652 строки
25KB

  1. import logging
  2. try:
  3. import queue
  4. except ImportError: # pragma: no cover
  5. import Queue as queue
  6. import signal
  7. import threading
  8. import time
  9. import six
  10. from six.moves import urllib
  11. try:
  12. import requests
  13. except ImportError: # pragma: no cover
  14. requests = None
  15. try:
  16. import websocket
  17. except ImportError: # pragma: no cover
  18. websocket = None
  19. from . import exceptions
  20. from . import packet
  21. from . import payload
  22. default_logger = logging.getLogger('engineio.client')
  23. connected_clients = []
  24. if six.PY2: # pragma: no cover
  25. ConnectionError = OSError
  26. def signal_handler(sig, frame):
  27. """SIGINT handler.
  28. Disconnect all active clients and then invoke the original signal handler.
  29. """
  30. for client in connected_clients[:]:
  31. if client.is_asyncio_based():
  32. client.start_background_task(client.disconnect, abort=True)
  33. else:
  34. client.disconnect(abort=True)
  35. return original_signal_handler(sig, frame)
  36. original_signal_handler = signal.signal(signal.SIGINT, signal_handler)
  37. class Client(object):
  38. """An Engine.IO client.
  39. This class implements a fully compliant Engine.IO web client with support
  40. for websocket and long-polling transports.
  41. :param logger: To enable logging set to ``True`` or pass a logger object to
  42. use. To disable logging set to ``False``. The default is
  43. ``False``.
  44. :param json: An alternative json module to use for encoding and decoding
  45. packets. Custom json modules must have ``dumps`` and ``loads``
  46. functions that are compatible with the standard library
  47. versions.
  48. :param request_timeout: A timeout in seconds for requests. The default is
  49. 5 seconds.
  50. """
  51. event_names = ['connect', 'disconnect', 'message']
  52. def __init__(self, logger=False, json=None, request_timeout=5):
  53. self.handlers = {}
  54. self.base_url = None
  55. self.transports = None
  56. self.current_transport = None
  57. self.sid = None
  58. self.upgrades = None
  59. self.ping_interval = None
  60. self.ping_timeout = None
  61. self.pong_received = True
  62. self.http = None
  63. self.ws = None
  64. self.read_loop_task = None
  65. self.write_loop_task = None
  66. self.ping_loop_task = None
  67. self.ping_loop_event = self.create_event()
  68. self.queue = None
  69. self.state = 'disconnected'
  70. if json is not None:
  71. packet.Packet.json = json
  72. if not isinstance(logger, bool):
  73. self.logger = logger
  74. else:
  75. self.logger = default_logger
  76. if not logging.root.handlers and \
  77. self.logger.level == logging.NOTSET:
  78. if logger:
  79. self.logger.setLevel(logging.INFO)
  80. else:
  81. self.logger.setLevel(logging.ERROR)
  82. self.logger.addHandler(logging.StreamHandler())
  83. self.request_timeout = request_timeout
  84. def is_asyncio_based(self):
  85. return False
  86. def on(self, event, handler=None):
  87. """Register an event handler.
  88. :param event: The event name. Can be ``'connect'``, ``'message'`` or
  89. ``'disconnect'``.
  90. :param handler: The function that should be invoked to handle the
  91. event. When this parameter is not given, the method
  92. acts as a decorator for the handler function.
  93. Example usage::
  94. # as a decorator:
  95. @eio.on('connect')
  96. def connect_handler():
  97. print('Connection request')
  98. # as a method:
  99. def message_handler(msg):
  100. print('Received message: ', msg)
  101. eio.send('response')
  102. eio.on('message', message_handler)
  103. """
  104. if event not in self.event_names:
  105. raise ValueError('Invalid event')
  106. def set_handler(handler):
  107. self.handlers[event] = handler
  108. return handler
  109. if handler is None:
  110. return set_handler
  111. set_handler(handler)
  112. def connect(self, url, headers={}, transports=None,
  113. engineio_path='engine.io'):
  114. """Connect to an Engine.IO server.
  115. :param url: The URL of the Engine.IO server. It can include custom
  116. query string parameters if required by the server.
  117. :param headers: A dictionary with custom headers to send with the
  118. connection request.
  119. :param transports: The list of allowed transports. Valid transports
  120. are ``'polling'`` and ``'websocket'``. If not
  121. given, the polling transport is connected first,
  122. then an upgrade to websocket is attempted.
  123. :param engineio_path: The endpoint where the Engine.IO server is
  124. installed. The default value is appropriate for
  125. most cases.
  126. Example usage::
  127. eio = engineio.Client()
  128. eio.connect('http://localhost:5000')
  129. """
  130. if self.state != 'disconnected':
  131. raise ValueError('Client is not in a disconnected state')
  132. valid_transports = ['polling', 'websocket']
  133. if transports is not None:
  134. if isinstance(transports, six.string_types):
  135. transports = [transports]
  136. transports = [transport for transport in transports
  137. if transport in valid_transports]
  138. if not transports:
  139. raise ValueError('No valid transports provided')
  140. self.transports = transports or valid_transports
  141. self.queue = self.create_queue()
  142. return getattr(self, '_connect_' + self.transports[0])(
  143. url, headers, engineio_path)
  144. def wait(self):
  145. """Wait until the connection with the server ends.
  146. Client applications can use this function to block the main thread
  147. during the life of the connection.
  148. """
  149. if self.read_loop_task:
  150. self.read_loop_task.join()
  151. def send(self, data, binary=None):
  152. """Send a message to a client.
  153. :param data: The data to send to the client. Data can be of type
  154. ``str``, ``bytes``, ``list`` or ``dict``. If a ``list``
  155. or ``dict``, the data will be serialized as JSON.
  156. :param binary: ``True`` to send packet as binary, ``False`` to send
  157. as text. If not given, unicode (Python 2) and str
  158. (Python 3) are sent as text, and str (Python 2) and
  159. bytes (Python 3) are sent as binary.
  160. """
  161. self._send_packet(packet.Packet(packet.MESSAGE, data=data,
  162. binary=binary))
  163. def disconnect(self, abort=False):
  164. """Disconnect from the server.
  165. :param abort: If set to ``True``, do not wait for background tasks
  166. associated with the connection to end.
  167. """
  168. if self.state == 'connected':
  169. self._send_packet(packet.Packet(packet.CLOSE))
  170. self.queue.put(None)
  171. self.state = 'disconnecting'
  172. self._trigger_event('disconnect', run_async=False)
  173. if self.current_transport == 'websocket':
  174. self.ws.close()
  175. if not abort:
  176. self.read_loop_task.join()
  177. self.state = 'disconnected'
  178. try:
  179. connected_clients.remove(self)
  180. except ValueError: # pragma: no cover
  181. pass
  182. self._reset()
  183. def transport(self):
  184. """Return the name of the transport currently in use.
  185. The possible values returned by this function are ``'polling'`` and
  186. ``'websocket'``.
  187. """
  188. return self.current_transport
  189. def start_background_task(self, target, *args, **kwargs):
  190. """Start a background task.
  191. This is a utility function that applications can use to start a
  192. background task.
  193. :param target: the target function to execute.
  194. :param args: arguments to pass to the function.
  195. :param kwargs: keyword arguments to pass to the function.
  196. This function returns an object compatible with the `Thread` class in
  197. the Python standard library. The `start()` method on this object is
  198. already called by this function.
  199. """
  200. th = threading.Thread(target=target, args=args, kwargs=kwargs)
  201. th.start()
  202. return th
  203. def sleep(self, seconds=0):
  204. """Sleep for the requested amount of time."""
  205. return time.sleep(seconds)
  206. def create_queue(self, *args, **kwargs):
  207. """Create a queue object."""
  208. q = queue.Queue(*args, **kwargs)
  209. q.Empty = queue.Empty
  210. return q
  211. def create_event(self, *args, **kwargs):
  212. """Create an event object."""
  213. return threading.Event(*args, **kwargs)
  214. def _reset(self):
  215. self.state = 'disconnected'
  216. self.sid = None
  217. def _connect_polling(self, url, headers, engineio_path):
  218. """Establish a long-polling connection to the Engine.IO server."""
  219. if requests is None: # pragma: no cover
  220. # not installed
  221. self.logger.error('requests package is not installed -- cannot '
  222. 'send HTTP requests!')
  223. return
  224. self.base_url = self._get_engineio_url(url, engineio_path, 'polling')
  225. self.logger.info('Attempting polling connection to ' + self.base_url)
  226. r = self._send_request(
  227. 'GET', self.base_url + self._get_url_timestamp(), headers=headers,
  228. timeout=self.request_timeout)
  229. if r is None:
  230. self._reset()
  231. raise exceptions.ConnectionError(
  232. 'Connection refused by the server')
  233. if r.status_code != 200:
  234. raise exceptions.ConnectionError(
  235. 'Unexpected status code {} in server response'.format(
  236. r.status_code))
  237. try:
  238. p = payload.Payload(encoded_payload=r.content)
  239. except ValueError:
  240. six.raise_from(exceptions.ConnectionError(
  241. 'Unexpected response from server'), None)
  242. open_packet = p.packets[0]
  243. if open_packet.packet_type != packet.OPEN:
  244. raise exceptions.ConnectionError(
  245. 'OPEN packet not returned by server')
  246. self.logger.info(
  247. 'Polling connection accepted with ' + str(open_packet.data))
  248. self.sid = open_packet.data['sid']
  249. self.upgrades = open_packet.data['upgrades']
  250. self.ping_interval = open_packet.data['pingInterval'] / 1000.0
  251. self.ping_timeout = open_packet.data['pingTimeout'] / 1000.0
  252. self.current_transport = 'polling'
  253. self.base_url += '&sid=' + self.sid
  254. self.state = 'connected'
  255. connected_clients.append(self)
  256. self._trigger_event('connect', run_async=False)
  257. for pkt in p.packets[1:]:
  258. self._receive_packet(pkt)
  259. if 'websocket' in self.upgrades and 'websocket' in self.transports:
  260. # attempt to upgrade to websocket
  261. if self._connect_websocket(url, headers, engineio_path):
  262. # upgrade to websocket succeeded, we're done here
  263. return
  264. # start background tasks associated with this client
  265. self.ping_loop_task = self.start_background_task(self._ping_loop)
  266. self.write_loop_task = self.start_background_task(self._write_loop)
  267. self.read_loop_task = self.start_background_task(
  268. self._read_loop_polling)
  269. def _connect_websocket(self, url, headers, engineio_path):
  270. """Establish or upgrade to a WebSocket connection with the server."""
  271. if websocket is None: # pragma: no cover
  272. # not installed
  273. self.logger.warning('websocket-client package not installed, only '
  274. 'polling transport is available')
  275. return False
  276. websocket_url = self._get_engineio_url(url, engineio_path, 'websocket')
  277. if self.sid:
  278. self.logger.info(
  279. 'Attempting WebSocket upgrade to ' + websocket_url)
  280. upgrade = True
  281. websocket_url += '&sid=' + self.sid
  282. else:
  283. upgrade = False
  284. self.base_url = websocket_url
  285. self.logger.info(
  286. 'Attempting WebSocket connection to ' + websocket_url)
  287. # get the cookies from the long-polling connection so that they can
  288. # also be sent the the WebSocket route
  289. cookies = None
  290. if self.http:
  291. cookies = '; '.join(["{}={}".format(cookie.name, cookie.value)
  292. for cookie in self.http.cookies])
  293. try:
  294. ws = websocket.create_connection(
  295. websocket_url + self._get_url_timestamp(), header=headers,
  296. cookie=cookies)
  297. except ConnectionError:
  298. if upgrade:
  299. self.logger.warning(
  300. 'WebSocket upgrade failed: connection error')
  301. return False
  302. else:
  303. raise exceptions.ConnectionError('Connection error')
  304. if upgrade:
  305. p = packet.Packet(packet.PING,
  306. data=six.text_type('probe')).encode()
  307. try:
  308. ws.send(p)
  309. except Exception as e: # pragma: no cover
  310. self.logger.warning(
  311. 'WebSocket upgrade failed: unexpected send exception: %s',
  312. str(e))
  313. return False
  314. try:
  315. p = ws.recv()
  316. except Exception as e: # pragma: no cover
  317. self.logger.warning(
  318. 'WebSocket upgrade failed: unexpected recv exception: %s',
  319. str(e))
  320. return False
  321. pkt = packet.Packet(encoded_packet=p)
  322. if pkt.packet_type != packet.PONG or pkt.data != 'probe':
  323. self.logger.warning(
  324. 'WebSocket upgrade failed: no PONG packet')
  325. return False
  326. p = packet.Packet(packet.UPGRADE).encode()
  327. try:
  328. ws.send(p)
  329. except Exception as e: # pragma: no cover
  330. self.logger.warning(
  331. 'WebSocket upgrade failed: unexpected send exception: %s',
  332. str(e))
  333. return False
  334. self.current_transport = 'websocket'
  335. self.logger.info('WebSocket upgrade was successful')
  336. else:
  337. try:
  338. p = ws.recv()
  339. except Exception as e: # pragma: no cover
  340. raise exceptions.ConnectionError(
  341. 'Unexpected recv exception: ' + str(e))
  342. open_packet = packet.Packet(encoded_packet=p)
  343. if open_packet.packet_type != packet.OPEN:
  344. raise exceptions.ConnectionError('no OPEN packet')
  345. self.logger.info(
  346. 'WebSocket connection accepted with ' + str(open_packet.data))
  347. self.sid = open_packet.data['sid']
  348. self.upgrades = open_packet.data['upgrades']
  349. self.ping_interval = open_packet.data['pingInterval'] / 1000.0
  350. self.ping_timeout = open_packet.data['pingTimeout'] / 1000.0
  351. self.current_transport = 'websocket'
  352. self.state = 'connected'
  353. connected_clients.append(self)
  354. self._trigger_event('connect', run_async=False)
  355. self.ws = ws
  356. # start background tasks associated with this client
  357. self.ping_loop_task = self.start_background_task(self._ping_loop)
  358. self.write_loop_task = self.start_background_task(self._write_loop)
  359. self.read_loop_task = self.start_background_task(
  360. self._read_loop_websocket)
  361. return True
  362. def _receive_packet(self, pkt):
  363. """Handle incoming packets from the server."""
  364. packet_name = packet.packet_names[pkt.packet_type] \
  365. if pkt.packet_type < len(packet.packet_names) else 'UNKNOWN'
  366. self.logger.info(
  367. 'Received packet %s data %s', packet_name,
  368. pkt.data if not isinstance(pkt.data, bytes) else '<binary>')
  369. if pkt.packet_type == packet.MESSAGE:
  370. self._trigger_event('message', pkt.data, run_async=True)
  371. elif pkt.packet_type == packet.PONG:
  372. self.pong_received = True
  373. elif pkt.packet_type == packet.CLOSE:
  374. self.disconnect(abort=True)
  375. elif pkt.packet_type == packet.NOOP:
  376. pass
  377. else:
  378. self.logger.error('Received unexpected packet of type %s',
  379. pkt.packet_type)
  380. def _send_packet(self, pkt):
  381. """Queue a packet to be sent to the server."""
  382. if self.state != 'connected':
  383. return
  384. self.queue.put(pkt)
  385. self.logger.info(
  386. 'Sending packet %s data %s',
  387. packet.packet_names[pkt.packet_type],
  388. pkt.data if not isinstance(pkt.data, bytes) else '<binary>')
  389. def _send_request(
  390. self, method, url, headers=None, body=None,
  391. timeout=None): # pragma: no cover
  392. if self.http is None:
  393. self.http = requests.Session()
  394. try:
  395. return self.http.request(method, url, headers=headers, data=body,
  396. timeout=timeout)
  397. except requests.exceptions.RequestException as exc:
  398. self.logger.info('HTTP %s request to %s failed with error %s.',
  399. method, url, exc)
  400. def _trigger_event(self, event, *args, **kwargs):
  401. """Invoke an event handler."""
  402. run_async = kwargs.pop('run_async', False)
  403. if event in self.handlers:
  404. if run_async:
  405. return self.start_background_task(self.handlers[event], *args)
  406. else:
  407. try:
  408. return self.handlers[event](*args)
  409. except:
  410. self.logger.exception(event + ' handler error')
  411. def _get_engineio_url(self, url, engineio_path, transport):
  412. """Generate the Engine.IO connection URL."""
  413. engineio_path = engineio_path.strip('/')
  414. parsed_url = urllib.parse.urlparse(url)
  415. if transport == 'polling':
  416. scheme = 'http'
  417. elif transport == 'websocket':
  418. scheme = 'ws'
  419. else: # pragma: no cover
  420. raise ValueError('invalid transport')
  421. if parsed_url.scheme in ['https', 'wss']:
  422. scheme += 's'
  423. return ('{scheme}://{netloc}/{path}/?{query}'
  424. '{sep}transport={transport}&EIO=3').format(
  425. scheme=scheme, netloc=parsed_url.netloc,
  426. path=engineio_path, query=parsed_url.query,
  427. sep='&' if parsed_url.query else '',
  428. transport=transport)
  429. def _get_url_timestamp(self):
  430. """Generate the Engine.IO query string timestamp."""
  431. return '&t=' + str(time.time())
  432. def _ping_loop(self):
  433. """This background task sends a PING to the server at the requested
  434. interval.
  435. """
  436. self.pong_received = True
  437. self.ping_loop_event.clear()
  438. while self.state == 'connected':
  439. if not self.pong_received:
  440. self.logger.info(
  441. 'PONG response has not been received, aborting')
  442. if self.ws:
  443. self.ws.shutdown()
  444. self.queue.put(None)
  445. break
  446. self.pong_received = False
  447. self._send_packet(packet.Packet(packet.PING))
  448. self.ping_loop_event.wait(timeout=self.ping_interval)
  449. self.logger.info('Exiting ping task')
  450. def _read_loop_polling(self):
  451. """Read packets by polling the Engine.IO server."""
  452. while self.state == 'connected':
  453. self.logger.info(
  454. 'Sending polling GET request to ' + self.base_url)
  455. r = self._send_request(
  456. 'GET', self.base_url + self._get_url_timestamp(),
  457. timeout=max(self.ping_interval, self.ping_timeout) + 5)
  458. if r is None:
  459. self.logger.warning(
  460. 'Connection refused by the server, aborting')
  461. self.queue.put(None)
  462. break
  463. if r.status_code != 200:
  464. self.logger.warning('Unexpected status code %s in server '
  465. 'response, aborting', r.status_code)
  466. self.queue.put(None)
  467. break
  468. try:
  469. p = payload.Payload(encoded_payload=r.content)
  470. except ValueError:
  471. self.logger.warning(
  472. 'Unexpected packet from server, aborting')
  473. self.queue.put(None)
  474. break
  475. for pkt in p.packets:
  476. self._receive_packet(pkt)
  477. self.logger.info('Waiting for write loop task to end')
  478. self.write_loop_task.join()
  479. self.logger.info('Waiting for ping loop task to end')
  480. self.ping_loop_event.set()
  481. self.ping_loop_task.join()
  482. if self.state == 'connected':
  483. self._trigger_event('disconnect', run_async=False)
  484. try:
  485. connected_clients.remove(self)
  486. except ValueError: # pragma: no cover
  487. pass
  488. self._reset()
  489. self.logger.info('Exiting read loop task')
  490. def _read_loop_websocket(self):
  491. """Read packets from the Engine.IO WebSocket connection."""
  492. while self.state == 'connected':
  493. p = None
  494. try:
  495. p = self.ws.recv()
  496. except websocket.WebSocketConnectionClosedException:
  497. self.logger.warning(
  498. 'WebSocket connection was closed, aborting')
  499. self.queue.put(None)
  500. break
  501. except Exception as e:
  502. self.logger.info(
  503. 'Unexpected error "%s", aborting', str(e))
  504. self.queue.put(None)
  505. break
  506. if isinstance(p, six.text_type): # pragma: no cover
  507. p = p.encode('utf-8')
  508. pkt = packet.Packet(encoded_packet=p)
  509. self._receive_packet(pkt)
  510. self.logger.info('Waiting for write loop task to end')
  511. self.write_loop_task.join()
  512. self.logger.info('Waiting for ping loop task to end')
  513. self.ping_loop_event.set()
  514. self.ping_loop_task.join()
  515. if self.state == 'connected':
  516. self._trigger_event('disconnect', run_async=False)
  517. try:
  518. connected_clients.remove(self)
  519. except ValueError: # pragma: no cover
  520. pass
  521. self._reset()
  522. self.logger.info('Exiting read loop task')
  523. def _write_loop(self):
  524. """This background task sends packages to the server as they are
  525. pushed to the send queue.
  526. """
  527. while self.state == 'connected':
  528. # to simplify the timeout handling, use the maximum of the
  529. # ping interval and ping timeout as timeout, with an extra 5
  530. # seconds grace period
  531. timeout = max(self.ping_interval, self.ping_timeout) + 5
  532. packets = None
  533. try:
  534. packets = [self.queue.get(timeout=timeout)]
  535. except self.queue.Empty:
  536. self.logger.error('packet queue is empty, aborting')
  537. break
  538. if packets == [None]:
  539. self.queue.task_done()
  540. packets = []
  541. else:
  542. while True:
  543. try:
  544. packets.append(self.queue.get(block=False))
  545. except self.queue.Empty:
  546. break
  547. if packets[-1] is None:
  548. packets = packets[:-1]
  549. self.queue.task_done()
  550. break
  551. if not packets:
  552. # empty packet list returned -> connection closed
  553. break
  554. if self.current_transport == 'polling':
  555. p = payload.Payload(packets=packets)
  556. r = self._send_request(
  557. 'POST', self.base_url, body=p.encode(),
  558. headers={'Content-Type': 'application/octet-stream'},
  559. timeout=self.request_timeout)
  560. for pkt in packets:
  561. self.queue.task_done()
  562. if r is None:
  563. self.logger.warning(
  564. 'Connection refused by the server, aborting')
  565. break
  566. if r.status_code != 200:
  567. self.logger.warning('Unexpected status code %s in server '
  568. 'response, aborting', r.status_code)
  569. self._reset()
  570. break
  571. else:
  572. # websocket
  573. try:
  574. for pkt in packets:
  575. encoded_packet = pkt.encode(always_bytes=False)
  576. if pkt.binary:
  577. self.ws.send_binary(encoded_packet)
  578. else:
  579. self.ws.send(encoded_packet)
  580. self.queue.task_done()
  581. except websocket.WebSocketConnectionClosedException:
  582. self.logger.warning(
  583. 'WebSocket connection was closed, aborting')
  584. break
  585. self.logger.info('Exiting write loop task')