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