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

websockets: whitespace, PEP8

Browse Source
This commit is contained in:
Aldo Cortesi 2015-04-17 13:57:39 +12:00
parent 3bbafa24bd
commit 488c25d812
1 changed files with 96 additions and 73 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

169
netlib/websockets/websockets.py
View File

@ -1,31 +1,34 @@
from __future__ import absolute_import from __future__ import absolute_import
from base64 import b64encode import base64
from hashlib import sha1 import hashlib
from mimetools import Message import mimetools
from netlib import tcp import StringIO
from netlib import utils
from StringIO import StringIO
import os import os
import SocketServer
import struct import struct
import io import io
# Colleciton of utility functions that implement small portions of the RFC6455 WebSockets Protocol from .. import utils
# Useful for building WebSocket clients and servers.
# Colleciton of utility functions that implement small portions of the RFC6455
# WebSockets Protocol Useful for building WebSocket clients and servers.
# #
# Emphassis is on readabilty, simplicity and modularity, not performance or completeness # Emphassis is on readabilty, simplicity and modularity, not performance or
# # completeness
# This is a work in progress and does not yet contain all the utilites need to create fully complient client/servers
# #
# This is a work in progress and does not yet contain all the utilites need to
# create fully complient client/servers #
# Spec: https://tools.ietf.org/html/rfc6455 # Spec: https://tools.ietf.org/html/rfc6455
# The magic sha that websocket servers must know to prove they understand RFC6455 # The magic sha that websocket servers must know to prove they understand
# RFC6455
websockets_magic = '258EAFA5-E914-47DA-95CA-C5AB0DC85B11' websockets_magic = '258EAFA5-E914-47DA-95CA-C5AB0DC85B11'
class WebSocketFrameValidationException(Exception): class WebSocketFrameValidationException(Exception):
pass pass
class WebSocketsFrame(object): class WebSocketsFrame(object):
""" """
Represents one websockets frame. Represents one websockets frame.
@ -33,7 +36,7 @@ class WebSocketsFrame(object):
from_bytes() is also avaliable. from_bytes() is also avaliable.
WebSockets Frame as defined in RFC6455 WebSockets Frame as defined in RFC6455
0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
+-+-+-+-+-------+-+-------------+-------------------------------+ +-+-+-+-+-------+-+-------------+-------------------------------+
|F|R|R|R| opcode|M| Payload len | Extended payload length | |F|R|R|R| opcode|M| Payload len | Extended payload length |
@ -62,7 +65,7 @@ class WebSocketsFrame(object):
rsv1 = 0, # decimal integer 1 or 0 rsv1 = 0, # decimal integer 1 or 0
rsv2 = 0, # decimal integer 1 or 0 rsv2 = 0, # decimal integer 1 or 0
rsv3 = 0, # decimal integer 1 or 0 rsv3 = 0, # decimal integer 1 or 0
payload = None, # bytestring payload = None, # bytestring
masking_key = None, # 32 bit byte string masking_key = None, # 32 bit byte string
actual_payload_length = None, # any decimal integer actual_payload_length = None, # any decimal integer
): ):
@ -81,18 +84,17 @@ class WebSocketsFrame(object):
@classmethod @classmethod
def from_bytes(cls, bytestring): def from_bytes(cls, bytestring):
""" """
Construct a websocket frame from an in-memory bytestring Construct a websocket frame from an in-memory bytestring to construct
to construct a frame from a stream of bytes, use from_byte_stream() directly a frame from a stream of bytes, use from_byte_stream() directly
""" """
return cls.from_byte_stream(io.BytesIO(bytestring).read) return cls.from_byte_stream(io.BytesIO(bytestring).read)
@classmethod @classmethod
def default(cls, message, from_client = False): def default(cls, message, from_client = False):
""" """
Construct a basic websocket frame from some default values. Construct a basic websocket frame from some default values.
Creates a non-fragmented text frame. Creates a non-fragmented text frame.
""" """
length_code, actual_length = get_payload_length_pair(message) length_code, actual_length = get_payload_length_pair(message)
if from_client: if from_client:
@ -103,7 +105,7 @@ class WebSocketsFrame(object):
mask_bit = 0 mask_bit = 0
masking_key = None masking_key = None
payload = message payload = message
return cls( return cls(
fin = 1, # final frame fin = 1, # final frame
opcode = 1, # text opcode = 1, # text
@ -117,10 +119,10 @@ class WebSocketsFrame(object):
def is_valid(self): def is_valid(self):
""" """
Validate websocket frame invariants, call at anytime to ensure the WebSocketsFrame Validate websocket frame invariants, call at anytime to ensure the
has not been corrupted. WebSocketsFrame has not been corrupted.
""" """
try: try:
assert 0 <= self.fin <= 1 assert 0 <= self.fin <= 1
assert 0 <= self.rsv1 <= 1 assert 0 <= self.rsv1 <= 1
assert 0 <= self.rsv2 <= 1 assert 0 <= self.rsv2 <= 1
@ -128,18 +130,18 @@ class WebSocketsFrame(object):
assert 1 <= self.opcode <= 4 assert 1 <= self.opcode <= 4
assert 0 <= self.mask_bit <= 1 assert 0 <= self.mask_bit <= 1
assert 1 <= self.payload_length_code <= 127 assert 1 <= self.payload_length_code <= 127
if self.mask_bit == 1: if self.mask_bit == 1:
assert 1 <= len(self.masking_key) <= 4 assert 1 <= len(self.masking_key) <= 4
else: else:
assert self.masking_key == None assert self.masking_key is None
assert self.actual_payload_length == len(self.payload) assert self.actual_payload_length == len(self.payload)
if self.payload is not None and self.masking_key is not None: if self.payload is not None and self.masking_key is not None:
assert apply_mask(self.payload, self.masking_key) == self.decoded_payload assert apply_mask(self.payload, self.masking_key) == self.decoded_payload
return True return True
except AssertionError: except AssertionError:
return False return False
@ -165,30 +167,32 @@ class WebSocketsFrame(object):
def to_bytes(self): def to_bytes(self):
""" """
Serialize the frame back into the wire format, returns a bytestring Serialize the frame back into the wire format, returns a bytestring If
If you haven't checked is_valid_frame() then there's no guarentees that the you haven't checked is_valid_frame() then there's no guarentees that
serialized bytes will be correct. see safe_to_bytes() the serialized bytes will be correct. see safe_to_bytes()
""" """
max_16_bit_int = (1 << 16) max_16_bit_int = (1 << 16)
max_64_bit_int = (1 << 63) max_64_bit_int = (1 << 63)
# break down of the bit-math used to construct the first byte from the frame's integer values # break down of the bit-math used to construct the first byte from the
# first shift the significant bit into the correct position # frame's integer values first shift the significant bit into the
# correct position
# 00000001 << 7 = 10000000 # 00000001 << 7 = 10000000
# ... # ...
# then combine: # then combine:
# #
# 10000000 fin # 10000000 fin
# 01000000 res1 # 01000000 res1
# 00100000 res2 # 00100000 res2
# 00010000 res3 # 00010000 res3
# 00000001 opcode # 00000001 opcode
# -------- OR # -------- OR
# 11110001 = first_byte # 11110001 = first_byte
first_byte = (self.fin << 7) | (self.rsv1 << 6) | (self.rsv2 << 4) | (self.rsv3 << 4) | self.opcode first_byte = (self.fin << 7) | (self.rsv1 << 6) |\
(self.rsv2 << 4) | (self.rsv3 << 4) | self.opcode
second_byte = (self.mask_bit << 7) | self.payload_length_code second_byte = (self.mask_bit << 7) | self.payload_length_code
bytes = chr(first_byte) + chr(second_byte) bytes = chr(first_byte) + chr(second_byte)
@ -199,11 +203,13 @@ class WebSocketsFrame(object):
elif self.actual_payload_length < max_16_bit_int: elif self.actual_payload_length < max_16_bit_int:
# '!H' pack as 16 bit unsigned short # '!H' pack as 16 bit unsigned short
bytes += struct.pack('!H', self.actual_payload_length) # add 2 byte extended payload length # 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 < max_64_bit_int:
# '!Q' = pack as 64 bit unsigned long long # '!Q' = pack as 64 bit unsigned long long
bytes += struct.pack('!Q', self.actual_payload_length) # add 8 bytes extended payload length # add 8 bytes extended payload length
bytes += struct.pack('!Q', self.actual_payload_length)
if self.masking_key is not None: if self.masking_key is not None:
bytes += self.masking_key bytes += self.masking_key
@ -212,43 +218,46 @@ class WebSocketsFrame(object):
return bytes return bytes
@classmethod @classmethod
def from_byte_stream(cls, read_bytes): def from_byte_stream(cls, read_bytes):
""" """
read a websockets frame sent by a server or client read a websockets frame sent by a server or client
read_bytes is a function that can be backed read_bytes is a function that can be backed
by sockets or by any byte reader. So this by sockets or by any byte reader. So this
function may be used to read frames from disk/wire/memory function may be used to read frames from disk/wire/memory
""" """
first_byte = utils.bytes_to_int(read_bytes(1)) first_byte = utils.bytes_to_int(read_bytes(1))
second_byte = utils.bytes_to_int(read_bytes(1)) second_byte = utils.bytes_to_int(read_bytes(1))
fin = first_byte >> 7 # grab the left most bit # grab the left most bit
opcode = first_byte & 15 # grab right most 4 bits by and-ing with 00001111 fin = first_byte >> 7
mask_bit = second_byte >> 7 # grab left most bit # grab right most 4 bits by and-ing with 00001111
payload_length = second_byte & 127 # grab the next 7 bits opcode = first_byte & 15
# grab left most bit
mask_bit = second_byte >> 7
# grab the next 7 bits
payload_length = second_byte & 127
# payload_lengthy > 125 indicates you need to read more bytes # payload_lengthy > 125 indicates you need to read more bytes
# to get the actual payload length # to get the actual payload length
if payload_length <= 125: if payload_length <= 125:
actual_payload_length = payload_length actual_payload_length = payload_length
elif payload_length == 126: elif payload_length == 126:
actual_payload_length = utils.bytes_to_int(read_bytes(2)) actual_payload_length = utils.bytes_to_int(read_bytes(2))
elif payload_length == 127: elif payload_length == 127:
actual_payload_length = utils.bytes_to_int(read_bytes(8)) actual_payload_length = utils.bytes_to_int(read_bytes(8))
# masking key only present if mask bit set # masking key only present if mask bit set
if mask_bit == 1: if mask_bit == 1:
masking_key = read_bytes(4) masking_key = read_bytes(4)
else: else:
masking_key = None masking_key = None
payload = read_bytes(actual_payload_length) payload = read_bytes(actual_payload_length)
if mask_bit == 1: if mask_bit == 1:
decoded_payload = apply_mask(payload, masking_key) decoded_payload = apply_mask(payload, masking_key)
else: else:
@ -295,12 +304,15 @@ def apply_mask(message, masking_key):
result += chr(ord(char) ^ masks[len(result) % 4]) result += chr(ord(char) ^ masks[len(result) % 4])
return result return result
def random_masking_key(): def random_masking_key():
return os.urandom(4) return os.urandom(4)
def create_client_handshake(host, port, key, version, resource): def create_client_handshake(host, port, key, version, resource):
""" """
WebSockets connections are intiated by the client with a valid HTTP upgrade request WebSockets connections are intiated by the client with a valid HTTP
upgrade request
""" """
headers = [ headers = [
('Host', '%s:%s' % (host, port)), ('Host', '%s:%s' % (host, port)),
@ -312,10 +324,11 @@ def create_client_handshake(host, port, key, version, resource):
request = "GET %s HTTP/1.1" % resource request = "GET %s HTTP/1.1" % resource
return build_handshake(headers, request) return build_handshake(headers, request)
def create_server_handshake(key): def create_server_handshake(key):
""" """
The server response is a valid HTTP 101 response. The server response is a valid HTTP 101 response.
""" """
headers = [ headers = [
('Connection', 'Upgrade'), ('Connection', 'Upgrade'),
('Upgrade', 'websocket'), ('Upgrade', 'websocket'),
@ -332,12 +345,13 @@ def build_handshake(headers, request):
handshake.append(b'\r\n') handshake.append(b'\r\n')
return b'\r\n'.join(handshake) return b'\r\n'.join(handshake)
def read_handshake(read_bytes, num_bytes_per_read): def read_handshake(read_bytes, num_bytes_per_read):
""" """
From provided function that reads bytes, read in a From provided function that reads bytes, read in a
complete HTTP request, which terminates with a CLRF complete HTTP request, which terminates with a CLRF
""" """
response = b'' response = b''
doubleCLRF = b'\r\n\r\n' doubleCLRF = b'\r\n\r\n'
while True: while True:
bytes = read_bytes(num_bytes_per_read) bytes = read_bytes(num_bytes_per_read)
@ -348,14 +362,15 @@ def read_handshake(read_bytes, num_bytes_per_read):
break break
return response return response
def get_payload_length_pair(payload_bytestring): def get_payload_length_pair(payload_bytestring):
""" """
A websockets frame contains an initial length_code, and an optional A websockets frame contains an initial length_code, and an optional
extended length code to represent the actual length if length code is larger extended length code to represent the actual length if length code is
than 125 larger than 125
""" """
actual_length = len(payload_bytestring) actual_length = len(payload_bytestring)
if actual_length <= 125: if actual_length <= 125:
length_code = actual_length length_code = actual_length
elif actual_length >= 126 and actual_length <= 65535: elif actual_length >= 126 and actual_length <= 65535:
@ -364,6 +379,7 @@ def get_payload_length_pair(payload_bytestring):
length_code = 127 length_code = 127
return (length_code, actual_length) return (length_code, actual_length)
def process_handshake_from_client(handshake): def process_handshake_from_client(handshake):
headers = headers_from_http_message(handshake) headers = headers_from_http_message(handshake)
if headers.get("Upgrade", None) != "websocket": if headers.get("Upgrade", None) != "websocket":
@ -371,6 +387,7 @@ def process_handshake_from_client(handshake):
key = headers['Sec-WebSocket-Key'] key = headers['Sec-WebSocket-Key']
return key return key
def process_handshake_from_server(handshake, client_nounce): def process_handshake_from_server(handshake, client_nounce):
headers = headers_from_http_message(handshake) headers = headers_from_http_message(handshake)
if headers.get("Upgrade", None) != "websocket": if headers.get("Upgrade", None) != "websocket":
@ -378,12 +395,18 @@ def process_handshake_from_server(handshake, client_nounce):
key = headers['Sec-WebSocket-Accept'] key = headers['Sec-WebSocket-Accept']
return key return key
def headers_from_http_message(http_message): def headers_from_http_message(http_message):
return Message(StringIO(http_message.split('\r\n', 1)[1])) return mimetools.Message(
StringIO.StringIO(http_message.split('\r\n', 1)[1])
)
def create_server_nounce(client_nounce): def create_server_nounce(client_nounce):
return b64encode(sha1(client_nounce + websockets_magic).hexdigest().decode('hex')) return base64.b64encode(
hashlib.sha1(client_nounce + websockets_magic).hexdigest().decode('hex')
)
def create_client_nounce(): def create_client_nounce():
return b64encode(os.urandom(16)).decode('utf-8') return base64.b64encode(os.urandom(16)).decode('utf-8')