Grasscutters/mitmproxy
3
0
Fork 0
mirror of https://github.com/Grasscutters/mitmproxy.git synced 2024-11-27 10:26:23 +00:00
Code Issues Projects Releases Wiki Activity

Merge pull request #55 from Chandler/websockets

Browse Source 
small websockets cleanup
This commit is contained in:
Aldo Cortesi 2015-04-20 17:31:42 +12:00
commit 0141629c08
2 changed files with 77 additions and 45 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

58
netlib/websockets.py
View File

@ -25,6 +25,11 @@ from . import utils
websockets_magic = '258EAFA5-E914-47DA-95CA-C5AB0DC85B11'
class CONST(object):
MAX_16_BIT_INT = (1 << 16)
MAX_64_BIT_INT = (1 << 64)
class WebSocketFrameValidationException(Exception):
pass
@ -81,14 +86,6 @@ class Frame(object):
self.decoded_payload = decoded_payload
self.actual_payload_length = actual_payload_length
@classmethod
def from_bytes(cls, bytestring):
"""
Construct a websocket frame from an in-memory bytestring to construct
a frame from a stream of bytes, use from_byte_stream() directly
"""
return cls.from_byte_stream(io.BytesIO(bytestring).read)
@classmethod
def default(cls, message, from_client = False):
"""
@ -145,7 +142,7 @@ class Frame(object):
except AssertionError:
return False
def human_readable(self):
def human_readable(self): # pragma: nocover
return "\n".join([
("fin - " + str(self.fin)),
("rsv1 - " + str(self.rsv1)),
@ -160,6 +157,14 @@ class Frame(object):
("actual_payload_length - " + str(self.actual_payload_length))
])
@classmethod
def from_bytes(cls, bytestring):
"""
Construct a websocket frame from an in-memory bytestring
to construct a frame from a stream of bytes, use from_file() directly
"""
return cls.from_file(io.BytesIO(bytestring))
def safe_to_bytes(self):
if self.is_valid():
return self.to_bytes()
@ -172,8 +177,6 @@ class Frame(object):
If you haven't checked is_valid_frame() then there's no guarentees
that the serialized bytes will be correct. see safe_to_bytes()
"""
max_16_bit_int = (1 << 16)
max_64_bit_int = (1 << 63)
# break down of the bit-math used to construct the first byte from the
# frame's integer values first shift the significant bit into the
@ -199,11 +202,11 @@ class Frame(object):
if self.actual_payload_length < 126:
pass
elif self.actual_payload_length < max_16_bit_int:
elif self.actual_payload_length < CONST.MAX_16_BIT_INT:
# '!H' pack as 16 bit unsigned short
# add 2 byte extended payload length
bytes += struct.pack('!H', self.actual_payload_length)
elif self.actual_payload_length < max_64_bit_int:
elif self.actual_payload_length < CONST.MAX_64_BIT_INT:
# '!Q' = pack as 64 bit unsigned long long
# add 8 bytes extended payload length
bytes += struct.pack('!Q', self.actual_payload_length)
@ -214,17 +217,20 @@ class Frame(object):
bytes += self.payload # already will be encoded if neccessary
return bytes
def to_file(self, writer):
writer.write(self.to_bytes())
writer.flush()
@classmethod
def from_byte_stream(cls, read_bytes):
def from_file(cls, reader):
"""
read a websockets frame sent by a server or client
read_bytes is a function that can be backed
by sockets or by any byte reader. So this
function may be used to read frames from disk/wire/memory
reader is a "file like" object that could be backed by a network stream or a disk
or an in memory stream reader
"""
first_byte = utils.bytes_to_int(read_bytes(1))
second_byte = utils.bytes_to_int(read_bytes(1))
first_byte = utils.bytes_to_int(reader.read(1))
second_byte = utils.bytes_to_int(reader.read(1))
# grab the left most bit
fin = first_byte >> 7
@ -241,18 +247,18 @@ class Frame(object):
actual_payload_length = payload_length
elif payload_length == 126:
actual_payload_length = utils.bytes_to_int(read_bytes(2))
actual_payload_length = utils.bytes_to_int(reader.read(2))
elif payload_length == 127:
actual_payload_length = utils.bytes_to_int(read_bytes(8))
actual_payload_length = utils.bytes_to_int(reader.read(8))
# masking key only present if mask bit set
if mask_bit == 1:
masking_key = read_bytes(4)
masking_key = reader.read(4)
else:
masking_key = None
payload = read_bytes(actual_payload_length)
payload = reader.read(actual_payload_length)
if mask_bit == 1:
decoded_payload = apply_mask(payload, masking_key)
@ -344,7 +350,7 @@ def build_handshake(headers, request):
return b'\r\n'.join(handshake)
def read_handshake(read_bytes, num_bytes_per_read):
def read_handshake(reader, num_bytes_per_read):
"""
From provided function that reads bytes, read in a
complete HTTP request, which terminates with a CLRF
@ -352,7 +358,7 @@ def read_handshake(read_bytes, num_bytes_per_read):
response = b''
doubleCLRF = b'\r\n\r\n'
while True:
bytes = read_bytes(num_bytes_per_read)
bytes = reader.read(num_bytes_per_read)
if not bytes:
break
response += bytes
@ -386,7 +392,7 @@ def process_handshake_from_client(handshake):
return key
def process_handshake_from_server(handshake, client_nounce):
def process_handshake_from_server(handshake):
headers = headers_from_http_message(handshake)
if headers.get("Upgrade", None) != "websocket":
return

58
test/test_websockets.py
View File

@ -1,6 +1,7 @@
from netlib import tcp
from netlib import test
from netlib import websockets
import io
import os
from nose.tools import raises
@ -20,16 +21,15 @@ class WebSocketsEchoHandler(tcp.BaseHandler):
self.read_next_message()
def read_next_message(self):
decoded = websockets.Frame.from_byte_stream(self.rfile.read).decoded_payload
decoded = websockets.Frame.from_file(self.rfile).decoded_payload
self.on_message(decoded)
def send_message(self, message):
frame = websockets.Frame.default(message, from_client = False)
self.wfile.write(frame.safe_to_bytes())
self.wfile.flush()
frame.to_file(self.wfile)
def handshake(self):
client_hs = websockets.read_handshake(self.rfile.read, 1)
client_hs = websockets.read_handshake(self.rfile, 1)
key = websockets.process_handshake_from_client(client_hs)
response = websockets.create_server_handshake(key)
self.wfile.write(response)
@ -62,22 +62,18 @@ class WebSocketsClient(tcp.TCPClient):
self.wfile.write(handshake)
self.wfile.flush()
server_handshake = websockets.read_handshake(self.rfile.read, 1)
server_nounce = websockets.process_handshake_from_server(
server_handshake, self.client_nounce
)
server_handshake = websockets.read_handshake(self.rfile, 1)
server_nounce = websockets.process_handshake_from_server(server_handshake)
if not server_nounce == websockets.create_server_nounce(self.client_nounce):
self.close()
def read_next_message(self):
return websockets.Frame.from_byte_stream(self.rfile.read).payload
return websockets.Frame.from_file(self.rfile).payload
def send_message(self, message):
frame = websockets.Frame.default(message, from_client = True)
self.wfile.write(frame.safe_to_bytes())
self.wfile.flush()
frame.to_file(self.wfile)
class TestWebSockets(test.ServerTestBase):
handler = WebSocketsEchoHandler
@ -128,10 +124,10 @@ class TestWebSockets(test.ServerTestBase):
frame = websockets.Frame.default(
self.random_bytes(num_bytes), is_client
)
assert frame == websockets.Frame.from_bytes(frame.to_bytes())
assert frame == websockets.Frame.from_bytes(frame.safe_to_bytes())
bytes = b'\x81\x11cba'
assert websockets.Frame.from_bytes(bytes).to_bytes() == bytes
bytes = b'\x81\x03cba'
assert websockets.Frame.from_bytes(bytes).safe_to_bytes() == bytes
@raises(websockets.WebSocketFrameValidationException)
def test_safe_to_bytes(self):
@ -139,10 +135,40 @@ class TestWebSockets(test.ServerTestBase):
frame.actual_payload_length = 1 # corrupt the frame
frame.safe_to_bytes()
def test_handshake(self):
bad_upgrade = "not_websockets"
bad_header_handshake = websockets.build_handshake([
('Host', '%s:%s' % ("a", "b")),
('Connection', "c"),
('Upgrade', bad_upgrade),
('Sec-WebSocket-Key', "d"),
('Sec-WebSocket-Version', "e")
], "f")
# check behavior when required header values are missing
assert None == websockets.process_handshake_from_server(bad_header_handshake)
assert None == websockets.process_handshake_from_client(bad_header_handshake)
key = "test_key"
client_handshake = websockets.create_client_handshake("a","b",key,"d","e")
assert key == websockets.process_handshake_from_client(client_handshake)
server_handshake = websockets.create_server_handshake(key)
assert websockets.create_server_nounce(key) == websockets.process_handshake_from_server(server_handshake)
handshake = websockets.create_client_handshake("a","b","c","d","e")
stream = io.BytesIO(handshake)
assert handshake == websockets.read_handshake(stream, 1)
# ensure readhandshake doesn't loop forever on empty stream
empty_stream = io.BytesIO("")
assert "" == websockets.read_handshake(empty_stream, 1)
class BadHandshakeHandler(WebSocketsEchoHandler):
def handshake(self):
client_hs = websockets.read_handshake(self.rfile.read, 1)
client_hs = websockets.read_handshake(self.rfile, 1)
websockets.process_handshake_from_client(client_hs)
response = websockets.create_server_handshake("malformed_key")
self.wfile.write(response)