mirror of
https://github.com/Grasscutters/mitmproxy.git
synced 2024-11-23 00:01:36 +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)
|
||||
|
||||
|
||||
class WebSocketFrameValidationException(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class Frame(object):
|
||||
"""
|
||||
Represents one websockets frame.
|
||||
@ -95,7 +91,8 @@ class Frame(object):
|
||||
|
||||
if from_client:
|
||||
mask_bit = 1
|
||||
masking_key = random_masking_key()
|
||||
# Random masking key
|
||||
masking_key = os.urandom(4)
|
||||
payload = apply_mask(message, masking_key)
|
||||
else:
|
||||
mask_bit = 0
|
||||
@ -164,12 +161,6 @@ class Frame(object):
|
||||
"""
|
||||
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):
|
||||
"""
|
||||
Serialize the frame back into the wire format, returns a bytestring
|
||||
@ -308,10 +299,6 @@ def apply_mask(message, masking_key):
|
||||
return result
|
||||
|
||||
|
||||
def random_masking_key():
|
||||
return os.urandom(4)
|
||||
|
||||
|
||||
def client_handshake_headers(key=None, version=VERSION):
|
||||
"""
|
||||
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
|
||||
import io
|
||||
from netlib import tcp, test, websockets, http
|
||||
import os
|
||||
from nose.tools import raises
|
||||
|
||||
@ -120,17 +119,11 @@ class TestWebSockets(test.ServerTestBase):
|
||||
self.random_bytes(num_bytes), is_client
|
||||
)
|
||||
assert frame == websockets.Frame.from_bytes(
|
||||
frame.safe_to_bytes()
|
||||
frame.to_bytes()
|
||||
)
|
||||
|
||||
bytes = b'\x81\x03cba'
|
||||
assert websockets.Frame.from_bytes(bytes).safe_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()
|
||||
assert websockets.Frame.from_bytes(bytes).to_bytes() == bytes
|
||||
|
||||
def test_check_server_handshake(self):
|
||||
resp = http.Response(
|
||||
|
Loading…
Reference in New Issue
Block a user