Grasscutters/mitmproxy
3
0
Fork 0
mirror of https://github.com/Grasscutters/mitmproxy.git synced 2024-11-23 08:11:00 +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

101
netlib/websockets/websockets.py
View File

@ -1,31 +1,34 @@
from __future__ import absolute_import
from base64 import b64encode
from hashlib import sha1
from mimetools import Message
from netlib import tcp
from netlib import utils
from StringIO import StringIO
import base64
import hashlib
import mimetools
import StringIO
import os
import SocketServer
import struct
import io
# Colleciton of utility functions that implement small portions of the RFC6455 WebSockets Protocol
# Useful for building WebSocket clients and servers.
from .. import utils
# 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
#
# This is a work in progress and does not yet contain all the utilites need to create fully complient client/servers
# 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 #
# 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'
class WebSocketFrameValidationException(Exception):
pass
class WebSocketsFrame(object):
"""
Represents one websockets frame.
@ -81,12 +84,11 @@ class WebSocketsFrame(object):
@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
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):
"""
@ -117,8 +119,8 @@ class WebSocketsFrame(object):
def is_valid(self):
"""
Validate websocket frame invariants, call at anytime to ensure the WebSocketsFrame
has not been corrupted.
Validate websocket frame invariants, call at anytime to ensure the
WebSocketsFrame has not been corrupted.
"""
try:
assert 0 <= self.fin <= 1
@ -132,7 +134,7 @@ class WebSocketsFrame(object):
if self.mask_bit == 1:
assert 1 <= len(self.masking_key) <= 4
else:
assert self.masking_key == None
assert self.masking_key is None
assert self.actual_payload_length == len(self.payload)
@ -165,16 +167,17 @@ class WebSocketsFrame(object):
def to_bytes(self):
"""
Serialize the frame back into the wire format, returns a bytestring
If you haven't checked is_valid_frame() then there's no guarentees that the
serialized bytes will be correct. see safe_to_bytes()
Serialize the frame back into the wire format, returns a bytestring 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 correct position
# 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
# correct position
# 00000001 << 7 = 10000000
# ...
# then combine:
@ -187,7 +190,8 @@ class WebSocketsFrame(object):
# -------- OR
# 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
@ -199,11 +203,13 @@ class WebSocketsFrame(object):
elif self.actual_payload_length < max_16_bit_int:
# '!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:
# '!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:
bytes += self.masking_key
@ -212,7 +218,6 @@ class WebSocketsFrame(object):
return bytes
@classmethod
def from_byte_stream(cls, read_bytes):
"""
@ -225,10 +230,14 @@ class WebSocketsFrame(object):
first_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
opcode = first_byte & 15 # grab right most 4 bits by and-ing with 00001111
mask_bit = second_byte >> 7 # grab left most bit
payload_length = second_byte & 127 # grab the next 7 bits
# grab the left most bit
fin = first_byte >> 7
# grab right most 4 bits by and-ing with 00001111
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
# to get the actual payload length
@ -295,12 +304,15 @@ def apply_mask(message, masking_key):
result += chr(ord(char) ^ masks[len(result) % 4])
return result
def random_masking_key():
return os.urandom(4)
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 = [
('Host', '%s:%s' % (host, port)),
@ -312,6 +324,7 @@ def create_client_handshake(host, port, key, version, resource):
request = "GET %s HTTP/1.1" % resource
return build_handshake(headers, request)
def create_server_handshake(key):
"""
The server response is a valid HTTP 101 response.
@ -332,6 +345,7 @@ def build_handshake(headers, request):
handshake.append(b'\r\n')
return b'\r\n'.join(handshake)
def read_handshake(read_bytes, num_bytes_per_read):
"""
From provided function that reads bytes, read in a
@ -348,11 +362,12 @@ def read_handshake(read_bytes, num_bytes_per_read):
break
return response
def get_payload_length_pair(payload_bytestring):
"""
A websockets frame contains an initial length_code, and an optional
extended length code to represent the actual length if length code is larger
than 125
extended length code to represent the actual length if length code is
larger than 125
"""
actual_length = len(payload_bytestring)
@ -364,6 +379,7 @@ def get_payload_length_pair(payload_bytestring):
length_code = 127
return (length_code, actual_length)
def process_handshake_from_client(handshake):
headers = headers_from_http_message(handshake)
if headers.get("Upgrade", None) != "websocket":
@ -371,6 +387,7 @@ def process_handshake_from_client(handshake):
key = headers['Sec-WebSocket-Key']
return key
def process_handshake_from_server(handshake, client_nounce):
headers = headers_from_http_message(handshake)
if headers.get("Upgrade", None) != "websocket":
@ -378,12 +395,18 @@ def process_handshake_from_server(handshake, client_nounce):
key = headers['Sec-WebSocket-Accept']
return key
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):
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():
return b64encode(os.urandom(16)).decode('utf-8')
return base64.b64encode(os.urandom(16)).decode('utf-8')