diff options
author | Maximilian Hils <git@maximilianhils.com> | 2017-01-23 21:21:44 +0100 |
---|---|---|
committer | GitHub <noreply@github.com> | 2017-01-23 21:21:44 +0100 |
commit | c46dd1e29d4faddb2292a8d54fa01dadea0eb6f9 (patch) | |
tree | cb056e1b8bee88df8f425817068f62342676e4d6 | |
parent | 6a7eeef0ee1207e92a2c57b630cfdfaf527595a5 (diff) | |
parent | 5792e2c483e622896790bb5f08ac67002e18cc14 (diff) | |
download | mitmproxy-c46dd1e29d4faddb2292a8d54fa01dadea0eb6f9.tar.gz mitmproxy-c46dd1e29d4faddb2292a8d54fa01dadea0eb6f9.tar.bz2 mitmproxy-c46dd1e29d4faddb2292a8d54fa01dadea0eb6f9.zip |
Merge pull request #1955 from Kriechi/test++
increase addon test coverage
-rw-r--r-- | mitmproxy/test/tflow.py | 56 | ||||
-rw-r--r-- | test/mitmproxy/addons/test_disable_h2c_upgrade.py | 17 | ||||
-rw-r--r-- | test/mitmproxy/addons/test_dumper.py | 20 |
3 files changed, 92 insertions, 1 deletions
diff --git a/mitmproxy/test/tflow.py b/mitmproxy/test/tflow.py index 959c9a2c..edf4d7a7 100644 --- a/mitmproxy/test/tflow.py +++ b/mitmproxy/test/tflow.py @@ -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): @@ -26,6 +28,60 @@ def ttcpflow(client_conn=True, server_conn=True, messages=True, err=None): 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): """ @type client_conn: bool | None | mitmproxy.proxy.connection.ClientConnection diff --git a/test/mitmproxy/addons/test_disable_h2c_upgrade.py b/test/mitmproxy/addons/test_disable_h2c_upgrade.py new file mode 100644 index 00000000..6cab713d --- /dev/null +++ b/test/mitmproxy/addons/test_disable_h2c_upgrade.py @@ -0,0 +1,17 @@ +from mitmproxy.addons import disable_h2c_upgrade +from mitmproxy.test import tflow + + +class TestTermLog: + def test_simple(self): + a = disable_h2c_upgrade.DisableH2CleartextUpgrade() + + f = tflow.tflow() + f.request.headers['upgrade'] = 'h2c' + f.request.headers['connection'] = 'foo' + f.request.headers['http2-settings'] = 'bar' + + a.request(f) + assert 'upgrade' not in f.request.headers + assert 'connection' not in f.request.headers + assert 'http2-settings' not in f.request.headers diff --git a/test/mitmproxy/addons/test_dumper.py b/test/mitmproxy/addons/test_dumper.py index 7ac17a7c..8fa8a22a 100644 --- a/test/mitmproxy/addons/test_dumper.py +++ b/test/mitmproxy/addons/test_dumper.py @@ -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,21 @@ 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() |