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"""
22from __future__ import print_function
23
24import socket
25import struct
26import threading
27
28import six
29
30# websocket modules
31from ._abnf import *
32from ._exceptions import *
33from ._handshake import *
34from ._http import *
35from ._logging import *
36from ._socket import *
37from ._utils import *
38
39__all__ = ['WebSocket', 'create_connection']
40
41"""
42websocket python client.
43=========================
44
45This version support only hybi-13.
46Please see http://tools.ietf.org/html/rfc6455 for protocol.
47"""
48
49
50class WebSocket(object):
51    """
52    Low level WebSocket interface.
53    This class is based on
54      The WebSocket protocol draft-hixie-thewebsocketprotocol-76
55      http://tools.ietf.org/html/draft-hixie-thewebsocketprotocol-76
56
57    We can connect to the websocket server and send/receive data.
58    The following example is an echo client.
59
60    >>> import websocket
61    >>> ws = websocket.WebSocket()
62    >>> ws.connect("ws://echo.websocket.org")
63    >>> ws.send("Hello, Server")
64    >>> ws.recv()
65    'Hello, Server'
66    >>> ws.close()
67
68    get_mask_key: a callable to produce new mask keys, see the set_mask_key
69      function's docstring for more details
70    sockopt: values for socket.setsockopt.
71        sockopt must be tuple and each element is argument of sock.setsockopt.
72    sslopt: dict object for ssl socket option.
73    fire_cont_frame: fire recv event for each cont frame. default is False
74    enable_multithread: if set to True, lock send method.
75    skip_utf8_validation: skip utf8 validation.
76    """
77
78    def __init__(self, get_mask_key=None, sockopt=None, sslopt=None,
79                 fire_cont_frame=False, enable_multithread=False,
80                 skip_utf8_validation=False, **_):
81        """
82        Initialize WebSocket object.
83        """
84        self.sock_opt = sock_opt(sockopt, sslopt)
85        self.handshake_response = None
86        self.sock = None
87
88        self.connected = False
89        self.get_mask_key = get_mask_key
90        # These buffer over the build-up of a single frame.
91        self.frame_buffer = frame_buffer(self._recv, skip_utf8_validation)
92        self.cont_frame = continuous_frame(
93            fire_cont_frame, skip_utf8_validation)
94
95        if enable_multithread:
96            self.lock = threading.Lock()
97        else:
98            self.lock = NoLock()
99
100    def __iter__(self):
101        """
102        Allow iteration over websocket, implying sequential `recv` executions.
103        """
104        while True:
105            yield self.recv()
106
107    def __next__(self):
108        return self.recv()
109
110    def next(self):
111        return self.__next__()
112
113    def fileno(self):
114        return self.sock.fileno()
115
116    def set_mask_key(self, func):
117        """
118        set function to create musk key. You can customize mask key generator.
119        Mainly, this is for testing purpose.
120
121        func: callable object. the func takes 1 argument as integer.
122              The argument means length of mask key.
123              This func must return string(byte array),
124              which length is argument specified.
125        """
126        self.get_mask_key = func
127
128    def gettimeout(self):
129        """
130        Get the websocket timeout(second).
131        """
132        return self.sock_opt.timeout
133
134    def settimeout(self, timeout):
135        """
136        Set the timeout to the websocket.
137
138        timeout: timeout time(second).
139        """
140        self.sock_opt.timeout = timeout
141        if self.sock:
142            self.sock.settimeout(timeout)
143
144    timeout = property(gettimeout, settimeout)
145
146    def getsubprotocol(self):
147        """
148        get subprotocol
149        """
150        if self.handshake_response:
151            return self.handshake_response.subprotocol
152        else:
153            return None
154
155    subprotocol = property(getsubprotocol)
156
157    def getstatus(self):
158        """
159        get handshake status
160        """
161        if self.handshake_response:
162            return self.handshake_response.status
163        else:
164            return None
165
166    status = property(getstatus)
167
168    def getheaders(self):
169        """
170        get handshake response header
171        """
172        if self.handshake_response:
173            return self.handshake_response.headers
174        else:
175            return None
176
177    headers = property(getheaders)
178
179    def connect(self, url, **options):
180        """
181        Connect to url. url is websocket url scheme.
182        ie. ws://host:port/resource
183        You can customize using 'options'.
184        If you set "header" list object, you can set your own custom header.
185
186        >>> ws = WebSocket()
187        >>> ws.connect("ws://echo.websocket.org/",
188                ...     header=["User-Agent: MyProgram",
189                ...             "x-custom: header"])
190
191        timeout: socket timeout time. This value is integer.
192                 if you set None for this value,
193                 it means "use default_timeout value"
194
195        options: "header" -> custom http header list or dict.
196                 "cookie" -> cookie value.
197                 "origin" -> custom origin url.
198                 "host"   -> custom host header string.
199                 "http_proxy_host" - http proxy host name.
200                 "http_proxy_port" - http proxy port. If not set, set to 80.
201                 "http_no_proxy"   - host names, which doesn't use proxy.
202                 "http_proxy_auth" - http proxy auth information.
203                                     tuple of username and password.
204                                     default is None
205                 "subprotocols" - array of available sub protocols.
206                                  default is None.
207                 "socket" - pre-initialized stream socket.
208
209        """
210        self.sock, addrs = connect(url, self.sock_opt, proxy_info(**options),
211                                   options.pop('socket', None))
212
213        try:
214            self.handshake_response = handshake(self.sock, *addrs, **options)
215            self.connected = True
216        except:
217            if self.sock:
218                self.sock.close()
219                self.sock = None
220            raise
221
222    def send(self, payload, opcode=ABNF.OPCODE_TEXT):
223        """
224        Send the data as string.
225
226        payload: Payload must be utf-8 string or unicode,
227                  if the opcode is OPCODE_TEXT.
228                  Otherwise, it must be string(byte array)
229
230        opcode: operation code to send. Please see OPCODE_XXX.
231        """
232
233        frame = ABNF.create_frame(payload, opcode)
234        return self.send_frame(frame)
235
236    def send_frame(self, frame):
237        """
238        Send the data frame.
239
240        frame: frame data created  by ABNF.create_frame
241
242        >>> ws = create_connection("ws://echo.websocket.org/")
243        >>> frame = ABNF.create_frame("Hello", ABNF.OPCODE_TEXT)
244        >>> ws.send_frame(frame)
245        >>> cont_frame = ABNF.create_frame("My name is ", ABNF.OPCODE_CONT, 0)
246        >>> ws.send_frame(frame)
247        >>> cont_frame = ABNF.create_frame("Foo Bar", ABNF.OPCODE_CONT, 1)
248        >>> ws.send_frame(frame)
249
250        """
251        if self.get_mask_key:
252            frame.get_mask_key = self.get_mask_key
253        data = frame.format()
254        length = len(data)
255        trace("send: " + repr(data))
256
257        with self.lock:
258            while data:
259                l = self._send(data)
260                data = data[l:]
261
262        return length
263
264    def send_binary(self, payload):
265        return self.send(payload, ABNF.OPCODE_BINARY)
266
267    def ping(self, payload=""):
268        """
269        send ping data.
270
271        payload: data payload to send server.
272        """
273        if isinstance(payload, six.text_type):
274            payload = payload.encode("utf-8")
275        self.send(payload, ABNF.OPCODE_PING)
276
277    def pong(self, payload):
278        """
279        send pong data.
280
281        payload: data payload to send server.
282        """
283        if isinstance(payload, six.text_type):
284            payload = payload.encode("utf-8")
285        self.send(payload, ABNF.OPCODE_PONG)
286
287    def recv(self):
288        """
289        Receive string data(byte array) from the server.
290
291        return value: string(byte array) value.
292        """
293        opcode, data = self.recv_data()
294        if six.PY3 and opcode == ABNF.OPCODE_TEXT:
295            return data.decode("utf-8")
296        elif opcode == ABNF.OPCODE_TEXT or opcode == ABNF.OPCODE_BINARY:
297            return data
298        else:
299            return ''
300
301    def recv_data(self, control_frame=False):
302        """
303        Receive data with operation code.
304
305        control_frame: a boolean flag indicating whether to return control frame
306        data, defaults to False
307
308        return  value: tuple of operation code and string(byte array) value.
309        """
310        opcode, frame = self.recv_data_frame(control_frame)
311        return opcode, frame.data
312
313    def recv_data_frame(self, control_frame=False):
314        """
315        Receive data with operation code.
316
317        control_frame: a boolean flag indicating whether to return control frame
318        data, defaults to False
319
320        return  value: tuple of operation code and string(byte array) value.
321        """
322        while True:
323            frame = self.recv_frame()
324            if not frame:
325                # handle error:
326                # 'NoneType' object has no attribute 'opcode'
327                raise WebSocketProtocolException(
328                    "Not a valid frame %s" % frame)
329            elif frame.opcode in (ABNF.OPCODE_TEXT, ABNF.OPCODE_BINARY, ABNF.OPCODE_CONT):
330                self.cont_frame.validate(frame)
331                self.cont_frame.add(frame)
332
333                if self.cont_frame.is_fire(frame):
334                    return self.cont_frame.extract(frame)
335
336            elif frame.opcode == ABNF.OPCODE_CLOSE:
337                self.send_close()
338                return frame.opcode, frame
339            elif frame.opcode == ABNF.OPCODE_PING:
340                if len(frame.data) < 126:
341                    self.pong(frame.data)
342                else:
343                    raise WebSocketProtocolException(
344                        "Ping message is too long")
345                if control_frame:
346                    return frame.opcode, frame
347            elif frame.opcode == ABNF.OPCODE_PONG:
348                if control_frame:
349                    return frame.opcode, frame
350
351    def recv_frame(self):
352        """
353        receive data as frame from server.
354
355        return value: ABNF frame object.
356        """
357        return self.frame_buffer.recv_frame()
358
359    def send_close(self, status=STATUS_NORMAL, reason=six.b("")):
360        """
361        send close data to the server.
362
363        status: status code to send. see STATUS_XXX.
364
365        reason: the reason to close. This must be string or bytes.
366        """
367        if status < 0 or status >= ABNF.LENGTH_16:
368            raise ValueError("code is invalid range")
369        self.connected = False
370        self.send(struct.pack('!H', status) + reason, ABNF.OPCODE_CLOSE)
371
372    def close(self, status=STATUS_NORMAL, reason=six.b(""), timeout=3):
373        """
374        Close Websocket object
375
376        status: status code to send. see STATUS_XXX.
377
378        reason: the reason to close. This must be string.
379
380        timeout: timeout until receive a close frame.
381            If None, it will wait forever until receive a close frame.
382        """
383        if self.connected:
384            if status < 0 or status >= ABNF.LENGTH_16:
385                raise ValueError("code is invalid range")
386
387            try:
388                self.connected = False
389                self.send(struct.pack('!H', status) +
390                          reason, ABNF.OPCODE_CLOSE)
391                sock_timeout = self.sock.gettimeout()
392                self.sock.settimeout(timeout)
393                try:
394                    frame = self.recv_frame()
395                    if isEnabledForError():
396                        recv_status = struct.unpack("!H", frame.data)[0]
397                        if recv_status != STATUS_NORMAL:
398                            error("close status: " + repr(recv_status))
399                except:
400                    pass
401                self.sock.settimeout(sock_timeout)
402                self.sock.shutdown(socket.SHUT_RDWR)
403            except:
404                pass
405
406        self.shutdown()
407
408    def abort(self):
409        """
410        Low-level asynchronous abort, wakes up other threads that are waiting in recv_*
411        """
412        if self.connected:
413            self.sock.shutdown(socket.SHUT_RDWR)
414
415    def shutdown(self):
416        """close socket, immediately."""
417        if self.sock:
418            self.sock.close()
419            self.sock = None
420            self.connected = False
421
422    def _send(self, data):
423        return send(self.sock, data)
424
425    def _recv(self, bufsize):
426        try:
427            return recv(self.sock, bufsize)
428        except WebSocketConnectionClosedException:
429            if self.sock:
430                self.sock.close()
431            self.sock = None
432            self.connected = False
433            raise
434
435
436def create_connection(url, timeout=None, class_=WebSocket, **options):
437    """
438    connect to url and return websocket object.
439
440    Connect to url and return the WebSocket object.
441    Passing optional timeout parameter will set the timeout on the socket.
442    If no timeout is supplied,
443    the global default timeout setting returned by getdefauttimeout() is used.
444    You can customize using 'options'.
445    If you set "header" list object, you can set your own custom header.
446
447    >>> conn = create_connection("ws://echo.websocket.org/",
448         ...     header=["User-Agent: MyProgram",
449         ...             "x-custom: header"])
450
451
452    timeout: socket timeout time. This value is integer.
453             if you set None for this value,
454             it means "use default_timeout value"
455
456    class_: class to instantiate when creating the connection. It has to implement
457            settimeout and connect. It's __init__ should be compatible with
458            WebSocket.__init__, i.e. accept all of it's kwargs.
459    options: "header" -> custom http header list or dict.
460             "cookie" -> cookie value.
461             "origin" -> custom origin url.
462             "host"   -> custom host header string.
463             "http_proxy_host" - http proxy host name.
464             "http_proxy_port" - http proxy port. If not set, set to 80.
465             "http_no_proxy"   - host names, which doesn't use proxy.
466             "http_proxy_auth" - http proxy auth information.
467                                    tuple of username and password.
468                                    default is None
469             "enable_multithread" -> enable lock for multithread.
470             "sockopt" -> socket options
471             "sslopt" -> ssl option
472             "subprotocols" - array of available sub protocols.
473                              default is None.
474             "skip_utf8_validation" - skip utf8 validation.
475             "socket" - pre-initialized stream socket.
476    """
477    sockopt = options.pop("sockopt", [])
478    sslopt = options.pop("sslopt", {})
479    fire_cont_frame = options.pop("fire_cont_frame", False)
480    enable_multithread = options.pop("enable_multithread", False)
481    skip_utf8_validation = options.pop("skip_utf8_validation", False)
482    websock = class_(sockopt=sockopt, sslopt=sslopt,
483                     fire_cont_frame=fire_cont_frame,
484                     enable_multithread=enable_multithread,
485                     skip_utf8_validation=skip_utf8_validation, **options)
486    websock.settimeout(timeout if timeout is not None else getdefaulttimeout())
487    websock.connect(url, **options)
488    return websock
489