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