kylefang/mitmproxy
1
0
Fork 0
mirror of https://github.com/zhigang1992/mitmproxy.git synced 2026-01-13 09:19:49 +08:00
Code Issues Packages Projects Releases Wiki Activity

websockets, tcp, version: coverage++

Browse Source
This commit is contained in:
Thomas Kriechbaumer
2017-02-08 15:00:31 +01:00
parent 2134b5b06a
commit 1847cf175c
10 changed files with 171 additions and 66 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
mitmproxy/addons/dumper.py
View File

@@ -238,7 +238,7 @@ class Dumper:
def websocket_message(self, f):
if self.match(f):
message = f.messages[-1]
self.echo(message.info)
self.echo(f.message_info(message))
if self.flow_detail >= 3:
self._echo_message(message)

9
mitmproxy/proxy/protocol/websocket.py
View File

@@ -8,7 +8,7 @@ from mitmproxy import flow
from mitmproxy.proxy.protocol import base
from mitmproxy.net import tcp
from mitmproxy.net import websockets
from mitmproxy.websocket import WebSocketFlow, WebSocketBinaryMessage, WebSocketTextMessage
from mitmproxy.websocket import WebSocketFlow, WebSocketMessage
class WebSocketLayer(base.Layer):
@@ -65,12 +65,7 @@ class WebSocketLayer(base.Layer):
compressed_message = fb[0].header.rsv1
fb.clear()
if message_type == websockets.OPCODE.TEXT:
t = WebSocketTextMessage
else:
t = WebSocketBinaryMessage
websocket_message = t(self.flow, not is_server, payload)
websocket_message = WebSocketMessage(message_type, not is_server, payload)
length = len(websocket_message.content)
self.flow.messages.append(websocket_message)
self.channel.ask("websocket_message", self.flow)

6
mitmproxy/tcp.py
View File

