def __init__(self, handshake_reply, protocols=None):
self.stream = Stream(always_mask=False)
self.handshake_reply = handshake_reply
self.handshake_sent = False
self.protocols = protocols
self.client_terminated = False
self.server_terminated = False
self.reading_buffer_size = DEFAULT_READING_SIZE
def __init__(self, url, protocols=None, version='8'):
self.stream = Stream()
self.url = url
self.protocols = protocols
self.version = version
self.key = b64encode(os.urandom(16))
self.client_terminated = False
self.server_terminated = False
def __init__(self, sock, protocols=None, extensions=None, environ=None):
""" The ``sock`` is an opened connection
resulting from the websocket handshake.
If ``protocols`` is provided, it is a list of protocols
negotiated during the handshake as is ``extensions``.
If ``environ`` is provided, it is a copy of the WSGI environ
dictionnary from the underlying WSGI server.
"""
self.stream = Stream(always_mask=False)
"""
Underlying websocket stream that performs the websocket
parsing to high level objects. By default this stream
never masks its messages. Clients using this class should
set the ``stream.always_mask`` fields to ``True``
and ``stream.expect_masking`` fields to ``False``.
"""
self.protocols = protocols
"""
List of protocols supported by this endpoint.
Unused for now.
"""
self.extensions = extensions
"""
List of extensions supported by this endpoint.
Unused for now.
"""
self.sock = sock
"""
Underlying connection.
"""
self.client_terminated = False
"""
Indicates if the client has been marked as terminated.
"""
self.server_terminated = False
"""
Indicates if the server has been marked as terminated.
"""
self.reading_buffer_size = DEFAULT_READING_SIZE
"""
Current connection reading buffer size.
"""
self.sender = self.sock.sendall
self.environ = environ
"""
def test_helper_with_binary_message(self):
msg = os.urandom(16)
s = Stream()
m = s.binary_message(msg)
self.assertIsInstance(m, BinaryMessage)
self.assertTrue(m.is_binary)
self.assertFalse(m.is_text)
self.assertEqual(m.opcode, OPCODE_BINARY)
self.assertIsInstance(m.data, bytes)
self.assertEqual(len(m), 16)
self.assertEqual(len(m.data), 16)
self.assertEqual(m.data, msg)
def test_helper_with_bytes_text_message(self):
s = Stream()
m = s.text_message('hello there!')
self.assertIsInstance(m, TextMessage)
self.assertFalse(m.is_binary)
self.assertTrue(m.is_text)
self.assertEqual(m.opcode, OPCODE_TEXT)
self.assertEqual(m.encoding, 'utf-8')
self.assertIsInstance(m.data, bytes)
self.assertEqual(len(m), 12)
self.assertEqual(len(m.data), 12)
self.assertEqual(m.data, b'hello there!')
def __init__(self, sock, environ, protocols=None, extensions=None):
self.stream = Stream()
self.protocols = protocols
self.extensions = extensions
self.environ = environ
self.sock = sock
self.sock.settimeout(30.0)
self.client_terminated = False
self.server_terminated = False
self._lock = Semaphore()
def __init__(self, sock, handshake_reply, protocols=None):
self.stream = Stream(always_mask=False)
self.handshake_reply = handshake_reply
self.handshake_sent = False
self.protocols = protocols
self.sock = sock
self.client_terminated = False
self.server_terminated = False
self.reading_buffer_size = DEFAULT_READING_SIZE
self.sender = self.sock.sendall
# This was initially a loop that used callbacks in ws4py
# Here it was turned into a generator, the callback replaced by yield
self.runner = self._run()
def test_text_message_received(self):
msg = b'hello there'
f = Frame(opcode=OPCODE_TEXT, body=msg, fin=1, masking_key=os.urandom(4)).build()
s = Stream()
self.assertEqual(len(s.messages), 0)
s.parser.send(f)
self.assertEqual(len(s.messages), 1)
def test_binary_message_with_continuation_received(self):
msg = os.urandom(16)
key = os.urandom(4)
f = Frame(opcode=OPCODE_BINARY, body=msg, fin=0,
masking_key=key).build()
s = Stream()
self.assertEqual(s.has_message, False)
s.parser.send(enc(f))
self.assertEqual(s.has_message, False)
for i in range(3):
f = Frame(opcode=OPCODE_CONTINUATION,
body=msg,
fin=0,
masking_key=key).build()
s.parser.send(enc(f))
self.assertEqual(s.has_message, False)
self.assertEqual(s.message.completed, False)
self.assertEqual(s.message.opcode, OPCODE_BINARY)
f = Frame(opcode=OPCODE_CONTINUATION, body=msg, fin=1,
masking_key=key).build()
s.parser.send(enc(f))
self.assertEqual(s.has_message, True)
self.assertEqual(s.message.completed, True)
self.assertEqual(s.message.opcode, OPCODE_BINARY)
def test_text_message_with_continuation_received(self):
msg = 'hello there'
f = Frame(opcode=OPCODE_TEXT,
body=msg,
fin=0,
masking_key=os.urandom(4)).build()
s = Stream()
self.assertEqual(s.has_message, False)
s.parser.send(enc(f))
self.assertEqual(s.message.completed, False)
for i in range(3):
f = Frame(opcode=OPCODE_CONTINUATION,
body=msg,
fin=0,
masking_key=os.urandom(4)).build()
s.parser.send(enc(f))
self.assertEqual(s.has_message, False)
self.assertEqual(s.message.completed, False)
self.assertEqual(s.message.opcode, OPCODE_TEXT)
f = Frame(opcode=OPCODE_CONTINUATION,
body=msg,
fin=1,
masking_key=os.urandom(4)).build()
s.parser.send(enc(f))
self.assertEqual(s.has_message, True)
self.assertEqual(s.message.completed, True)
self.assertEqual(s.message.opcode, OPCODE_TEXT)
def test_ping_message_received(self):
msg = 'ping me'
f = Frame(opcode=OPCODE_PING, body=msg, fin=1).build()
s = Stream()
self.assertEqual(len(s.pings), 0)
s.parser.send(f)
self.assertEqual(len(s.pings), 1)
def test_pong_message_received(self):
msg = b'pong!'
f = Frame(opcode=OPCODE_PONG, body=msg, fin=1, masking_key=os.urandom(4)).build()
s = Stream()
self.assertEqual(len(s.pongs), 0)
s.parser.send(f)
self.assertEqual(len(s.pongs), 1)
def test_text_message_with_continuation_and_ping_in_between(self):
msg = b'hello there'
key = os.urandom(4)
f = Frame(opcode=OPCODE_TEXT, body=msg, fin=0, masking_key=os.urandom(4)).build()
s = Stream()
self.assertEqual(s.has_message, False)
s.parser.send(f)
self.assertEqual(s.message.completed, False)
for i in range(3):
f = Frame(opcode=OPCODE_CONTINUATION, body=msg, fin=0, masking_key=os.urandom(4)).build()
s.parser.send(f)
self.assertEqual(s.has_message, False)
self.assertEqual(s.message.completed, False)
self.assertEqual(s.message.opcode, OPCODE_TEXT)
f = Frame(opcode=OPCODE_PING, body=b'ping me', fin=1, masking_key=os.urandom(4)).build()
self.assertEqual(len(s.pings), i)
s.parser.send(f)
self.assertEqual(len(s.pings), i+1)
f = Frame(opcode=OPCODE_CONTINUATION, body=msg, fin=1, masking_key=os.urandom(4)).build()
s.parser.send(f)
self.assertEqual(s.has_message, True)
self.assertEqual(s.message.opcode, OPCODE_TEXT)
self.assertEqual(s.message.completed, True)
def test_text_message_received(self):
msg = b'hello there'
f = Frame(opcode=OPCODE_TEXT, body=msg, fin=1, masking_key=os.urandom(4)).build()
s = Stream()
self.assertEqual(s.has_message, False)
s.parser.send(f)
self.assertEqual(s.message.completed, True)
def test_binary_message_received(self):
msg = os.urandom(16)
f = Frame(opcode=OPCODE_BINARY, body=msg, fin=1, masking_key=os.urandom(4)).build()
s = Stream()
self.assertEqual(s.has_message, False)
s.parser.send(f)
self.assertEqual(s.message.completed, True)
def test_closing_parser_should_release_resources(self):
f = Frame(opcode=OPCODE_TEXT,
body=b'hello',
fin=1,
masking_key=os.urandom(4)).build()
s = Stream()
s.parser.send(f)
s.parser.close()
def __init__(self, url, protocols=None, version='8'):
self.stream = Stream()
self.url = url
self.protocols = protocols
self.version = version
self.key = b64encode(os.urandom(16))
self.client_terminated = False
self.server_terminated = False
def test_using_masking_key_when_unexpected(self):
f = Frame(opcode=OPCODE_TEXT, body=b'hello', fin=1, masking_key=os.urandom(4)).build()
s = Stream(expect_masking=False)
s.parser.send(f)
next(s.parser)
self.assertNotEqual(s.errors, [])
self.assertIsInstance(s.errors[0], CloseControlMessage)
self.assertEqual(s.errors[0].code, 1002)
def __init__(self, handshake_reply, protocols=None):
self.stream = Stream(always_mask=False)
self.handshake_reply = handshake_reply
self.handshake_sent = False
self.protocols = protocols
self.client_terminated = False
self.server_terminated = False
self.reading_buffer_size = DEFAULT_READING_SIZE
def test_incremental_text_message_received(self):
msg = 'hello there'
f = Frame(opcode=OPCODE_TEXT, body=msg, fin=1).build()
s = Stream()
self.assertEqual(s.has_message, False)
for byte in f:
s.parser.send(byte)
self.assertEqual(s.has_message, True)
def test_invalid_encoded_bytes(self):
f = Frame(opcode=OPCODE_TEXT, body=b'h\xc3llo',
fin=1, masking_key=os.urandom(4)).build()
s = Stream()
s.parser.send(f)
next(s.parser)
self.assertNotEqual(s.errors, [])
self.assertIsInstance(s.errors[0], CloseControlMessage)
self.assertEqual(s.errors[0].code, 1007)
def test_incremental_text_message_received(self):
msg = b'hello there'
f = Frame(opcode=OPCODE_TEXT, body=msg, fin=1, masking_key=os.urandom(4)).build()
s = Stream()
self.assertEqual(s.has_message, False)
bytes = f
for index, byte in enumerate(bytes):
s.parser.send(bytes[index:index+1])
self.assertEqual(s.has_message, True)
def test_empty_close_message(self):
f = Frame(opcode=OPCODE_CLOSE,
body=enc(''),
fin=1,
masking_key=os.urandom(4)).build()
s = Stream()
self.assertEqual(s.closing, None)
s.parser.send(enc(f))
self.assertEqual(type(s.closing), CloseControlMessage)
def test_continuation_frame_before_message_started_is_invalid(self):
f = Frame(opcode=OPCODE_CONTINUATION, body=b'hello',
fin=1, masking_key=os.urandom(4)).build()
s = Stream()
s.parser.send(f)
next(s.parser)
self.assertNotEqual(s.errors, [])
self.assertIsInstance(s.errors[0], CloseControlMessage)
self.assertEqual(s.errors[0].code, 1002)
def __init__(self, sock, protocols=None, extensions=None, environ=None):
""" The ``sock`` is an opened connection
resulting from the websocket handshake.
If ``protocols`` is provided, it is a list of protocols
negotiated during the handshake as is ``extensions``.
If ``environ`` is provided, it is a copy of the WSGI environ
dictionnary from the underlying WSGI server.
"""
self.stream = Stream(always_mask=False)
"""
Underlying websocket stream that performs the websocket
parsing to high level objects. By default this stream
never masks its messages. Clients using this class should
set the ``stream.always_mask`` fields to ``True``
and ``stream.expect_masking`` fields to ``False``.
"""
self.protocols = protocols
"""
List of protocols supported by this endpoint.
Unused for now.
"""
self.extensions = extensions
"""
List of extensions supported by this endpoint.
Unused for now.
"""
self.sock = sock
"""
Underlying connection.
"""
self.client_terminated = False
"""
Indicates if the client has been marked as terminated.
"""
self.server_terminated = False
"""
Indicates if the server has been marked as terminated.
"""
self.reading_buffer_size = DEFAULT_READING_SIZE
"""
Current connection reading buffer size.
"""
self.sender = self.sock.sendall
self.environ = environ
"""
def test_close_message_of_size_one_are_invalid(self):
payload = b'*'
f = Frame(opcode=OPCODE_CLOSE, body=payload,
fin=1, masking_key=os.urandom(4)).build()
s = Stream()
self.assertEqual(len(s.errors), 0)
self.assertEqual(s.closing, None)
s.parser.send(f)
self.assertEqual(type(s.closing), CloseControlMessage)
self.assertEqual(s.closing.code, 1002)
def test_close_message_received(self):
payload = struct.pack("!H", 1000) + b'hello'
f = Frame(opcode=OPCODE_CLOSE, body=payload,
fin=1, masking_key=os.urandom(4)).build()
s = Stream()
self.assertEqual(s.closing, None)
s.parser.send(f)
self.assertEqual(type(s.closing), CloseControlMessage)
self.assertEqual(s.closing.code, 1000)
self.assertEqual(s.closing.reason, b'hello')
def test_ping_message_received(self):
msg = enc('ping me')
f = Frame(opcode=OPCODE_PING,
body=msg,
fin=1,
masking_key=os.urandom(4)).build()
s = Stream()
self.assertEqual(len(s.pings), 0)
s.parser.send(enc(f))
self.assertEqual(len(s.pings), 1)
def test_invalid_close_message_reason_encoding(self):
payload = struct.pack("!H", 1000) + b'h\xc3llo'
f = Frame(opcode=OPCODE_CLOSE, body=payload,
fin=1, masking_key=os.urandom(4)).build()
s = Stream()
self.assertEqual(len(s.errors), 0)
self.assertEqual(s.closing, None)
s.parser.send(f)
self.assertEqual(s.closing, None)
self.assertEqual(type(s.errors[0]), CloseControlMessage)
self.assertEqual(s.errors[0].code, 1007)
def test_protocol_exception_from_frame_parsing(self):
payload = struct.pack("!H", 1000) + b'hello'
f = Frame(opcode=OPCODE_CLOSE, body=payload,
fin=1, masking_key=os.urandom(4))
f.rsv1 = 1
f = f.build()
s = Stream()
self.assertEqual(len(s.errors), 0)
self.assertEqual(s.closing, None)
s.parser.send(f)
self.assertEqual(s.closing, None)
self.assertEqual(type(s.errors[0]), CloseControlMessage)
self.assertEqual(s.errors[0].code, 1002)
def test_too_large_close_message(self):
payload = struct.pack("!H", 1000) + b'*' * 330
f = Frame(opcode=OPCODE_CLOSE, body=payload,
fin=1, masking_key=os.urandom(4)).build()
s = Stream()
self.assertEqual(len(s.errors), 0)
self.assertEqual(s.closing, None)
s.parser.send(f)
self.assertEqual(s.closing, None)
self.assertEqual(len(s.errors), 1)
self.assertEqual(type(s.errors[0]), CloseControlMessage)
self.assertEqual(s.errors[0].code, 1002)
def __init__(self, sock, protocols, extensions):
"""
A handler appropriate for servers. This handler
runs the connection's read and parsing in a thread,
meaning that incoming messages will be alerted in a different
thread from the one that created the handler.
@param sock: opened connection after the websocket upgrade
@param protocols: list of protocols from the handshake
@param extensions: list of extensions from the handshake
"""
self.stream = Stream()
self.protocols = protocols
self.extensions = extensions
self.sock = sock
self.sock.settimeout(30.0)
self.client_terminated = False
self.server_terminated = False
self._th = threading.Thread(target=self._receive)
def __init__(self, sock, handshake_reply, protocols=None):
self.stream = Stream(always_mask=False)
self.handshake_reply = handshake_reply
self.handshake_sent = False
self.protocols = protocols
self.sock = sock
self.client_terminated = False
self.server_terminated = False
self.reading_buffer_size = DEFAULT_READING_SIZE
self.sender = self.sock.sendall
# This was initially a loop that used callbacks in ws4py
# Here it was turned into a generator, the callback replaced by yield
self.runner = self._run()
def __init__(self, sock, environ, protocols=None, extensions=None):
self.stream = Stream()
self.protocols = protocols
self.extensions = extensions
self.environ = environ
self.sock = sock
self.sock.settimeout(30.0)
self.client_terminated = False
self.server_terminated = False
self._lock = Semaphore()
def test_binary_and_text_messages_cannot_interleave(self):
s = Stream()
f = Frame(opcode=OPCODE_TEXT, body=b'hello',
fin=0, masking_key=os.urandom(4)).build()
s.parser.send(f)
next(s.parser)
f = Frame(opcode=OPCODE_BINARY, body=os.urandom(7),
fin=1, masking_key=os.urandom(4)).build()
s.parser.send(f)
next(s.parser)
self.assertNotEqual(s.errors, [])
self.assertIsInstance(s.errors[0], CloseControlMessage)
self.assertEqual(s.errors[0].code, 1002)
def __init__(self, sock, protocols, extensions):
"""
A handler appropriate for servers. This handler
runs the connection's read and parsing in a thread,
meaning that incoming messages will be alerted in a different
thread from the one that created the handler.
@param sock: opened connection after the websocket upgrade
@param protocols: list of protocols from the handshake
@param extensions: list of extensions from the handshake
"""
self.stream = Stream()
self.protocols = protocols
self.extensions = extensions
self.sock = sock
self.sock.settimeout(30.0)
self.client_terminated = False
self.server_terminated = False
self._th = threading.Thread(target=self._receive)
class WebSocket(object):
def __init__(self, sock, environ, protocols=None, extensions=None):
self.stream = Stream()
self.protocols = protocols
self.extensions = extensions
self.environ = environ
self.sock = sock
self.sock.settimeout(30.0)
self.client_terminated = False
self.server_terminated = False
self._lock = Semaphore()
def close(self, code=1000, reason=''):
"""
Call this method to initiate the websocket connection
closing by sending a close frame to the connected peer.
Once this method is called, the server_terminated
attribute is set. Calling this method several times is
safe as the closing frame will be sent only the first
time.
@param code: status code describing why the connection is closed
@param reason: a human readable message describing why the connection is closed
"""
if not self.server_terminated:
self.server_terminated = True
self.write_to_connection(self.stream.close(code=code, reason=reason))
self.close_connection()
@property
def terminated(self):
"""
Returns True if both the client and server have been
marked as terminated.
"""
return self.client_terminated is True and self.server_terminated is True
def write_to_connection(self, bytes):
"""
Writes the provided bytes to the underlying connection.
@param bytes: data tio send out
"""
return self.sock.sendall(bytes)
def read_from_connection(self, amount):
"""
Reads bytes from the underlying connection.
@param amount: quantity to read (if possible)
"""
return self.sock.recv(amount)
def close_connection(self):
"""
Shutdowns then closes the underlying connection.
"""
try:
self.sock.shutdown(socket.SHUT_RDWR)
finally:
self.sock.close()
def send(self, payload, binary=False):
"""
Sends the given payload out.
If payload is some bytes or a bytearray,
then it is sent as a single message not fragmented.
If payload is a generator, each chunk is sent as part of
fragmented message.
@param payload: string, bytes, bytearray or a generator
@param binary: if set, handles the payload as a binary message
"""
if isinstance(payload, basestring) or isinstance(payload, bytearray):
if not binary:
self.write_to_connection(self.stream.text_message(payload).single())
else:
self.write_to_connection(self.stream.binary_message(payload).single())
elif type(payload) == types.GeneratorType:
bytes = payload.next()
first = True
for chunk in payload:
if not binary:
self.write_to_connection(self.stream.text_message(bytes).fragment(first=first))
else:
self.write_to_connection(self.stream.binary_message(payload).fragment(first=first))
bytes = chunk
first = False
if not binary:
self.write_to_connection(self.stream.text_message(bytes).fragment(last=True))
else:
self.write_to_connection(self.stream.text_message(bytes).fragment(last=True))
def receive(self, message_obj=False):
"""
Performs the operation of reading from the underlying
connection in order to feed the stream of bytes.
We start with a small size of two bytes to be read
from the connection so that we can quickly parse an
incoming frame header. Then the stream indicates
whatever size must be read from the connection since
it knows the frame payload length.
Note that we perform some automatic opererations:
* On a closing message, we respond with a closing
message and finally close the connection
* We respond to pings with pong messages.
* Whenever an error is raised by the stream parsing,
we initiate the closing of the connection with the
appropriate error code.
"""
next_size = 2
#try:
while not self.terminated:
bytes = self.read_from_connection(next_size)
if not bytes and next_size > 0:
raise IOError()
message = None
with self._lock:
s = self.stream
next_size = s.parser.send(bytes)
if s.closing is not None:
if not self.server_terminated:
next_size = 2
self.close(s.closing.code, s.closing.reason)
else:
self.client_terminated = True
raise IOError()
elif s.errors:
errors = s.errors[:]
for error in s.errors:
self.close(error.code, error.reason)
s.errors.remove(error)
raise IOError()
elif s.has_message:
if message_obj:
message = s.message
s.message = None
else:
message = str(s.message)
s.message.data = None
s.message = None
for ping in s.pings:
self.write_to_connection(s.pong(str(ping.data)))
s.pings = []
s.pongs = []
if message is not None:
return message
class WebSocketHandler(object):
def __init__(self, sock, protocols, extensions):
"""
A handler appropriate for servers. This handler
runs the connection's read and parsing in a thread,
meaning that incoming messages will be alerted in a different
thread from the one that created the handler.
@param sock: opened connection after the websocket upgrade
@param protocols: list of protocols from the handshake
@param extensions: list of extensions from the handshake
"""
self.stream = Stream()
self.protocols = protocols
self.extensions = extensions
self.sock = sock
self.sock.settimeout(30.0)
self.client_terminated = False
self.server_terminated = False
self._th = threading.Thread(target=self._receive)
def opened(self):
"""
Called by the server when the upgrade handshake
has succeeeded. Starts the internal loop that
reads bytes from the connection and hands it over
to the stream for parsing.
"""
self._th.start()
def close(self, code=1000, reason=''):
"""
Call this method to initiate the websocket connection
closing by sending a close frame to the connected peer.
Once this method is called, the server_terminated
attribute is set. Calling this method several times is
safe as the closing frame will be sent only the first
time.
@param code: status code describing why the connection is closed
@param reason: a human readable message describing why the connection is closed
"""
if not self.server_terminated:
self.server_terminated = True
self.write_to_connection(self.stream.close(code=code, reason=reason))
def closed(self, code, reason=None):
"""
Called by the server when the websocket connection
is finally closed.
@param code: status code
@param reason: human readable message of the closing exchange
"""
pass
@property
def terminated(self):
"""
Returns True if both the client and server have been
marked as terminated.
"""
return self.client_terminated is True and self.server_terminated is True
def write_to_connection(self, bytes):
"""
Writes the provided bytes to the underlying connection.
@param bytes: data tio send out
"""
return self.sock.sendall(bytes)
def read_from_connection(self, amount):
"""
Reads bytes from the underlying connection.
@param amount: quantity to read (if possible)
"""
return self.sock.recv(amount)
def close_connection(self):
"""
Shutdowns then closes the underlying connection.
"""
try:
self.sock.shutdown(socket.SHUT_RDWR)
self.sock.close()
except:
pass
def ponged(self, pong):
"""
Pong message received on the stream.
@param pong: messaging.PongControlMessage instance
"""
pass
def received_message(self, m):
"""
Message received on the stream.
@param pong: messaging.TextMessage or messaging.BinaryMessage instance
"""
pass
def send(self, payload, binary=False):
"""
Sends the given payload out.
If payload is some bytes or a bytearray,
then it is sent as a single message not fragmented.
If payload is a generator, each chunk is sent as part of
fragmented message.
@param payload: string, bytes, bytearray or a generator
@param binary: if set, handles the payload as a binary message
"""
if isinstance(payload, basestring) or isinstance(payload, bytearray):
if not binary:
self.write_to_connection(self.stream.text_message(payload).single())
else:
self.write_to_connection(self.stream.binary_message(payload).single())
elif type(payload) == types.GeneratorType:
bytes = payload.next()
first = True
for chunk in payload:
if not binary:
self.write_to_connection(self.stream.text_message(bytes).fragment(first=first))
else:
self.write_to_connection(self.stream.binary_message(payload).fragment(first=first))
bytes = chunk
first = False
if not binary:
self.write_to_connection(self.stream.text_message(bytes).fragment(last=True))
else:
self.write_to_connection(self.stream.text_message(bytes).fragment(last=True))
def _receive(self):
"""
Performs the operation of reading from the underlying
connection in order to feed the stream of bytes.
We start with a small size of two bytes to be read
from the connection so that we can quickly parse an
incoming frame header. Then the stream indicates
whatever size must be read from the connection since
it knows the frame payload length.
Note that we perform some automatic opererations:
* On a closing message, we respond with a closing
message and finally close the connection
* We respond to pings with pong messages.
* Whenever an error is raised by the stream parsing,
we initiate the closing of the connection with the
appropiate error code.
"""
next_size = 2
try:
self.sock.setblocking(1)
while not self.terminated:
bytes = self.read_from_connection(next_size)
if not bytes and next_size > 0:
break
s = self.stream
next_size = s.parser.send(bytes)
if s.closing is not None:
if not self.server_terminated:
next_size = 2
self.close(s.closing.code, s.closing.reason)
else:
self.client_terminated = True
break
elif s.errors:
errors = s.errors[:]
for error in errors:
self.close(error.code, error.reason)
s.errors.remove(error)
break
elif s.has_message:
self.received_message(s.message)
s.message.data = None
s.message = None
for ping in s.pings:
self.write_to_connection(s.pong(str(ping.data)))
s.pings = []
for pong in s.pongs:
self.ponged(pong)
s.pongs = []
except:
print "".join(traceback.format_exception(*exc_info()))
finally:
self.client_terminated = self.server_terminated = True
self.close_connection()
if self.stream.closing:
self.closed(self.stream.closing.code, self.stream.closing.reason)
else:
self.closed(1006)
def __init__(self, sock, protocols=None, extensions=None, environ=None, heartbeat_freq=None):
""" The ``sock`` is an opened connection
resulting from the websocket handshake.
If ``protocols`` is provided, it is a list of protocols
negotiated during the handshake as is ``extensions``.
If ``environ`` is provided, it is a copy of the WSGI environ
dictionnary from the underlying WSGI server.
"""
self.stream = Stream(always_mask=False)
"""
Underlying websocket stream that performs the websocket
parsing to high level objects. By default this stream
never masks its messages. Clients using this class should
set the ``stream.always_mask`` fields to ``True``
and ``stream.expect_masking`` fields to ``False``.
"""
self.protocols = protocols
"""
List of protocols supported by this endpoint.
Unused for now.
"""
self.extensions = extensions
"""
List of extensions supported by this endpoint.
Unused for now.
"""
self.sock = sock
"""
Underlying connection.
"""
self._is_secure = hasattr(sock, '_ssl') or hasattr(sock, '_sslobj')
"""
Tell us if the socket is secure or not.
"""
self.client_terminated = False
"""
Indicates if the client has been marked as terminated.
"""
self.server_terminated = False
"""
Indicates if the server has been marked as terminated.
"""
self.reading_buffer_size = DEFAULT_READING_SIZE
"""
Current connection reading buffer size.
"""
self.environ = environ
"""
WSGI environ dictionary.
"""
self.heartbeat_freq = heartbeat_freq
"""
At which interval the heartbeat will be running.
Set this to `0` or `None` to disable it entirely.
"""
self._local_address = None
self._peer_address = None
def test_helper_masked_pong_message(self):
s = Stream(always_mask=True)
m = s.pong('sos')
self.assertIsInstance(m, bytes)
self.assertEqual(len(m), 9)
def test_helper_pong_message(self):
s = Stream()
m = s.pong('sos')
self.assertIsInstance(m, bytes)
self.assertEqual(len(m), 5)
class WebSocket(object):
""" Represents a websocket endpoint and provides a high level interface to drive the endpoint. """
def __init__(self, sock, protocols=None, extensions=None, environ=None, heartbeat_freq=None):
""" The ``sock`` is an opened connection
resulting from the websocket handshake.
If ``protocols`` is provided, it is a list of protocols
negotiated during the handshake as is ``extensions``.
If ``environ`` is provided, it is a copy of the WSGI environ
dictionnary from the underlying WSGI server.
"""
self.stream = Stream(always_mask=False)
"""
Underlying websocket stream that performs the websocket
parsing to high level objects. By default this stream
never masks its messages. Clients using this class should
set the ``stream.always_mask`` fields to ``True``
and ``stream.expect_masking`` fields to ``False``.
"""
self.protocols = protocols
"""
List of protocols supported by this endpoint.
Unused for now.
"""
self.extensions = extensions
"""
List of extensions supported by this endpoint.
Unused for now.
"""
self.sock = sock
"""
Underlying connection.
"""
self._is_secure = hasattr(sock, '_ssl') or hasattr(sock, '_sslobj')
"""
Tell us if the socket is secure or not.
"""
self.client_terminated = False
"""
Indicates if the client has been marked as terminated.
"""
self.server_terminated = False
"""
Indicates if the server has been marked as terminated.
"""
self.reading_buffer_size = DEFAULT_READING_SIZE
"""
Current connection reading buffer size.
"""
self.environ = environ
"""
WSGI environ dictionary.
"""
self.heartbeat_freq = heartbeat_freq
"""
At which interval the heartbeat will be running.
Set this to `0` or `None` to disable it entirely.
"""
self._local_address = None
self._peer_address = None
@property
def local_address(self):
"""
Local endpoint address as a tuple
"""
if not self._local_address:
self._local_address = self.sock.getsockname()
if len(self._local_address) == 4:
self._local_address = self._local_address[:2]
return self._local_address
@property
def peer_address(self):
"""
Peer endpoint address as a tuple
"""
if not self._peer_address:
self._peer_address = self.sock.getpeername()
if len(self._peer_address) == 4:
self._peer_address = self._peer_address[:2]
return self._peer_address
def opened(self):
"""
Called by the server when the upgrade handshake
has succeeeded.
"""
pass
def close(self, code=1000, reason=''):
"""
Call this method to initiate the websocket connection
closing by sending a close frame to the connected peer.
The ``code`` is the status code representing the
termination's reason.
Once this method is called, the ``server_terminated``
attribute is set. Calling this method several times is
safe as the closing frame will be sent only the first
time.
.. seealso:: Defined Status Codes http://tools.ietf.org/html/rfc6455#section-7.4.1
"""
if not self.server_terminated:
self.server_terminated = True
try:
self._write(self.stream.close(code=code, reason=reason).single(mask=self.stream.always_mask))
except Exception as ex:
logger.error("Error when terminating the connection: %s", str(ex))
def closed(self, code, reason=None):
"""
Called when the websocket stream and connection are finally closed.
The provided ``code`` is status set by the other point and
``reason`` is a human readable message.
.. seealso:: Defined Status Codes http://tools.ietf.org/html/rfc6455#section-7.4.1
"""
pass
@property
def terminated(self):
"""
Returns ``True`` if both the client and server have been
marked as terminated.
"""
return self.client_terminated is True and self.server_terminated is True
@property
def connection(self):
return self.sock
def close_connection(self):
"""
Shutdowns then closes the underlying connection.
"""
if self.sock:
try:
self.sock.shutdown(socket.SHUT_RDWR)
self.sock.close()
except:
pass
finally:
self.sock = None
def ping(self, message):
"""
Send a ping message to the remote peer.
The given `message` must be a unicode string.
"""
self.send(PingControlMessage(message))
def ponged(self, pong):
"""
Pong message, as a :class:`messaging.PongControlMessage` instance,
received on the stream.
"""
pass
def received_message(self, message):
"""
Called whenever a complete ``message``, binary or text,
is received and ready for application's processing.
The passed message is an instance of :class:`messaging.TextMessage`
or :class:`messaging.BinaryMessage`.
.. note:: You should override this method in your subclass.
"""
pass
def unhandled_error(self, error):
"""
Called whenever a socket, or an OS, error is trapped
by ws4py but not managed by it. The given error is
an instance of `socket.error` or `OSError`.
Note however that application exceptions will not go
through this handler. Instead, do make sure you
protect your code appropriately in `received_message`
or `send`.
The default behaviour of this handler is to log
the error with a message.
"""
logger.exception("Failed to receive data")
def _write(self, b):
"""
Trying to prevent a write operation
on an already closed websocket stream.
This cannot be bullet proof but hopefully
will catch almost all use cases.
"""
if self.terminated or self.sock is None:
raise RuntimeError("Cannot send on a terminated websocket")
self.sock.sendall(b)
def send(self, payload, binary=False):
"""
Sends the given ``payload`` out.
If ``payload`` is some bytes or a bytearray,
then it is sent as a single message not fragmented.
If ``payload`` is a generator, each chunk is sent as part of
fragmented message.
If ``binary`` is set, handles the payload as a binary message.
"""
message_sender = self.stream.binary_message if binary else self.stream.text_message
if isinstance(payload, basestring) or isinstance(payload, bytearray):
m = message_sender(payload).single(mask=self.stream.always_mask)
self._write(m)
elif isinstance(payload, Message):
data = payload.single(mask=self.stream.always_mask)
self._write(data)
elif type(payload) == types.GeneratorType:
bytes = next(payload)
first = True
for chunk in payload:
self._write(message_sender(bytes).fragment(first=first, mask=self.stream.always_mask))
bytes = chunk
first = False
self._write(message_sender(bytes).fragment(last=True, mask=self.stream.always_mask))
else:
raise ValueError("Unsupported type '%s' passed to send()" % type(payload))
def _get_from_pending(self):
"""
The SSL socket object provides the same interface
as the socket interface but behaves differently.
When data is sent over a SSL connection
more data may be read than was requested from by
the ws4py websocket object.
In that case, the data may have been indeed read
from the underlying real socket, but not read by the
application which will expect another trigger from the
manager's polling mechanism as if more data was still on the
wire. This will happen only when new data is
sent by the other peer which means there will be
some delay before the initial read data is handled
by the application.
Due to this, we have to rely on a non-public method
to query the internal SSL socket buffer if it has indeed
more data pending in its buffer.
Now, some people in the Python community
`discourage <https://bugs.python.org/issue21430>`_
this usage of the ``pending()`` method because it's not
the right way of dealing with such use case. They advise
`this approach <https://docs.python.org/dev/library/ssl.html#notes-on-non-blocking-sockets>`_
instead. Unfortunately, this applies only if the
application can directly control the poller which is not
the case with the WebSocket abstraction here.
We therefore rely on this `technic <http://stackoverflow.com/questions/3187565/select-and-ssl-in-python>`_
which seems to be valid anyway.
This is a bit of a shame because we have to process
more data than what wanted initially.
"""
data = b""
pending = self.sock.pending()
while pending:
data += self.sock.recv(pending)
pending = self.sock.pending()
return data
def once(self):
"""
Performs the operation of reading from the underlying
connection in order to feed the stream of bytes.
We start with a small size of two bytes to be read
from the connection so that we can quickly parse an
incoming frame header. Then the stream indicates
whatever size must be read from the connection since
it knows the frame payload length.
It returns `False` if an error occurred at the
socket level or during the bytes processing. Otherwise,
it returns `True`.
"""
if self.terminated:
logger.debug("WebSocket is already terminated")
return False
try:
b = self.sock.recv(self.reading_buffer_size)
# This will only make sense with secure sockets.
if self._is_secure:
b += self._get_from_pending()
except (socket.error, OSError, pyOpenSSLError) as e:
self.unhandled_error(e)
return False
else:
if not self.process(b):
return False
return True
def terminate(self):
"""
Completes the websocket by calling the `closed`
method either using the received closing code
and reason, or when none was received, using
the special `1006` code.
Finally close the underlying connection for
good and cleanup resources by unsetting
the `environ` and `stream` attributes.
"""
s = self.stream
try:
if s.closing is None:
self.closed(1006, "Going away")
else:
self.closed(s.closing.code, s.closing.reason)
finally:
self.client_terminated = self.server_terminated = True
self.close_connection()
# Cleaning up resources
s._cleanup()
self.stream = None
self.environ = None
def process(self, bytes):
""" Takes some bytes and process them through the
internal stream's parser. If a message of any kind is
found, performs one of these actions:
* A closing message will initiate the closing handshake
* Errors will initiate a closing handshake
* A message will be passed to the ``received_message`` method
* Pings will see pongs be sent automatically
* Pongs will be passed to the ``ponged`` method
The process should be terminated when this method
returns ``False``.
"""
s = self.stream
if not bytes and self.reading_buffer_size > 0:
return False
self.reading_buffer_size = s.parser.send(bytes) or DEFAULT_READING_SIZE
if s.closing is not None:
logger.debug("Closing message received (%d) '%s'" % (s.closing.code, s.closing.reason))
if not self.server_terminated:
self.close(s.closing.code, s.closing.reason)
else:
self.client_terminated = True
return False
if s.errors:
for error in s.errors:
logger.debug("Error message received (%d) '%s'" % (error.code, error.reason))
self.close(error.code, error.reason)
s.errors = []
return False
if s.has_message:
self.received_message(s.message)
if s.message is not None:
s.message.data = None
s.message = None
return True
if s.pings:
for ping in s.pings:
self._write(s.pong(ping.data))
s.pings = []
if s.pongs:
for pong in s.pongs:
self.ponged(pong)
s.pongs = []
return True
def run(self):
"""
Performs the operation of reading from the underlying
connection in order to feed the stream of bytes.
We start with a small size of two bytes to be read
from the connection so that we can quickly parse an
incoming frame header. Then the stream indicates
whatever size must be read from the connection since
it knows the frame payload length.
Note that we perform some automatic opererations:
* On a closing message, we respond with a closing
message and finally close the connection
* We respond to pings with pong messages.
* Whenever an error is raised by the stream parsing,
we initiate the closing of the connection with the
appropiate error code.
This method is blocking and should likely be run
in a thread.
"""
self.sock.setblocking(True)
with Heartbeat(self, frequency=self.heartbeat_freq):
s = self.stream
try:
self.opened()
while not self.terminated:
if not self.once():
break
finally:
self.terminate()
class WebSocket(object):
def __init__(self, handshake_reply, protocols=None):
self.stream = Stream(always_mask=False)
self.handshake_reply = handshake_reply
self.handshake_sent = False
self.protocols = protocols
self.client_terminated = False
self.server_terminated = False
self.reading_buffer_size = DEFAULT_READING_SIZE
def init(self, sender):
# This was initially a loop that used callbacks in ws4py
# Here it was turned into a generator, the callback replaced by yield
self.sender = sender
self.sender(self.handshake_reply)
self.handshake_sent = True
def send(self, payload, binary=False):
"""
Sends the given ``payload`` out.
If ``payload`` is some bytes or a bytearray,
then it is sent as a single message not fragmented.
If ``payload`` is a generator, each chunk is sent as part of
fragmented message.
If ``binary`` is set, handles the payload as a binary message.
"""
message_sender = self.stream.binary_message if binary else self.stream.text_message
if isinstance(payload, basestring) or isinstance(payload, bytearray):
self.sender(message_sender(payload).single(mask=self.stream.always_mask))
elif isinstance(payload, Message):
self.sender(payload.single(mask=self.stream.always_mask))
elif type(payload) == types.GeneratorType:
bytes = payload.next()
first = True
for chunk in payload:
self.sender(message_sender(bytes).fragment(first=first, mask=self.stream.always_mask))
bytes = chunk
first = False
self.sender(message_sender(bytes).fragment(last=True, mask=self.stream.always_mask))
else:
raise ValueError("Unsupported type '%s' passed to send()" % type(payload))
def dataReceived(self, data):
"""
Performs the operation of reading from the underlying
connection in order to feed the stream of bytes.
We start with a small size of two bytes to be read
from the connection so that we can quickly parse an
incoming frame header. Then the stream indicates
whatever size must be read from the connection since
it knows the frame payload length.
Note that we perform some automatic operations:
* On a closing message, we respond with a closing
message and finally close the connection
* We respond to pings with pong messages.
* Whenever an error is raised by the stream parsing,
we initiate the closing of the connection with the
appropiate error code.
"""
s = self.stream
self.reading_buffer_size = s.parser.send(data) or DEFAULT_READING_SIZE
if s.closing is not None:
if not self.server_terminated:
self.close(s.closing.code, s.closing.reason)
else:
self.client_terminated = True
return None
if s.errors:
for error in s.errors:
self.close(error.code, error.reason)
s.errors = []
return None
if s.has_message:
msg = s.message
return msg
s.message = None
else:
if s.pings:
for ping in s.pings:
self.sender(s.pong(ping.data))
s.pings = []
if s.pongs:
s.pongs = []
return None
def close(self, code=1000, reason=''):
"""
Call this method to initiate the websocket connection
closing by sending a close frame to the connected peer.
The ``code`` is the status code representing the
termination's reason.
Once this method is called, the ``server_terminated``
attribute is set. Calling this method several times is
safe as the closing frame will be sent only the first
time.
.. seealso:: Defined Status Codes http://tools.ietf.org/html/rfc6455#section-7.4.1
"""
if not self.server_terminated:
self.server_terminated = True
self.sender(self.stream.close(code=code, reason=reason).single(mask=self.stream.always_mask))
@property
def terminated(self):
"""
Returns ``True`` if both the client and server have been
marked as terminated.
"""
return self.client_terminated is True and self.server_terminated is True
def __iter__(self):
return self.runner
class WebSocket(object):
""" Represents a websocket endpoint and provides a high level interface to drive the endpoint. """
def __init__(self, sock, protocols=None, extensions=None, environ=None, heartbeat_freq=None):
""" The ``sock`` is an opened connection
resulting from the websocket handshake.
If ``protocols`` is provided, it is a list of protocols
negotiated during the handshake as is ``extensions``.
If ``environ`` is provided, it is a copy of the WSGI environ
dictionnary from the underlying WSGI server.
"""
self.stream = Stream(always_mask=False)
"""
Underlying websocket stream that performs the websocket
parsing to high level objects. By default this stream
never masks its messages. Clients using this class should
set the ``stream.always_mask`` fields to ``True``
and ``stream.expect_masking`` fields to ``False``.
"""
self.protocols = protocols
"""
List of protocols supported by this endpoint.
Unused for now.
"""
self.extensions = extensions
"""
List of extensions supported by this endpoint.
Unused for now.
"""
self.sock = sock
"""
Underlying connection.
"""
self.client_terminated = False
"""
Indicates if the client has been marked as terminated.
"""
self.server_terminated = False
"""
Indicates if the server has been marked as terminated.
"""
self.reading_buffer_size = DEFAULT_READING_SIZE
"""
Current connection reading buffer size.
"""
self.environ = environ
"""
WSGI environ dictionary.
"""
self.heartbeat_freq = heartbeat_freq
"""
At which interval the heartbeat will be running.
Set this to `0` or `None` to disable it entirely.
"""
self._local_address = None
self._peer_address = None
@property
def local_address(self):
"""
Local endpoint address as a tuple
"""
if not self._local_address:
self._local_address = self.sock.getsockname()
if len(self._local_address) == 4:
self._local_address = self._local_address[:2]
return self._local_address
@property
def peer_address(self):
"""
Peer endpoint address as a tuple
"""
if not self._peer_address:
self._peer_address = self.sock.getpeername()
if len(self._peer_address) == 4:
self._peer_address = self._peer_address[:2]
return self._peer_address
def opened(self):
"""
Called by the server when the upgrade handshake
has succeeeded.
"""
pass
def close(self, code=1000, reason=''):
"""
Call this method to initiate the websocket connection
closing by sending a close frame to the connected peer.
The ``code`` is the status code representing the
termination's reason.
Once this method is called, the ``server_terminated``
attribute is set. Calling this method several times is
safe as the closing frame will be sent only the first
time.
.. seealso:: Defined Status Codes http://tools.ietf.org/html/rfc6455#section-7.4.1
"""
if not self.server_terminated:
self.server_terminated = True
self._write(self.stream.close(code=code, reason=reason).single(mask=self.stream.always_mask))
def closed(self, code, reason=None):
"""
Called when the websocket stream and connection are finally closed.
The provided ``code`` is status set by the other point and
``reason`` is a human readable message.
.. seealso:: Defined Status Codes http://tools.ietf.org/html/rfc6455#section-7.4.1
"""
pass
@property
def terminated(self):
"""
Returns ``True`` if both the client and server have been
marked as terminated.
"""
return self.client_terminated is True and self.server_terminated is True
@property
def connection(self):
return self.sock
def close_connection(self):
"""
Shutdowns then closes the underlying connection.
"""
if self.sock:
try:
self.sock.shutdown(socket.SHUT_RDWR)
self.sock.close()
except:
pass
finally:
self.sock = None
def ponged(self, pong):
"""
Pong message, as a :class:`messaging.PongControlMessage` instance,
received on the stream.
"""
pass
def received_message(self, message):
"""
Called whenever a complete ``message``, binary or text,
is received and ready for application's processing.
The passed message is an instance of :class:`messaging.TextMessage`
or :class:`messaging.BinaryMessage`.
.. note:: You should override this method in your subclass.
"""
pass
def _write(self, b):
"""
Trying to prevent a write operation
on an already closed websocket stream.
This cannot be bullet proof but hopefully
will catch almost all use cases.
"""
if self.terminated or self.sock is None:
raise RuntimeError("Cannot send on a terminated websocket")
self.sock.sendall(b)
def send(self, payload, binary=False):
"""
Sends the given ``payload`` out.
If ``payload`` is some bytes or a bytearray,
then it is sent as a single message not fragmented.
If ``payload`` is a generator, each chunk is sent as part of
fragmented message.
If ``binary`` is set, handles the payload as a binary message.
"""
message_sender = self.stream.binary_message if binary else self.stream.text_message
if isinstance(payload, basestring) or isinstance(payload, bytearray):
m = message_sender(payload).single(mask=self.stream.always_mask)
self._write(m)
elif isinstance(payload, Message):
data = payload.single(mask=self.stream.always_mask)
self._write(data)
elif type(payload) == types.GeneratorType:
bytes = next(payload)
first = True
for chunk in payload:
self._write(message_sender(bytes).fragment(first=first, mask=self.stream.always_mask))
bytes = chunk
first = False
self._write(message_sender(bytes).fragment(last=True, mask=self.stream.always_mask))
else:
raise ValueError("Unsupported type '%s' passed to send()" % type(payload))
def once(self):
"""
Performs the operation of reading from the underlying
connection in order to feed the stream of bytes.
We start with a small size of two bytes to be read
from the connection so that we can quickly parse an
incoming frame header. Then the stream indicates
whatever size must be read from the connection since
it knows the frame payload length.
It returns `False` if an error occurred at the
socket level or during the bytes processing. Otherwise,
it returns `True`.
"""
if self.terminated:
logger.debug("WebSocket is already terminated")
return False
try:
b = self.sock.recv(self.reading_buffer_size)
except socket.error:
logger.exception("Failed to receive data")
return False
else:
if not self.process(b):
return False
return True
def terminate(self):
"""
Completes the websocket by calling the `closed`
method either using the received closing code
and reason, or when none was received, using
the special `1006` code.
Finally close the underlying connection for
good and cleanup resources by unsetting
the `environ` and `stream` attributes.
"""
s = self.stream
self.client_terminated = self.server_terminated = True
try:
if not s.closing:
self.closed(1006, "Going away")
else:
self.closed(s.closing.code, s.closing.reason)
finally:
self.close_connection()
# Cleaning up resources
s._cleanup()
self.stream = None
self.environ = None
def process(self, bytes):
""" Takes some bytes and process them through the
internal stream's parser. If a message of any kind is
found, performs one of these actions:
* A closing message will initiate the closing handshake
* Errors will initiate a closing handshake
* A message will be passed to the ``received_message`` method
* Pings will see pongs be sent automatically
* Pongs will be passed to the ``ponged`` method
The process should be terminated when this method
returns ``False``.
"""
s = self.stream
if not bytes and self.reading_buffer_size > 0:
return False
self.reading_buffer_size = s.parser.send(bytes) or DEFAULT_READING_SIZE
if s.closing is not None:
logger.debug("Closing message received (%d) '%s'" % (s.closing.code, s.closing.reason))
if not self.server_terminated:
self.close(s.closing.code, s.closing.reason)
else:
self.client_terminated = True
s = None
return False
if s.errors:
for error in s.errors:
logger.debug("Error message received (%d) '%s'" % (error.code, error.reason))
self.close(error.code, error.reason)
s.errors = []
s = None
return False
if s.has_message:
self.received_message(s.message)
if s.message is not None:
s.message.data = None
s.message = None
s = None
return True
if s.pings:
for ping in s.pings:
self._write(s.pong(ping.data))
s.pings = []
if s.pongs:
for pong in s.pongs:
self.ponged(pong)
s.pongs = []
s = None
return True
def run(self):
"""
Performs the operation of reading from the underlying
connection in order to feed the stream of bytes.
We start with a small size of two bytes to be read
from the connection so that we can quickly parse an
incoming frame header. Then the stream indicates
whatever size must be read from the connection since
it knows the frame payload length.
Note that we perform some automatic opererations:
* On a closing message, we respond with a closing
message and finally close the connection
* We respond to pings with pong messages.
* Whenever an error is raised by the stream parsing,
we initiate the closing of the connection with the
appropiate error code.
This method is blocking and should likely be run
in a thread.
"""
self.sock.setblocking(True)
with Heartbeat(self, frequency=self.heartbeat_freq):
s = self.stream
try:
self.opened()
while not self.terminated:
if not self.once():
break
finally:
self.terminate()
class WebSocketBaseClient(object):
def __init__(self, url, protocols=None, version='8'):
self.stream = Stream()
self.url = url
self.protocols = protocols
self.version = version
self.key = b64encode(os.urandom(16))
self.client_terminated = False
self.server_terminated = False
@property
def handshake_headers(self):
parts = urlsplit(self.url)
host = parts.netloc
if ':' in host:
host, port = parts.netloc.split(':')
headers = [
('Host', host),
('Connection', 'Upgrade'),
('Upgrade', 'websocket'),
('Sec-WebSocket-Key', self.key),
('Sec-WebSocket-Origin', self.url),
('Sec-WebSocket-Version', self.version)
]
if self.protocols:
headers.append(('Sec-WebSocket-Protocol', ','.join(self.protocols)))
return headers
@property
def handshake_request(self):
parts = urlsplit(self.url)
headers = self.handshake_headers
request = ["GET %s HTTP/1.1" % parts.path]
for header, value in headers:
request.append("%s: %s" % (header, value))
request.append('\r\n')
return '\r\n'.join(request)
def process_response_line(self, response_line):
protocol, code, status = response_line.split(' ', 2)
if code != '101':
raise HandshakeError("Invalid response status: %s %s" % (code, status))
def process_handshake_header(self, headers):
protocols = []
extensions = []
headers = headers.strip()
for header_line in headers.split('\r\n'):
header, value = header_line.split(':', 1)
header = header.strip().lower()
value = value.strip().lower()
if header == 'upgrade' and value != 'websocket':
raise HandshakeError("Invalid Upgrade header: %s" % value)
elif header == 'connection' and value != 'upgrade':
raise HandshakeError("Invalid Connection header: %s" % value)
elif header == 'sec-websocket-accept':
match = b64encode(sha1(self.key + WS_KEY).digest())
if value != match.lower():
raise HandshakeError("Invalid challenge response: %s" % value)
elif header == 'sec-websocket-protocol':
protocols = ','.join(value)
elif header == 'sec-websocket-extensions':
extensions = ','.join(value)
return protocols, extensions
def opened(self, protocols, extensions):
pass
def received_message(self, m):
pass
def closed(self, code, reason=None):
pass
@property
def terminated(self):
return self.client_terminated is True and self.server_terminated is True
def close(self, reason='', code=1000):
if not self.client_terminated:
self.client_terminated = True
self.write_to_connection(self.stream.close(code=code, reason=reason))
def connect(self):
raise NotImplemented()
def write_to_connection(self, bytes):
raise NotImplemented()
def read_from_connection(self, amount):
raise NotImplemented()
def close_connection(self):
raise NotImplemented()
def send(self, payload, binary=False):
if isinstance(payload, basestring):
if not binary:
self.write_to_connection(self.stream.text_message(payload).single(mask=True))
else:
self.write_to_connection(self.stream.binary_message(payload).single(mask=True))
elif isinstance(payload, dict):
self.write_to_connection(self.stream.text_message(json.dumps(payload)).single(mask=True))
elif type(payload) == types.GeneratorType:
bytes = payload.next()
first = True
for chunk in payload:
if not binary:
self.write_to_connection(self.stream.text_message(bytes).fragment(first=first, mask=True))
else:
self.write_to_connection(self.stream.binary_message(payload).fragment(first=first, mask=True))
bytes = chunk
first = False
if not binary:
self.write_to_connection(self.stream.text_message(bytes).fragment(last=True, mask=True))
else:
self.write_to_connection(self.stream.text_message(bytes).fragment(last=True, mask=True))
