mirror of
https://github.com/Grasscutters/mitmproxy.git
synced 2024-11-27 02:24:18 +00:00
Whitespace, interface simplification
- safe_tobytes doesn't buy us much - move masking key generation inline
This commit is contained in:
parent
3e0a71ea34
commit
1b509d5aea
@ -29,10 +29,6 @@ class CONST(object):
|
|||||||
MAX_64_BIT_INT = (1 << 64)
|
MAX_64_BIT_INT = (1 << 64)
|
||||||
|
|
||||||
|
|
||||||
class WebSocketFrameValidationException(Exception):
|
|
||||||
pass
|
|
||||||
|
|
||||||
|
|
||||||
class Frame(object):
|
class Frame(object):
|
||||||
"""
|
"""
|
||||||
Represents one websockets frame.
|
Represents one websockets frame.
|
||||||
@ -95,7 +91,8 @@ class Frame(object):
|
|||||||
|
|
||||||
if from_client:
|
if from_client:
|
||||||
mask_bit = 1
|
mask_bit = 1
|
||||||
masking_key = random_masking_key()
|
# Random masking key
|
||||||
|
masking_key = os.urandom(4)
|
||||||
payload = apply_mask(message, masking_key)
|
payload = apply_mask(message, masking_key)
|
||||||
else:
|
else:
|
||||||
mask_bit = 0
|
mask_bit = 0
|
||||||
@ -164,12 +161,6 @@ class Frame(object):
|
|||||||
"""
|
"""
|
||||||
return cls.from_file(io.BytesIO(bytestring))
|
return cls.from_file(io.BytesIO(bytestring))
|
||||||
|
|
||||||
def safe_to_bytes(self):
|
|
||||||
if self.is_valid():
|
|
||||||
return self.to_bytes()
|
|
||||||
else:
|
|
||||||
raise WebSocketFrameValidationException()
|
|
||||||
|
|
||||||
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
|
||||||
@ -308,10 +299,6 @@ def apply_mask(message, masking_key):
|
|||||||
return result
|
return result
|
||||||
|
|
||||||
|
|
||||||
def random_masking_key():
|
|
||||||
return os.urandom(4)
|
|
||||||
|
|
||||||
|
|
||||||
def client_handshake_headers(key=None, version=VERSION):
|
def client_handshake_headers(key=None, version=VERSION):
|
||||||
"""
|
"""
|
||||||
Create the headers for a valid HTTP upgrade request. If Key is not
|
Create the headers for a valid HTTP upgrade request. If Key is not
|
||||||
|
@ -1,5 +1,4 @@
|
|||||||
from netlib import tcp, test, websockets, http, odict
|
from netlib import tcp, test, websockets, http
|
||||||
import io
|
|
||||||
import os
|
import os
|
||||||
from nose.tools import raises
|
from nose.tools import raises
|
||||||
|
|
||||||
@ -120,17 +119,11 @@ class TestWebSockets(test.ServerTestBase):
|
|||||||
self.random_bytes(num_bytes), is_client
|
self.random_bytes(num_bytes), is_client
|
||||||
)
|
)
|
||||||
assert frame == websockets.Frame.from_bytes(
|
assert frame == websockets.Frame.from_bytes(
|
||||||
frame.safe_to_bytes()
|
frame.to_bytes()
|
||||||
)
|
)
|
||||||
|
|
||||||
bytes = b'\x81\x03cba'
|
bytes = b'\x81\x03cba'
|
||||||
assert websockets.Frame.from_bytes(bytes).safe_to_bytes() == bytes
|
assert websockets.Frame.from_bytes(bytes).to_bytes() == bytes
|
||||||
|
|
||||||
@raises(websockets.WebSocketFrameValidationException)
|
|
||||||
def test_safe_to_bytes(self):
|
|
||||||
frame = websockets.Frame.default(self.random_bytes(8))
|
|
||||||
frame.actual_payload_length = 1 # corrupt the frame
|
|
||||||
frame.safe_to_bytes()
|
|
||||||
|
|
||||||
def test_check_server_handshake(self):
|
def test_check_server_handshake(self):
|
||||||
resp = http.Response(
|
resp = http.Response(
|
||||||
|
Loading…
Reference in New Issue
Block a user