@@ -9,8 +9,8 @@ from mitmproxy.types import serializable
class TCPMessage(serializable.Serializable):
def __init__(self, from_client, content, timestamp=None):
self.content = content
self.from_client = from_client
self.content = content
self.timestamp = timestamp or time.time()
@classmethod
@@ -21,9 +21,7 @@ class TCPMessage(serializable.Serializable):
return self.from_client, self.content, self.timestamp
def set_state(self, state):
self.from_client = state.pop("from_client")
self.content = state.pop("content")
self.timestamp = state.pop("timestamp")
self.from_client, self.content, self.timestamp = state
def __repr__(self):
return "{direction} {content}".format(

5
mitmproxy/test/tflow.py
View File

@@ -1,3 +1,4 @@
from mitmproxy.net import websockets
from mitmproxy.test import tutils
from mitmproxy import tcp
from mitmproxy import websocket
@@ -70,8 +71,8 @@ def twebsocketflow(client_conn=True, server_conn=True, messages=True, err=None,
if messages is True:
messages = [
websocket.WebSocketBinaryMessage(f, True, b"hello binary"),
websocket.WebSocketTextMessage(f, False, "hello text".encode()),
websocket.WebSocketMessage(websockets.OPCODE.BINARY, True, b"hello binary"),
websocket.WebSocketMessage(websockets.OPCODE.TEXT, False, "hello text".encode()),
]
if err is True:
err = terr()

2
mitmproxy/version.py
View File

@@ -3,5 +3,5 @@ VERSION = ".".join(str(i) for i in IVERSION)
PATHOD = "pathod " + VERSION
MITMPROXY = "mitmproxy " + VERSION
if __name__ == "__main__":
if __name__ == "__main__": # pragma: no cover
print(VERSION)

100
mitmproxy/websocket.py
View File

@@ -1,63 +1,38 @@
import time
from typing import List
from typing import List, Optional
from mitmproxy import flow
from mitmproxy.http import HTTPFlow
from mitmproxy.net import websockets
from mitmproxy.utils import strutils
from mitmproxy.types import serializable
from mitmproxy.utils import strutils
class WebSocketMessage(serializable.Serializable):
def __init__(self, flow, from_client, content, timestamp=None):
self.flow = flow
self.content = content
def __init__(self, type: int, from_client: bool, content: bytes, timestamp: Optional[int]=None):
self.type = type
self.from_client = from_client
self.timestamp = timestamp or time.time()
self.content = content
self.timestamp = timestamp or time.time() # type: int
@classmethod
def from_state(cls, state):
return cls(*state)
def get_state(self):
return self.from_client, self.content, self.timestamp
return self.type, self.from_client, self.content, self.timestamp
def set_state(self, state):
self.from_client = state.pop("from_client")
self.content = state.pop("content")
self.timestamp = state.pop("timestamp")
@property
def info(self):
return "{client} {direction} WebSocket {type} message {direction} {server}{endpoint}".format(
type=self.type,
client=repr(self.flow.client_conn.address),
server=repr(self.flow.server_conn.address),
direction="->" if self.from_client else "<-",
endpoint=self.flow.handshake_flow.request.path,
)
class WebSocketBinaryMessage(WebSocketMessage):
type = 'binary'
self.type, self.from_client, self.content, self.timestamp = state
def __repr__(self):
return "binary message: {}".format(strutils.bytes_to_escaped_str(self.content))
class WebSocketTextMessage(WebSocketMessage):
type = 'text'
def __repr__(self):
return "text message: {}".format(repr(self.content))
if self.type == websockets.OPCODE.TEXT:
return "text message: {}".format(repr(self.content))
else:
return "binary message: {}".format(strutils.bytes_to_escaped_str(self.content))
class WebSocketFlow(flow.Flow):
"""
A WebsocketFlow is a simplified representation of a Websocket session.
"""
@@ -70,18 +45,55 @@ class WebSocketFlow(flow.Flow):
self.close_message = '(message missing)'
self.close_reason = 'unknown status code'
self.handshake_flow = handshake_flow
self.client_key = websockets.get_client_key(self.handshake_flow.request.headers)
self.client_protocol = websockets.get_protocol(self.handshake_flow.request.headers)
self.client_extensions = websockets.get_extensions(self.handshake_flow.request.headers)
self.server_accept = websockets.get_server_accept(self.handshake_flow.response.headers)
self.server_protocol = websockets.get_protocol(self.handshake_flow.response.headers)
self.server_extensions = websockets.get_extensions(self.handshake_flow.response.headers)
_stateobject_attributes = flow.Flow._stateobject_attributes.copy()
_stateobject_attributes.update(
messages=List[WebSocketMessage],
close_sender=str,
close_code=str,
close_message=str,
close_reason=str,
handshake_flow=HTTPFlow,
)
@classmethod
def from_state(cls, state):
f = cls(None, None, None)
f.set_state(state)
return f
def __repr__(self):
return "WebSocketFlow ({} messages)".format(len(self.messages))
return "<WebSocketFlow ({} messages)>".format(len(self.messages))
@property
def client_key(self):
return websockets.get_client_key(self.handshake_flow.request.headers)
@property
def client_protocol(self):
return websockets.get_protocol(self.handshake_flow.request.headers)
@property
def client_extensions(self):
return websockets.get_extensions(self.handshake_flow.request.headers)
@property
def server_accept(self):
return websockets.get_server_accept(self.handshake_flow.response.headers)
@property
def server_protocol(self):
return websockets.get_protocol(self.handshake_flow.response.headers)
@property
def server_extensions(self):
return websockets.get_extensions(self.handshake_flow.response.headers)
def message_info(self, message: WebSocketMessage) -> str:
return "{client} {direction} WebSocket {type} message {direction} {server}{endpoint}".format(
type=message.type,
client=repr(self.client_conn.address),
server=repr(self.server_conn.address),
direction="->" if message.from_client else "<-",
endpoint=self.handshake_flow.request.path,
)

1
test/mitmproxy/addons/test_stickycookie.py
View File

@@ -39,7 +39,6 @@ class TestStickyCookie:
assert "cookie" not in f.request.headers
f = f.copy()
f.reply.acked = False
sc.request(f)
assert f.request.headers["cookie"] == "foo=bar"

11
test/mitmproxy/protocol/test_websocket.py
View File

@@ -179,16 +179,15 @@ class TestSimple(_WebSocketTest):
assert isinstance(self.master.state.flows[1], WebSocketFlow)
assert len(self.master.state.flows[1].messages) == 5
assert self.master.state.flows[1].messages[0].content == b'server-foobar'
assert self.master.state.flows[1].messages[0].type == 'text'
assert self.master.state.flows[1].messages[0].type == websockets.OPCODE.TEXT
assert self.master.state.flows[1].messages[1].content == b'client-foobar'
assert self.master.state.flows[1].messages[1].type == 'text'
assert self.master.state.flows[1].messages[1].type == websockets.OPCODE.TEXT
assert self.master.state.flows[1].messages[2].content == b'client-foobar'
assert self.master.state.flows[1].messages[2].type == 'text'
assert self.master.state.flows[1].messages[2].type == websockets.OPCODE.TEXT
assert self.master.state.flows[1].messages[3].content == b'\xde\xad\xbe\xef'
assert self.master.state.flows[1].messages[3].type == 'binary'
assert self.master.state.flows[1].messages[3].type == websockets.OPCODE.BINARY
assert self.master.state.flows[1].messages[4].content == b'\xde\xad\xbe\xef'
assert self.master.state.flows[1].messages[4].type == 'binary'
assert [m.info for m in self.master.state.flows[1].messages]
assert self.master.state.flows[1].messages[4].type == websockets.OPCODE.BINARY
class TestSimpleTLS(_WebSocketTest):

97
test/mitmproxy/test_flow.py
View File

@@ -10,6 +10,7 @@ from mitmproxy.exceptions import FlowReadException, Kill
from mitmproxy import flow
from mitmproxy import http
from mitmproxy import connections
from mitmproxy import tcp
from mitmproxy.proxy import ProxyConfig
from mitmproxy.proxy.server import DummyServer
from mitmproxy import master
@@ -156,8 +157,99 @@ class TestHTTPFlow:
assert f.response.raw_content == b"abarb"
class TestWebSocketFlow:
def test_copy(self):
f = tflow.twebsocketflow()
f.get_state()
f2 = f.copy()
a = f.get_state()
b = f2.get_state()
del a["id"]
del b["id"]
del a["handshake_flow"]["id"]
del b["handshake_flow"]["id"]
assert a == b
assert not f == f2
assert f is not f2
assert f.client_key == f2.client_key
assert f.client_protocol == f2.client_protocol
assert f.client_extensions == f2.client_extensions
assert f.server_accept == f2.server_accept
assert f.server_protocol == f2.server_protocol
assert f.server_extensions == f2.server_extensions
assert f.messages is not f2.messages
assert f.handshake_flow is not f2.handshake_flow
for m in f.messages:
m2 = m.copy()
m2.set_state(m2.get_state())
assert m is not m2
assert m.get_state() == m2.get_state()
f = tflow.twebsocketflow(err=True)
f2 = f.copy()
assert f is not f2
assert f.handshake_flow is not f2.handshake_flow
assert f.error.get_state() == f2.error.get_state()
assert f.error is not f2.error
def test_match(self):
f = tflow.twebsocketflow()
assert not flowfilter.match("~b nonexistent", f)
assert flowfilter.match(None, f)
assert not flowfilter.match("~b nonexistent", f)
f = tflow.twebsocketflow(err=True)
assert flowfilter.match("~e", f)
with pytest.raises(ValueError):
flowfilter.match("~", f)
def test_repr(self):
f = tflow.twebsocketflow()
assert 'WebSocketFlow' in repr(f)
assert 'binary message: ' in repr(f.messages[0])
assert 'text message: ' in repr(f.messages[1])
class TestTCPFlow:
def test_copy(self):
f = tflow.ttcpflow()
f.get_state()
f2 = f.copy()
a = f.get_state()
b = f2.get_state()
del a["id"]
del b["id"]
assert a == b
assert not f == f2
assert f is not f2
assert f.messages is not f2.messages
for m in f.messages:
assert m.get_state()
m2 = m.copy()
assert not m == m2
assert m is not m2
a = m.get_state()
b = m2.get_state()
assert a == b
m = tcp.TCPMessage(False, 'foo')
m.set_state(f.messages[0].get_state())
assert m.timestamp == f.messages[0].timestamp
f = tflow.ttcpflow(err=True)
f2 = f.copy()
assert f is not f2
assert f.error.get_state() == f2.error.get_state()
assert f.error is not f2.error
def test_match(self):
f = tflow.ttcpflow()
assert not flowfilter.match("~b nonexistent", f)
@@ -170,6 +262,11 @@ class TestTCPFlow:
with pytest.raises(ValueError):
flowfilter.match("~", f)
def test_repr(self):
f = tflow.ttcpflow()
assert 'TCPFlow' in repr(f)
assert '-> ' in repr(f.messages[0])
class TestSerialize:

4
tox.ini
View File

@@ -27,6 +27,10 @@ commands =
--full-cov=mitmproxy/io.py \
--full-cov=mitmproxy/log.py \
--full-cov=mitmproxy/options.py \
--full-cov=mitmproxy/stateobject.py \
--full-cov=mitmproxy/version.py \
--full-cov=mitmproxy/tcp.py \
--full-cov=mitmproxy/websocket.py \
--full-cov=pathod/ --no-full-cov=pathod/pathoc.py --no-full-cov=pathod/pathod.py --no-full-cov=pathod/test.py --no-full-cov=pathod/protocols/http2.py \
{posargs}
{env:CI_COMMANDS:python -c ""}