test websocket addon

This commit is contained in:
Thomas Kriechbaumer 2017-01-23 15:45:53 +01:00
parent 86174eb6ad
commit d674de298c
2 changed files with 73 additions and 1 deletions

View File

@ -1,9 +1,11 @@
from mitmproxy.test import tutils
from mitmproxy import tcp
from mitmproxy import websocket
from mitmproxy import controller
from mitmproxy import http
from mitmproxy import connections
from mitmproxy import flow
from mitmproxy.net import http as net_http
def ttcpflow(client_conn=True, server_conn=True, messages=True, err=None):
@ -25,6 +27,59 @@ def ttcpflow(client_conn=True, server_conn=True, messages=True, err=None):
f.reply = controller.DummyReply()
return f
def twebsocketflow(client_conn=True, server_conn=True, messages=True, err=None, handshake_flow=True):
if client_conn is True:
client_conn = tclient_conn()
if server_conn is True:
server_conn = tserver_conn()
if handshake_flow is True:
req = http.HTTPRequest(
"relative",
"GET",
"http",
"example.com",
"80",
"/ws",
"HTTP/1.1",
headers=net_http.Headers(
connection="upgrade",
upgrade="websocket",
sec_websocket_version="13",
sec_websocket_key="1234",
),
content=b''
)
resp = http.HTTPResponse(
"HTTP/1.1",
101,
reason=net_http.status_codes.RESPONSES.get(101),
headers=net_http.Headers(
connection='upgrade',
upgrade='websocket',
sec_websocket_accept=b'',
),
content=b'',
)
handshake_flow = http.HTTPFlow(client_conn, server_conn)
handshake_flow.request = req
handshake_flow.response = resp
f = websocket.WebSocketFlow(client_conn, server_conn, handshake_flow)
if messages is True:
messages = [
websocket.WebSocketBinaryMessage(f, True, b"hello binary"),
websocket.WebSocketTextMessage(f, False, "hello text".encode()),
]
if err is True:
err = terr()
f.messages = messages
f.error = err
f.reply = controller.DummyReply()
return f
def tflow(client_conn=True, server_conn=True, req=True, resp=None, err=None):
"""

View File

@ -157,7 +157,7 @@ def test_tcp():
d = dumper.Dumper(sio)
with taddons.context(options=dump.Options()) as ctx:
ctx.configure(d, flow_detail=3, showhost=True)
f = tflow.ttcpflow(client_conn=True, server_conn=True)
f = tflow.ttcpflow()
d.tcp_message(f)
assert "it's me" in sio.getvalue()
sio.truncate(0)
@ -165,3 +165,20 @@ def test_tcp():
f = tflow.ttcpflow(client_conn=True, err=True)
d.tcp_error(f)
assert "Error in TCP" in sio.getvalue()
def test_websocket():
sio = io.StringIO()
d = dumper.Dumper(sio)
with taddons.context(options=dump.Options()) as ctx:
ctx.configure(d, flow_detail=3, showhost=True)
f = tflow.twebsocketflow()
d.websocket_message(f)
assert "hello text" in sio.getvalue()
sio.truncate(0)
d.websocket_end(f)
assert "WebSocket connection closed by" in sio.getvalue()
f = tflow.twebsocketflow(client_conn=True, err=True)
d.websocket_error(f)
assert "Error in WebSocket" in sio.getvalue()