_app.py 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351
  1. """
  2. websocket - WebSocket client library for Python
  3. Copyright (C) 2010 Hiroki Ohtani(liris)
  4. This library is free software; you can redistribute it and/or
  5. modify it under the terms of the GNU Lesser General Public
  6. License as published by the Free Software Foundation; either
  7. version 2.1 of the License, or (at your option) any later version.
  8. This library is distributed in the hope that it will be useful,
  9. but WITHOUT ANY WARRANTY; without even the implied warranty of
  10. MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
  11. Lesser General Public License for more details.
  12. You should have received a copy of the GNU Lesser General Public
  13. License along with this library; if not, write to the Free Software
  14. Foundation, Inc., 51 Franklin Street, Fifth Floor,
  15. Boston, MA 02110-1335 USA
  16. """
  17. """
  18. WebSocketApp provides higher level APIs.
  19. """
  20. import inspect
  21. import select
  22. import sys
  23. import threading
  24. import time
  25. import traceback
  26. import six
  27. from ._abnf import ABNF
  28. from ._core import WebSocket, getdefaulttimeout
  29. from ._exceptions import *
  30. from . import _logging
  31. __all__ = ["WebSocketApp"]
  32. class Dispatcher:
  33. def __init__(self, app, ping_timeout):
  34. self.app = app
  35. self.ping_timeout = ping_timeout
  36. def read(self, sock, read_callback, check_callback):
  37. while self.app.sock.connected:
  38. r, w, e = select.select(
  39. (self.app.sock.sock, ), (), (), self.ping_timeout)
  40. if r:
  41. if not read_callback():
  42. break
  43. check_callback()
  44. class SSLDispacther:
  45. def __init__(self, app, ping_timeout):
  46. self.app = app
  47. self.ping_timeout = ping_timeout
  48. def read(self, sock, read_callback, check_callback):
  49. while self.app.sock.connected:
  50. r = self.select()
  51. if r:
  52. if not read_callback():
  53. break
  54. check_callback()
  55. def select(self):
  56. sock = self.app.sock.sock
  57. if sock.pending():
  58. return [sock,]
  59. r, w, e = select.select((sock, ), (), (), self.ping_timeout)
  60. return r
  61. class WebSocketApp(object):
  62. """
  63. Higher level of APIs are provided.
  64. The interface is like JavaScript WebSocket object.
  65. """
  66. def __init__(self, url, header=None,
  67. on_open=None, on_message=None, on_error=None,
  68. on_close=None, on_ping=None, on_pong=None,
  69. on_cont_message=None,
  70. keep_running=True, get_mask_key=None, cookie=None,
  71. subprotocols=None,
  72. on_data=None):
  73. """
  74. url: websocket url.
  75. header: custom header for websocket handshake.
  76. on_open: callable object which is called at opening websocket.
  77. this function has one argument. The argument is this class object.
  78. on_message: callable object which is called when received data.
  79. on_message has 2 arguments.
  80. The 1st argument is this class object.
  81. The 2nd argument is utf-8 string which we get from the server.
  82. on_error: callable object which is called when we get error.
  83. on_error has 2 arguments.
  84. The 1st argument is this class object.
  85. The 2nd argument is exception object.
  86. on_close: callable object which is called when closed the connection.
  87. this function has one argument. The argument is this class object.
  88. on_cont_message: callback object which is called when receive continued
  89. frame data.
  90. on_cont_message has 3 arguments.
  91. The 1st argument is this class object.
  92. The 2nd argument is utf-8 string which we get from the server.
  93. The 3rd argument is continue flag. if 0, the data continue
  94. to next frame data
  95. on_data: callback object which is called when a message received.
  96. This is called before on_message or on_cont_message,
  97. and then on_message or on_cont_message is called.
  98. on_data has 4 argument.
  99. The 1st argument is this class object.
  100. The 2nd argument is utf-8 string which we get from the server.
  101. The 3rd argument is data type. ABNF.OPCODE_TEXT or ABNF.OPCODE_BINARY will be came.
  102. The 4th argument is continue flag. if 0, the data continue
  103. keep_running: this parameter is obsolete and ignored.
  104. get_mask_key: a callable to produce new mask keys,
  105. see the WebSocket.set_mask_key's docstring for more information
  106. subprotocols: array of available sub protocols. default is None.
  107. """
  108. self.url = url
  109. self.header = header if header is not None else []
  110. self.cookie = cookie
  111. self.on_open = on_open
  112. self.on_message = on_message
  113. self.on_data = on_data
  114. self.on_error = on_error
  115. self.on_close = on_close
  116. self.on_ping = on_ping
  117. self.on_pong = on_pong
  118. self.on_cont_message = on_cont_message
  119. self.keep_running = False
  120. self.get_mask_key = get_mask_key
  121. self.sock = None
  122. self.last_ping_tm = 0
  123. self.last_pong_tm = 0
  124. self.subprotocols = subprotocols
  125. def send(self, data, opcode=ABNF.OPCODE_TEXT):
  126. """
  127. send message.
  128. data: message to send. If you set opcode to OPCODE_TEXT,
  129. data must be utf-8 string or unicode.
  130. opcode: operation code of data. default is OPCODE_TEXT.
  131. """
  132. if not self.sock or self.sock.send(data, opcode) == 0:
  133. raise WebSocketConnectionClosedException(
  134. "Connection is already closed.")
  135. def close(self, **kwargs):
  136. """
  137. close websocket connection.
  138. """
  139. self.keep_running = False
  140. if self.sock:
  141. self.sock.close(**kwargs)
  142. self.sock = None
  143. def _send_ping(self, interval, event):
  144. while not event.wait(interval):
  145. self.last_ping_tm = time.time()
  146. if self.sock:
  147. try:
  148. self.sock.ping()
  149. except Exception as ex:
  150. _logging.warning("send_ping routine terminated: {}".format(ex))
  151. break
  152. def run_forever(self, sockopt=None, sslopt=None,
  153. ping_interval=0, ping_timeout=None,
  154. http_proxy_host=None, http_proxy_port=None,
  155. http_no_proxy=None, http_proxy_auth=None,
  156. skip_utf8_validation=False,
  157. host=None, origin=None, dispatcher=None,
  158. suppress_origin = False, proxy_type=None):
  159. """
  160. run event loop for WebSocket framework.
  161. This loop is infinite loop and is alive during websocket is available.
  162. sockopt: values for socket.setsockopt.
  163. sockopt must be tuple
  164. and each element is argument of sock.setsockopt.
  165. sslopt: ssl socket optional dict.
  166. ping_interval: automatically send "ping" command
  167. every specified period(second)
  168. if set to 0, not send automatically.
  169. ping_timeout: timeout(second) if the pong message is not received.
  170. http_proxy_host: http proxy host name.
  171. http_proxy_port: http proxy port. If not set, set to 80.
  172. http_no_proxy: host names, which doesn't use proxy.
  173. skip_utf8_validation: skip utf8 validation.
  174. host: update host header.
  175. origin: update origin header.
  176. dispatcher: customize reading data from socket.
  177. suppress_origin: suppress outputting origin header.
  178. Returns
  179. -------
  180. False if caught KeyboardInterrupt
  181. True if other exception was raised during a loop
  182. """
  183. if ping_timeout is not None and ping_timeout <= 0:
  184. ping_timeout = None
  185. if ping_timeout and ping_interval and ping_interval <= ping_timeout:
  186. raise WebSocketException("Ensure ping_interval > ping_timeout")
  187. if not sockopt:
  188. sockopt = []
  189. if not sslopt:
  190. sslopt = {}
  191. if self.sock:
  192. raise WebSocketException("socket is already opened")
  193. thread = None
  194. self.keep_running = True
  195. self.last_ping_tm = 0
  196. self.last_pong_tm = 0
  197. def teardown(close_frame=None):
  198. """
  199. Tears down the connection.
  200. If close_frame is set, we will invoke the on_close handler with the
  201. statusCode and reason from there.
  202. """
  203. if thread and thread.isAlive():
  204. event.set()
  205. thread.join()
  206. self.keep_running = False
  207. if self.sock:
  208. self.sock.close()
  209. close_args = self._get_close_args(
  210. close_frame.data if close_frame else None)
  211. self._callback(self.on_close, *close_args)
  212. self.sock = None
  213. try:
  214. self.sock = WebSocket(
  215. self.get_mask_key, sockopt=sockopt, sslopt=sslopt,
  216. fire_cont_frame=self.on_cont_message is not None,
  217. skip_utf8_validation=skip_utf8_validation,
  218. enable_multithread=True if ping_interval else False)
  219. self.sock.settimeout(getdefaulttimeout())
  220. self.sock.connect(
  221. self.url, header=self.header, cookie=self.cookie,
  222. http_proxy_host=http_proxy_host,
  223. http_proxy_port=http_proxy_port, http_no_proxy=http_no_proxy,
  224. http_proxy_auth=http_proxy_auth, subprotocols=self.subprotocols,
  225. host=host, origin=origin, suppress_origin=suppress_origin,
  226. proxy_type=proxy_type)
  227. if not dispatcher:
  228. dispatcher = self.create_dispatcher(ping_timeout)
  229. self._callback(self.on_open)
  230. if ping_interval:
  231. event = threading.Event()
  232. thread = threading.Thread(
  233. target=self._send_ping, args=(ping_interval, event))
  234. thread.setDaemon(True)
  235. thread.start()
  236. def read():
  237. if not self.keep_running:
  238. return teardown()
  239. op_code, frame = self.sock.recv_data_frame(True)
  240. if op_code == ABNF.OPCODE_CLOSE:
  241. return teardown(frame)
  242. elif op_code == ABNF.OPCODE_PING:
  243. self._callback(self.on_ping, frame.data)
  244. elif op_code == ABNF.OPCODE_PONG:
  245. self.last_pong_tm = time.time()
  246. self._callback(self.on_pong, frame.data)
  247. elif op_code == ABNF.OPCODE_CONT and self.on_cont_message:
  248. self._callback(self.on_data, frame.data,
  249. frame.opcode, frame.fin)
  250. self._callback(self.on_cont_message,
  251. frame.data, frame.fin)
  252. else:
  253. data = frame.data
  254. if six.PY3 and op_code == ABNF.OPCODE_TEXT:
  255. data = data.decode("utf-8")
  256. self._callback(self.on_data, data, frame.opcode, True)
  257. self._callback(self.on_message, data)
  258. return True
  259. def check():
  260. if (ping_timeout):
  261. has_timeout_expired = time.time() - self.last_ping_tm > ping_timeout
  262. has_pong_not_arrived_after_last_ping = self.last_pong_tm - self.last_ping_tm < 0
  263. has_pong_arrived_too_late = self.last_pong_tm - self.last_ping_tm > ping_timeout
  264. if (self.last_ping_tm
  265. and has_timeout_expired
  266. and (has_pong_not_arrived_after_last_ping or has_pong_arrived_too_late)):
  267. raise WebSocketTimeoutException("ping/pong timed out")
  268. return True
  269. dispatcher.read(self.sock.sock, read, check)
  270. except (Exception, KeyboardInterrupt, SystemExit) as e:
  271. self._callback(self.on_error, e)
  272. if isinstance(e, SystemExit):
  273. # propagate SystemExit further
  274. raise
  275. teardown()
  276. return not isinstance(e, KeyboardInterrupt)
  277. def create_dispatcher(self, ping_timeout):
  278. timeout = ping_timeout or 10
  279. if self.sock.is_ssl():
  280. return SSLDispacther(self, timeout)
  281. return Dispatcher(self, timeout)
  282. def _get_close_args(self, data):
  283. """ this functions extracts the code, reason from the close body
  284. if they exists, and if the self.on_close except three arguments """
  285. # if the on_close callback is "old", just return empty list
  286. if sys.version_info < (3, 0):
  287. if not self.on_close or len(inspect.getargspec(self.on_close).args) != 3:
  288. return []
  289. else:
  290. if not self.on_close or len(inspect.getfullargspec(self.on_close).args) != 3:
  291. return []
  292. if data and len(data) >= 2:
  293. code = 256 * six.byte2int(data[0:1]) + six.byte2int(data[1:2])
  294. reason = data[2:].decode('utf-8')
  295. return [code, reason]
  296. return [None, None]
  297. def _callback(self, callback, *args):
  298. if callback:
  299. try:
  300. if inspect.ismethod(callback):
  301. callback(*args)
  302. else:
  303. callback(self, *args)
  304. except Exception as e:
  305. _logging.error("error from callback {}: {}".format(callback, e))
  306. if _logging.isEnabledForDebug():
  307. _, _, tb = sys.exc_info()
  308. traceback.print_tb(tb)