tools/telemetry/third_party/websocket-client/websocket.py

   1 """
   2 websocket - WebSocket client library for Python
   3
   4 Copyright (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., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
  19
  20 """
  21
  22
  23 import socket
  24
  25 try:
  26     import ssl
  27     from ssl import SSLError
  28     HAVE_SSL = True
  29 except ImportError:
  30     # dummy class of SSLError for ssl none-support environment.
  31     class SSLError(Exception):
  32         pass
  33
  34     HAVE_SSL = False
  35
  36 from urlparse import urlparse
  37 import os
  38 import array
  39 import struct
  40 import uuid
  41 import hashlib
  42 import base64
  43 import threading
  44 import time
  45 import logging
  46 import traceback
  47 import sys
  48
  49 """
  50 websocket python client.
  51 =========================
  52
  53 This version support only hybi-13.
  54 Please see http://tools.ietf.org/html/rfc6455 for protocol.
  55 """
  56
  57
  58 # websocket supported version.
  59 VERSION = 13
  60
  61 # closing frame status codes.
  62 STATUS_NORMAL = 1000
  63 STATUS_GOING_AWAY = 1001
  64 STATUS_PROTOCOL_ERROR = 1002
  65 STATUS_UNSUPPORTED_DATA_TYPE = 1003
  66 STATUS_STATUS_NOT_AVAILABLE = 1005
  67 STATUS_ABNORMAL_CLOSED = 1006
  68 STATUS_INVALID_PAYLOAD = 1007
  69 STATUS_POLICY_VIOLATION = 1008
  70 STATUS_MESSAGE_TOO_BIG = 1009
  71 STATUS_INVALID_EXTENSION = 1010
  72 STATUS_UNEXPECTED_CONDITION = 1011
  73 STATUS_TLS_HANDSHAKE_ERROR = 1015
  74
  75 logger = logging.getLogger()
  76
  77
  78 class WebSocketException(Exception):
  79     """
  80     websocket exeception class.
  81     """
  82     pass
  83
  84
  85 class WebSocketConnectionClosedException(WebSocketException):
  86     """
  87     If remote host closed the connection or some network error happened,
  88     this exception will be raised.
  89     """
  90     pass
  91
  92 class WebSocketTimeoutException(WebSocketException):
  93     """
  94     WebSocketTimeoutException will be raised at socket timeout during read/write data.
  95     """
  96     pass
  97
  98 default_timeout = None
  99 traceEnabled = False
 100
 101
 102 def enableTrace(tracable):
 103     """
 104     turn on/off the tracability.
 105
 106     tracable: boolean value. if set True, tracability is enabled.
 107     """
 108     global traceEnabled
 109     traceEnabled = tracable
 110     if tracable:
 111         if not logger.handlers:
 112             logger.addHandler(logging.StreamHandler())
 113         logger.setLevel(logging.DEBUG)
 114
 115
 116 def setdefaulttimeout(timeout):
 117     """
 118     Set the global timeout setting to connect.
 119
 120     timeout: default socket timeout time. This value is second.
 121     """
 122     global default_timeout
 123     default_timeout = timeout
 124
 125
 126 def getdefaulttimeout():
 127     """
 128     Return the global timeout setting(second) to connect.
 129     """
 130     return default_timeout
 131
 132
 133 def _parse_url(url):
 134     """
 135     parse url and the result is tuple of
 136     (hostname, port, resource path and the flag of secure mode)
 137
 138     url: url string.
 139     """
 140     if ":" not in url:
 141         raise ValueError("url is invalid")
 142
 143     scheme, url = url.split(":", 1)
 144
 145     parsed = urlparse(url, scheme="http")
 146     if parsed.hostname:
 147         hostname = parsed.hostname
 148     else:
 149         raise ValueError("hostname is invalid")
 150     port = 0
 151     if parsed.port:
 152         port = parsed.port
 153
 154     is_secure = False
 155     if scheme == "ws":
 156         if not port:
 157             port = 80
 158     elif scheme == "wss":
 159         is_secure = True
 160         if not port:
 161             port = 443
 162     else:
 163         raise ValueError("scheme %s is invalid" % scheme)
 164
 165     if parsed.path:
 166         resource = parsed.path
 167     else:
 168         resource = "/"
 169
 170     if parsed.query:
 171         resource += "?" + parsed.query
 172
 173     return (hostname, port, resource, is_secure)
 174
 175
 176 def create_connection(url, timeout=None, **options):
 177     """
 178     connect to url and return websocket object.
 179
 180     Connect to url and return the WebSocket object.
 181     Passing optional timeout parameter will set the timeout on the socket.
 182     If no timeout is supplied, the global default timeout setting returned by getdefauttimeout() is used.
 183     You can customize using 'options'.
 184     If you set "header" list object, you can set your own custom header.
 185
 186     >>> conn = create_connection("ws://echo.websocket.org/",
 187          ...     header=["User-Agent: MyProgram",
 188          ...             "x-custom: header"])
 189
 190
 191     timeout: socket timeout time. This value is integer.
 192              if you set None for this value, it means "use default_timeout value"
 193
 194     options: current support option is only "header".
 195              if you set header as dict value, the custom HTTP headers are added.
 196     """
 197     sockopt = options.get("sockopt", [])
 198     sslopt = options.get("sslopt", {})
 199     websock = WebSocket(sockopt=sockopt, sslopt=sslopt)
 200     websock.settimeout(timeout if timeout is not None else default_timeout)
 201     websock.connect(url, **options)
 202     return websock
 203
 204 _MAX_INTEGER = (1 << 32) -1
 205 _AVAILABLE_KEY_CHARS = range(0x21, 0x2f + 1) + range(0x3a, 0x7e + 1)
 206 _MAX_CHAR_BYTE = (1<<8) -1
 207
 208 # ref. Websocket gets an update, and it breaks stuff.
 209 # http://axod.blogspot.com/2010/06/websocket-gets-update-and-it-breaks.html
 210
 211
 212 def _create_sec_websocket_key():
 213     uid = uuid.uuid4()
 214     return base64.encodestring(uid.bytes).strip()
 215
 216
 217 _HEADERS_TO_CHECK = {
 218     "upgrade": "websocket",
 219     "connection": "upgrade",
 220     }
 221
 222
 223 class ABNF(object):
 224     """
 225     ABNF frame class.
 226     see http://tools.ietf.org/html/rfc5234
 227     and http://tools.ietf.org/html/rfc6455#section-5.2
 228     """
 229
 230     # operation code values.
 231     OPCODE_CONT   = 0x0
 232     OPCODE_TEXT   = 0x1
 233     OPCODE_BINARY = 0x2
 234     OPCODE_CLOSE  = 0x8
 235     OPCODE_PING   = 0x9
 236     OPCODE_PONG   = 0xa
 237
 238     # available operation code value tuple
 239     OPCODES = (OPCODE_CONT, OPCODE_TEXT, OPCODE_BINARY, OPCODE_CLOSE,
 240                 OPCODE_PING, OPCODE_PONG)
 241
 242     # opcode human readable string
 243     OPCODE_MAP = {
 244         OPCODE_CONT: "cont",
 245         OPCODE_TEXT: "text",
 246         OPCODE_BINARY: "binary",
 247         OPCODE_CLOSE: "close",
 248         OPCODE_PING: "ping",
 249         OPCODE_PONG: "pong"
 250         }
 251
 252     # data length threashold.
 253     LENGTH_7  = 0x7d
 254     LENGTH_16 = 1 << 16
 255     LENGTH_63 = 1 << 63
 256
 257     def __init__(self, fin=0, rsv1=0, rsv2=0, rsv3=0,
 258                  opcode=OPCODE_TEXT, mask=1, data=""):
 259         """
 260         Constructor for ABNF.
 261         please check RFC for arguments.
 262         """
 263         self.fin = fin
 264         self.rsv1 = rsv1
 265         self.rsv2 = rsv2
 266         self.rsv3 = rsv3
 267         self.opcode = opcode
 268         self.mask = mask
 269         self.data = data
 270         self.get_mask_key = os.urandom
 271
 272     def __str__(self):
 273         return "fin=" + str(self.fin) \
 274                 + " opcode=" + str(self.opcode) \
 275                 + " data=" + str(self.data)
 276
 277     @staticmethod
 278     def create_frame(data, opcode):
 279         """
 280         create frame to send text, binary and other data.
 281
 282         data: data to send. This is string value(byte array).
 283             if opcode is OPCODE_TEXT and this value is uniocde,
 284             data value is conveted into unicode string, automatically.
 285
 286         opcode: operation code. please see OPCODE_XXX.
 287         """
 288         if opcode == ABNF.OPCODE_TEXT and isinstance(data, unicode):
 289             data = data.encode("utf-8")
 290         # mask must be set if send data from client
 291         return ABNF(1, 0, 0, 0, opcode, 1, data)
 292
 293     def format(self):
 294         """
 295         format this object to string(byte array) to send data to server.
 296         """
 297         if any(x not in (0, 1) for x in [self.fin, self.rsv1, self.rsv2, self.rsv3]):
 298             raise ValueError("not 0 or 1")
 299         if self.opcode not in ABNF.OPCODES:
 300             raise ValueError("Invalid OPCODE")
 301         length = len(self.data)
 302         if length >= ABNF.LENGTH_63:
 303             raise ValueError("data is too long")
 304
 305         frame_header = chr(self.fin << 7
 306                            | self.rsv1 << 6 | self.rsv2 << 5 | self.rsv3 << 4
 307                            | self.opcode)
 308         if length < ABNF.LENGTH_7:
 309             frame_header += chr(self.mask << 7 | length)
 310         elif length < ABNF.LENGTH_16:
 311             frame_header += chr(self.mask << 7 | 0x7e)
 312             frame_header += struct.pack("!H", length)
 313         else:
 314             frame_header += chr(self.mask << 7 | 0x7f)
 315             frame_header += struct.pack("!Q", length)
 316
 317         if not self.mask:
 318             return frame_header + self.data
 319         else:
 320             mask_key = self.get_mask_key(4)
 321             return frame_header + self._get_masked(mask_key)
 322
 323     def _get_masked(self, mask_key):
 324         s = ABNF.mask(mask_key, self.data)
 325         return mask_key + "".join(s)
 326
 327     @staticmethod
 328     def mask(mask_key, data):
 329         """
 330         mask or unmask data. Just do xor for each byte
 331
 332         mask_key: 4 byte string(byte).
 333
 334         data: data to mask/unmask.
 335         """
 336         _m = array.array("B", mask_key)
 337         _d = array.array("B", data)
 338         for i in xrange(len(_d)):
 339             _d[i] ^= _m[i % 4]
 340         return _d.tostring()
 341
 342
 343 class WebSocket(object):
 344     """
 345     Low level WebSocket interface.
 346     This class is based on
 347       The WebSocket protocol draft-hixie-thewebsocketprotocol-76
 348       http://tools.ietf.org/html/draft-hixie-thewebsocketprotocol-76
 349
 350     We can connect to the websocket server and send/recieve data.
 351     The following example is a echo client.
 352
 353     >>> import websocket
 354     >>> ws = websocket.WebSocket()
 355     >>> ws.connect("ws://echo.websocket.org")
 356     >>> ws.send("Hello, Server")
 357     >>> ws.recv()
 358     'Hello, Server'
 359     >>> ws.close()
 360
 361     get_mask_key: a callable to produce new mask keys, see the set_mask_key
 362       function's docstring for more details
 363     sockopt: values for socket.setsockopt.
 364         sockopt must be tuple and each element is argument of sock.setscokopt.
 365     sslopt: dict object for ssl socket option.
 366     """
 367
 368     def __init__(self, get_mask_key=None, sockopt=None, sslopt=None):
 369         """
 370         Initalize WebSocket object.
 371         """
 372         if sockopt is None:
 373             sockopt = []
 374         if sslopt is None:
 375             sslopt = {}
 376         self.connected = False
 377         self.sock = socket.socket()
 378         for opts in sockopt:
 379             self.sock.setsockopt(*opts)
 380         self.sslopt = sslopt
 381         self.get_mask_key = get_mask_key
 382         # Buffers over the packets from the layer beneath until desired amount
 383         # bytes of bytes are received.
 384         self._recv_buffer = []
 385         # These buffer over the build-up of a single frame.
 386         self._frame_header = None
 387         self._frame_length = None
 388         self._frame_mask = None
 389         self._cont_data = None
 390
 391     def fileno(self):
 392         return self.sock.fileno()
 393
 394     def set_mask_key(self, func):
 395         """
 396         set function to create musk key. You can custumize mask key generator.
 397         Mainly, this is for testing purpose.
 398
 399         func: callable object. the fuct must 1 argument as integer.
 400               The argument means length of mask key.
 401               This func must be return string(byte array),
 402               which length is argument specified.
 403         """
 404         self.get_mask_key = func
 405
 406     def gettimeout(self):
 407         """
 408         Get the websocket timeout(second).
 409         """
 410         return self.sock.gettimeout()
 411
 412     def settimeout(self, timeout):
 413         """
 414         Set the timeout to the websocket.
 415
 416         timeout: timeout time(second).
 417         """
 418         self.sock.settimeout(timeout)
 419
 420     timeout = property(gettimeout, settimeout)
 421
 422     def connect(self, url, **options):
 423         """
 424         Connect to url. url is websocket url scheme. ie. ws://host:port/resource
 425         You can customize using 'options'.
 426         If you set "header" dict object, you can set your own custom header.
 427
 428         >>> ws = WebSocket()
 429         >>> ws.connect("ws://echo.websocket.org/",
 430                 ...     header={"User-Agent: MyProgram",
 431                 ...             "x-custom: header"})
 432
 433         timeout: socket timeout time. This value is integer.
 434                  if you set None for this value,
 435                  it means "use default_timeout value"
 436
 437         options: current support option is only "header".
 438                  if you set header as dict value,
 439                  the custom HTTP headers are added.
 440
 441         """
 442         hostname, port, resource, is_secure = _parse_url(url)
 443         # TODO: we need to support proxy
 444         self.sock.connect((hostname, port))
 445         if is_secure:
 446             if HAVE_SSL:
 447                 if self.sslopt is None:
 448                     sslopt = {}
 449                 else:
 450                     sslopt = self.sslopt
 451                 self.sock = ssl.wrap_socket(self.sock, **sslopt)
 452             else:
 453                 raise WebSocketException("SSL not available.")
 454
 455         self._handshake(hostname, port, resource, **options)
 456
 457     def _handshake(self, host, port, resource, **options):
 458         sock = self.sock
 459         headers = []
 460         headers.append("GET %s HTTP/1.1" % resource)
 461         headers.append("Upgrade: websocket")
 462         headers.append("Connection: Upgrade")
 463         if port == 80:
 464             hostport = host
 465         else:
 466             hostport = "%s:%d" % (host, port)
 467         headers.append("Host: %s" % hostport)
 468
 469         if "origin" in options:
 470             headers.append("Origin: %s" % options["origin"])
 471         else:
 472             headers.append("Origin: http://%s" % hostport)
 473
 474         key = _create_sec_websocket_key()
 475         headers.append("Sec-WebSocket-Key: %s" % key)
 476         headers.append("Sec-WebSocket-Version: %s" % VERSION)
 477         if "header" in options:
 478             headers.extend(options["header"])
 479
 480         headers.append("")
 481         headers.append("")
 482
 483         header_str = "\r\n".join(headers)
 484         self._send(header_str)
 485         if traceEnabled:
 486             logger.debug("--- request header ---")
 487             logger.debug(header_str)
 488             logger.debug("-----------------------")
 489
 490         status, resp_headers = self._read_headers()
 491         if status != 101:
 492             self.close()
 493             raise WebSocketException("Handshake Status %d" % status)
 494
 495         success = self._validate_header(resp_headers, key)
 496         if not success:
 497             self.close()
 498             raise WebSocketException("Invalid WebSocket Header")
 499
 500         self.connected = True
 501
 502     def _validate_header(self, headers, key):
 503         for k, v in _HEADERS_TO_CHECK.iteritems():
 504             r = headers.get(k, None)
 505             if not r:
 506                 return False
 507             r = r.lower()
 508             if v != r:
 509                 return False
 510
 511         result = headers.get("sec-websocket-accept", None)
 512         if not result:
 513             return False
 514         result = result.lower()
 515
 516         value = key + "258EAFA5-E914-47DA-95CA-C5AB0DC85B11"
 517         hashed = base64.encodestring(hashlib.sha1(value).digest()).strip().lower()
 518         return hashed == result
 519
 520     def _read_headers(self):
 521         status = None
 522         headers = {}
 523         if traceEnabled:
 524             logger.debug("--- response header ---")
 525
 526         while True:
 527             line = self._recv_line()
 528             if line == "\r\n":
 529                 break
 530             line = line.strip()
 531             if traceEnabled:
 532                 logger.debug(line)
 533             if not status:
 534                 status_info = line.split(" ", 2)
 535                 status = int(status_info[1])
 536             else:
 537                 kv = line.split(":", 1)
 538                 if len(kv) == 2:
 539                     key, value = kv
 540                     headers[key.lower()] = value.strip().lower()
 541                 else:
 542                     raise WebSocketException("Invalid header")
 543
 544         if traceEnabled:
 545             logger.debug("-----------------------")
 546
 547         return status, headers
 548
 549     def send(self, payload, opcode=ABNF.OPCODE_TEXT):
 550         """
 551         Send the data as string.
 552
 553         payload: Payload must be utf-8 string or unicoce,
 554                   if the opcode is OPCODE_TEXT.
 555                   Otherwise, it must be string(byte array)
 556
 557         opcode: operation code to send. Please see OPCODE_XXX.
 558         """
 559         frame = ABNF.create_frame(payload, opcode)
 560         if self.get_mask_key:
 561             frame.get_mask_key = self.get_mask_key
 562         data = frame.format()
 563         length = len(data)
 564         if traceEnabled:
 565             logger.debug("send: " + repr(data))
 566         while data:
 567             l = self._send(data)
 568             data = data[l:]
 569         return length
 570
 571     def send_binary(self, payload):
 572         return self.send(payload, ABNF.OPCODE_BINARY)
 573
 574     def ping(self, payload=""):
 575         """
 576         send ping data.
 577
 578         payload: data payload to send server.
 579         """
 580         self.send(payload, ABNF.OPCODE_PING)
 581
 582     def pong(self, payload):
 583         """
 584         send pong data.
 585
 586         payload: data payload to send server.
 587         """
 588         self.send(payload, ABNF.OPCODE_PONG)
 589
 590     def recv(self):
 591         """
 592         Receive string data(byte array) from the server.
 593
 594         return value: string(byte array) value.
 595         """
 596         opcode, data = self.recv_data()
 597         return data
 598
 599     def recv_data(self):
 600         """
 601         Recieve data with operation code.
 602
 603         return  value: tuple of operation code and string(byte array) value.
 604         """
 605         while True:
 606             frame = self.recv_frame()
 607             if not frame:
 608                 # handle error:
 609                 # 'NoneType' object has no attribute 'opcode'
 610                 raise WebSocketException("Not a valid frame %s" % frame)
 611             elif frame.opcode in (ABNF.OPCODE_TEXT, ABNF.OPCODE_BINARY, ABNF.OPCODE_CONT):
 612                 if frame.opcode == ABNF.OPCODE_CONT and not self._cont_data:
 613                     raise WebSocketException("Illegal frame")
 614                 if self._cont_data:
 615                     self._cont_data[1] += frame.data
 616                 else:
 617                     self._cont_data = [frame.opcode, frame.data]
 618
 619                 if frame.fin:
 620                     data = self._cont_data
 621                     self._cont_data = None
 622                     return data
 623             elif frame.opcode == ABNF.OPCODE_CLOSE:
 624                 self.send_close()
 625                 return (frame.opcode, None)
 626             elif frame.opcode == ABNF.OPCODE_PING:
 627                 self.pong(frame.data)
 628
 629     def recv_frame(self):
 630         """
 631         recieve data as frame from server.
 632
 633         return value: ABNF frame object.
 634         """
 635         # Header
 636         if self._frame_header is None:
 637             self._frame_header = self._recv_strict(2)
 638         b1 = ord(self._frame_header[0])
 639         fin = b1 >> 7 & 1
 640         rsv1 = b1 >> 6 & 1
 641         rsv2 = b1 >> 5 & 1
 642         rsv3 = b1 >> 4 & 1
 643         opcode = b1 & 0xf
 644         b2 = ord(self._frame_header[1])
 645         has_mask = b2 >> 7 & 1
 646         # Frame length
 647         if self._frame_length is None:
 648             length_bits = b2 & 0x7f
 649             if length_bits == 0x7e:
 650                 length_data = self._recv_strict(2)
 651                 self._frame_length = struct.unpack("!H", length_data)[0]
 652             elif length_bits == 0x7f:
 653                 length_data = self._recv_strict(8)
 654                 self._frame_length = struct.unpack("!Q", length_data)[0]
 655             else:
 656                 self._frame_length = length_bits
 657         # Mask
 658         if self._frame_mask is None:
 659             self._frame_mask = self._recv_strict(4) if has_mask else ""
 660         # Payload
 661         payload = self._recv_strict(self._frame_length)
 662         if has_mask:
 663             payload = ABNF.mask(self._frame_mask, payload)
 664         # Reset for next frame
 665         self._frame_header = None
 666         self._frame_length = None
 667         self._frame_mask = None
 668         return ABNF(fin, rsv1, rsv2, rsv3, opcode, has_mask, payload)
 669
 670
 671     def send_close(self, status=STATUS_NORMAL, reason=""):
 672         """
 673         send close data to the server.
 674
 675         status: status code to send. see STATUS_XXX.
 676
 677         reason: the reason to close. This must be string.
 678         """
 679         if status < 0 or status >= ABNF.LENGTH_16:
 680             raise ValueError("code is invalid range")
 681         self.send(struct.pack('!H', status) + reason, ABNF.OPCODE_CLOSE)
 682
 683     def close(self, status=STATUS_NORMAL, reason=""):
 684         """
 685         Close Websocket object
 686
 687         status: status code to send. see STATUS_XXX.
 688
 689         reason: the reason to close. This must be string.
 690         """
 691         if self.connected:
 692             if status < 0 or status >= ABNF.LENGTH_16:
 693                 raise ValueError("code is invalid range")
 694
 695             try:
 696                 self.send(struct.pack('!H', status) + reason, ABNF.OPCODE_CLOSE)
 697                 timeout = self.sock.gettimeout()
 698                 self.sock.settimeout(3)
 699                 try:
 700                     frame = self.recv_frame()
 701                     if logger.isEnabledFor(logging.ERROR):
 702                         recv_status = struct.unpack("!H", frame.data)[0]
 703                         if recv_status != STATUS_NORMAL:
 704                             logger.error("close status: " + repr(recv_status))
 705                 except:
 706                     pass
 707                 self.sock.settimeout(timeout)
 708                 self.sock.shutdown(socket.SHUT_RDWR)
 709             except:
 710                 pass
 711         self._closeInternal()
 712
 713     def _closeInternal(self):
 714         self.connected = False
 715         self.sock.close()
 716
 717     def _send(self, data):
 718         try:
 719             return self.sock.send(data)
 720         except socket.timeout as e:
 721             raise WebSocketTimeoutException(e.message)
 722         except Exception as e:
 723             if "timed out" in e.message:
 724                 raise WebSocketTimeoutException(e.message)
 725             else:
 726                 raise e
 727
 728     def _recv(self, bufsize):
 729         try:
 730             bytes = self.sock.recv(bufsize)
 731         except socket.timeout as e:
 732             raise WebSocketTimeoutException(e.message)
 733         except SSLError as e:
 734             if e.message == "The read operation timed out":
 735                 raise WebSocketTimeoutException(e.message)
 736             else:
 737                 raise
 738         if not bytes:
 739             raise WebSocketConnectionClosedException()
 740         return bytes
 741
 742
 743     def _recv_strict(self, bufsize):
 744         shortage = bufsize - sum(len(x) for x in self._recv_buffer)
 745         while shortage > 0:
 746             bytes = self._recv(shortage)
 747             self._recv_buffer.append(bytes)
 748             shortage -= len(bytes)
 749         unified = "".join(self._recv_buffer)
 750         if shortage == 0:
 751             self._recv_buffer = []
 752             return unified
 753         else:
 754             self._recv_buffer = [unified[bufsize:]]
 755             return unified[:bufsize]
 756
 757
 758     def _recv_line(self):
 759         line = []
 760         while True:
 761             c = self._recv(1)
 762             line.append(c)
 763             if c == "\n":
 764                 break
 765         return "".join(line)
 766
 767
 768 class WebSocketApp(object):
 769     """
 770     Higher level of APIs are provided.
 771     The interface is like JavaScript WebSocket object.
 772     """
 773     def __init__(self, url, header=[],
 774                  on_open=None, on_message=None, on_error=None,
 775                  on_close=None, keep_running=True, get_mask_key=None):
 776         """
 777         url: websocket url.
 778         header: custom header for websocket handshake.
 779         on_open: callable object which is called at opening websocket.
 780           this function has one argument. The arugment is this class object.
 781         on_message: callbale object which is called when recieved data.
 782          on_message has 2 arguments.
 783          The 1st arugment is this class object.
 784          The passing 2nd arugment is utf-8 string which we get from the server.
 785        on_error: callable object which is called when we get error.
 786          on_error has 2 arguments.
 787          The 1st arugment is this class object.
 788          The passing 2nd arugment is exception object.
 789        on_close: callable object which is called when closed the connection.
 790          this function has one argument. The arugment is this class object.
 791        keep_running: a boolean flag indicating whether the app's main loop should
 792          keep running, defaults to True
 793        get_mask_key: a callable to produce new mask keys, see the WebSocket.set_mask_key's
 794          docstring for more information
 795         """
 796         self.url = url
 797         self.header = header
 798         self.on_open = on_open
 799         self.on_message = on_message
 800         self.on_error = on_error
 801         self.on_close = on_close
 802         self.keep_running = keep_running
 803         self.get_mask_key = get_mask_key
 804         self.sock = None
 805
 806     def send(self, data, opcode=ABNF.OPCODE_TEXT):
 807         """
 808         send message.
 809         data: message to send. If you set opcode to OPCODE_TEXT, data must be utf-8 string or unicode.
 810         opcode: operation code of data. default is OPCODE_TEXT.
 811         """
 812         if self.sock.send(data, opcode) == 0:
 813             raise WebSocketConnectionClosedException()
 814
 815     def close(self):
 816         """
 817         close websocket connection.
 818         """
 819         self.keep_running = False
 820         self.sock.close()
 821
 822     def _send_ping(self, interval):
 823         while True:
 824             for i in range(interval):
 825                 time.sleep(1)
 826                 if not self.keep_running:
 827                     return
 828             self.sock.ping()
 829
 830     def run_forever(self, sockopt=None, sslopt=None, ping_interval=0):
 831         """
 832         run event loop for WebSocket framework.
 833         This loop is infinite loop and is alive during websocket is available.
 834         sockopt: values for socket.setsockopt.
 835             sockopt must be tuple and each element is argument of sock.setscokopt.
 836         sslopt: ssl socket optional dict.
 837         ping_interval: automatically send "ping" command every specified period(second)
 838             if set to 0, not send automatically.
 839         """
 840         if sockopt is None:
 841             sockopt = []
 842         if sslopt is None:
 843             sslopt = {}
 844         if self.sock:
 845             raise WebSocketException("socket is already opened")
 846         thread = None
 847
 848         try:
 849             self.sock = WebSocket(self.get_mask_key, sockopt=sockopt, sslopt=sslopt)
 850             self.sock.settimeout(default_timeout)
 851             self.sock.connect(self.url, header=self.header)
 852             self._callback(self.on_open)
 853
 854             if ping_interval:
 855                 thread = threading.Thread(target=self._send_ping, args=(ping_interval,))
 856                 thread.setDaemon(True)
 857                 thread.start()
 858
 859             while self.keep_running:
 860                 data = self.sock.recv()
 861                 if data is None:
 862                     break
 863                 self._callback(self.on_message, data)
 864         except Exception, e:
 865             self._callback(self.on_error, e)
 866         finally:
 867             if thread:
 868                 self.keep_running = False
 869             self.sock.close()
 870             self._callback(self.on_close)
 871             self.sock = None
 872
 873     def _callback(self, callback, *args):
 874         if callback:
 875             try:
 876                 callback(self, *args)
 877             except Exception, e:
 878                 logger.error(e)
 879                 if logger.isEnabledFor(logging.DEBUG):
 880                     _, _, tb = sys.exc_info()
 881                     traceback.print_tb(tb)
 882
 883
 884 if __name__ == "__main__":
 885     enableTrace(True)
 886     ws = create_connection("ws://echo.websocket.org/")
 887     print("Sending 'Hello, World'...")
 888     ws.send("Hello, World")
 889     print("Sent")
 890     print("Receiving...")
 891     result = ws.recv()
 892     print("Received '%s'" % result)
 893     ws.close()