mirror of
https://github.com/Grasscutters/mitmproxy.git
synced 2024-11-23 08:11:00 +00:00
275 lines
8.8 KiB
Python
275 lines
8.8 KiB
Python
import sys
|
|
import os
|
|
import struct
|
|
import io
|
|
|
|
from wsproto.frame_protocol import Opcode
|
|
|
|
from mitmproxy.net import tcp
|
|
from mitmproxy.utils import bits, human, strutils
|
|
|
|
|
|
MAX_16_BIT_INT = (1 << 16)
|
|
MAX_64_BIT_INT = (1 << 64)
|
|
|
|
DEFAULT = object()
|
|
|
|
|
|
class Masker:
|
|
"""
|
|
Data sent from the server must be masked to prevent malicious clients
|
|
from sending data over the wire in predictable patterns.
|
|
|
|
Servers do not have to mask data they send to the client.
|
|
https://tools.ietf.org/html/rfc6455#section-5.3
|
|
"""
|
|
|
|
def __init__(self, key):
|
|
self.key = key
|
|
self.offset = 0
|
|
|
|
def mask(self, offset, data):
|
|
datalen = len(data)
|
|
offset_mod = offset % 4
|
|
data = int.from_bytes(data, sys.byteorder)
|
|
num_keys = (datalen + offset_mod + 3) // 4
|
|
mask = int.from_bytes((self.key * num_keys)[offset_mod:datalen +
|
|
offset_mod], sys.byteorder)
|
|
return (data ^ mask).to_bytes(datalen, sys.byteorder)
|
|
|
|
def __call__(self, data):
|
|
ret = self.mask(self.offset, data)
|
|
self.offset += len(ret)
|
|
return ret
|
|
|
|
|
|
class FrameHeader:
|
|
|
|
def __init__(
|
|
self,
|
|
opcode=Opcode.TEXT,
|
|
payload_length=0,
|
|
fin=False,
|
|
rsv1=False,
|
|
rsv2=False,
|
|
rsv3=False,
|
|
masking_key=DEFAULT,
|
|
mask=DEFAULT,
|
|
length_code=DEFAULT
|
|
):
|
|
if not 0 <= opcode < 2 ** 4:
|
|
raise ValueError("opcode must be 0-16")
|
|
self.opcode = opcode
|
|
self.payload_length = payload_length
|
|
self.fin = fin
|
|
self.rsv1 = rsv1
|
|
self.rsv2 = rsv2
|
|
self.rsv3 = rsv3
|
|
|
|
if length_code is DEFAULT:
|
|
self.length_code = self._make_length_code(self.payload_length)
|
|
else:
|
|
self.length_code = length_code
|
|
|
|
if (mask is DEFAULT and masking_key is DEFAULT) or mask == 0 or mask is False:
|
|
self.mask = False
|
|
self.masking_key = b""
|
|
elif mask is DEFAULT:
|
|
self.mask = 1
|
|
self.masking_key = masking_key
|
|
elif masking_key is DEFAULT:
|
|
self.mask = mask
|
|
self.masking_key = os.urandom(4)
|
|
else:
|
|
self.mask = mask
|
|
self.masking_key = masking_key
|
|
|
|
if self.masking_key and len(self.masking_key) != 4:
|
|
raise ValueError("Masking key must be 4 bytes.")
|
|
|
|
@classmethod
|
|
def _make_length_code(self, length):
|
|
"""
|
|
A WebSocket frame contains an initial length_code, and an optional
|
|
extended length code to represent the actual length if length code is
|
|
larger than 125
|
|
"""
|
|
if length <= 125:
|
|
return length
|
|
elif length >= 126 and length <= 65535:
|
|
return 126
|
|
else:
|
|
return 127
|
|
|
|
def __repr__(self):
|
|
vals = [
|
|
"ws frame:",
|
|
Opcode(self.opcode).name.lower()
|
|
]
|
|
flags = []
|
|
for i in ["fin", "rsv1", "rsv2", "rsv3", "mask"]:
|
|
if getattr(self, i):
|
|
flags.append(i)
|
|
if flags:
|
|
vals.extend([":", "|".join(flags)])
|
|
if self.masking_key:
|
|
vals.append(":key=%s" % repr(self.masking_key))
|
|
if self.payload_length:
|
|
vals.append(" %s" % human.pretty_size(self.payload_length))
|
|
return "".join(vals)
|
|
|
|
def __bytes__(self):
|
|
first_byte = bits.setbit(0, 7, self.fin)
|
|
first_byte = bits.setbit(first_byte, 6, self.rsv1)
|
|
first_byte = bits.setbit(first_byte, 5, self.rsv2)
|
|
first_byte = bits.setbit(first_byte, 4, self.rsv3)
|
|
first_byte = first_byte | self.opcode
|
|
|
|
second_byte = bits.setbit(self.length_code, 7, self.mask)
|
|
|
|
b = bytes([first_byte, second_byte])
|
|
|
|
if self.payload_length < 126:
|
|
pass
|
|
elif self.payload_length < MAX_16_BIT_INT:
|
|
# '!H' pack as 16 bit unsigned short
|
|
# add 2 byte extended payload length
|
|
b += struct.pack('!H', self.payload_length)
|
|
elif self.payload_length < MAX_64_BIT_INT:
|
|
# '!Q' = pack as 64 bit unsigned long long
|
|
# add 8 bytes extended payload length
|
|
b += struct.pack('!Q', self.payload_length)
|
|
else:
|
|
raise ValueError("Payload length exceeds 64bit integer")
|
|
|
|
if self.masking_key:
|
|
b += self.masking_key
|
|
return b
|
|
|
|
@classmethod
|
|
def from_file(cls, fp):
|
|
"""
|
|
read a WebSocket frame header
|
|
"""
|
|
first_byte, second_byte = fp.safe_read(2)
|
|
fin = bits.getbit(first_byte, 7)
|
|
rsv1 = bits.getbit(first_byte, 6)
|
|
rsv2 = bits.getbit(first_byte, 5)
|
|
rsv3 = bits.getbit(first_byte, 4)
|
|
opcode = first_byte & 0xF
|
|
mask_bit = bits.getbit(second_byte, 7)
|
|
length_code = second_byte & 0x7F
|
|
|
|
# payload_length > 125 indicates you need to read more bytes
|
|
# to get the actual payload length
|
|
if length_code <= 125:
|
|
payload_length = length_code
|
|
elif length_code == 126:
|
|
payload_length, = struct.unpack("!H", fp.safe_read(2))
|
|
else: # length_code == 127:
|
|
payload_length, = struct.unpack("!Q", fp.safe_read(8))
|
|
|
|
# masking key only present if mask bit set
|
|
if mask_bit == 1:
|
|
masking_key = fp.safe_read(4)
|
|
else:
|
|
masking_key = None
|
|
|
|
return cls(
|
|
fin=fin,
|
|
rsv1=rsv1,
|
|
rsv2=rsv2,
|
|
rsv3=rsv3,
|
|
opcode=opcode,
|
|
mask=mask_bit,
|
|
length_code=length_code,
|
|
payload_length=payload_length,
|
|
masking_key=masking_key,
|
|
)
|
|
|
|
def __eq__(self, other):
|
|
if isinstance(other, FrameHeader):
|
|
return bytes(self) == bytes(other)
|
|
return False
|
|
|
|
|
|
class Frame:
|
|
"""
|
|
Represents a single WebSocket frame.
|
|
Constructor takes human readable forms of the frame components.
|
|
from_bytes() reads from a file-like object to create a new Frame.
|
|
|
|
WebSocket 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
|
|
+-+-+-+-+-------+-+-------------+-------------------------------+
|
|
|F|R|R|R| opcode|M| Payload len | Extended payload length |
|
|
|I|S|S|S| (4) |A| (7) | (16/64) |
|
|
|N|V|V|V| |S| | (if payload len==126/127) |
|
|
| |1|2|3| |K| | |
|
|
+-+-+-+-+-------+-+-------------+ - - - - - - - - - - - - - - - +
|
|
| Extended payload length continued, if payload len == 127 |
|
|
+ - - - - - - - - - - - - - - - +-------------------------------+
|
|
| |Masking-key, if MASK set to 1 |
|
|
+-------------------------------+-------------------------------+
|
|
| Masking-key (continued) | Payload Data |
|
|
+-------------------------------- - - - - - - - - - - - - - - - +
|
|
: Payload Data continued ... :
|
|
+ - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - +
|
|
| Payload Data continued ... |
|
|
+---------------------------------------------------------------+
|
|
"""
|
|
|
|
def __init__(self, payload=b"", **kwargs):
|
|
self.payload = payload
|
|
kwargs["payload_length"] = kwargs.get("payload_length", len(payload))
|
|
self.header = FrameHeader(**kwargs)
|
|
|
|
@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(tcp.Reader(io.BytesIO(bytestring)))
|
|
|
|
def __repr__(self):
|
|
ret = repr(self.header)
|
|
if self.payload:
|
|
ret = ret + "\nPayload:\n" + strutils.bytes_to_escaped_str(self.payload)
|
|
return ret
|
|
|
|
def __bytes__(self):
|
|
"""
|
|
Serialize the frame to wire format. Returns a string.
|
|
"""
|
|
b = bytes(self.header)
|
|
if self.header.masking_key:
|
|
b += Masker(self.header.masking_key)(self.payload)
|
|
else:
|
|
b += self.payload
|
|
return b
|
|
|
|
@classmethod
|
|
def from_file(cls, fp):
|
|
"""
|
|
read a WebSocket frame sent by a server or client
|
|
|
|
fp is a "file like" object that could be backed by a network
|
|
stream or a disk or an in memory stream reader
|
|
"""
|
|
header = FrameHeader.from_file(fp)
|
|
payload = fp.safe_read(header.payload_length)
|
|
|
|
if header.mask == 1 and header.masking_key:
|
|
payload = Masker(header.masking_key)(payload)
|
|
|
|
frame = cls(payload)
|
|
frame.header = header
|
|
return frame
|
|
|
|
def __eq__(self, other):
|
|
if isinstance(other, Frame):
|
|
return bytes(self) == bytes(other)
|
|
return False
|