Grasscutters/mitmproxy
3
0
Fork 0
mirror of https://github.com/Grasscutters/mitmproxy.git synced 2024-11-23 00:01:36 +00:00
Code Issues Projects Releases Wiki Activity

[sans-io] minor test improvements

Browse Source
This commit is contained in:
Maximilian Hils 2018-05-22 13:20:35 +02:00
parent 8938aec2c0
commit a860fe4a4b
2 changed files with 218 additions and 208 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

402
test/mitmproxy/proxy2/layers/test_tls.py
View File

@ -1,5 +1,6 @@
import os import os
import ssl import ssl
import typing
import pytest import pytest
@ -81,139 +82,10 @@ class SSLTest:
) )
def test_server_no_tls(tctx: context.Context): def _test_tls_client_server(
"""Test TLS layer without TLS""" tctx: context.Context,
layer = tls.ServerTLSLayer(tctx) alpn: typing.Optional[str]
playbook = tutils.playbook(layer) ) -> typing.Tuple[tutils.playbook[tls.ClientTLSLayer], SSLTest]:
# Handshake
assert (
playbook
>> events.DataReceived(tctx.client, b"Hello World")
<< commands.Hook("next_layer", tutils.Placeholder())
>> tutils.next_layer(tutils.EchoLayer)
<< commands.SendData(tctx.client, b"hello world")
)
def test_client_tls_only(tctx: context.Context):
"""Test TLS with client only"""
layer = tls.ClientTLSLayer(tctx)
playbook = tutils.playbook(layer)
tssl = SSLTest()
# Handshake
assert playbook
assert layer._handle_event == layer.state_wait_for_clienthello
def interact():
data = tutils.Placeholder()
assert (
playbook
>> events.DataReceived(tctx.client, tssl.out.read())
<< commands.SendData(tctx.client, data)
)
tssl.inc.write(data())
try:
tssl.obj.do_handshake()
except ssl.SSLWantReadError:
return False
else:
return True
# receive ClientHello, send ServerHello
with pytest.raises(ssl.SSLWantReadError):
tssl.obj.do_handshake()
assert not interact()
# Finish Handshake
assert interact()
tssl.obj.do_handshake()
assert layer._handle_event == layer.state_process
# Echo
echo(playbook, tssl, tctx.client)
assert (
playbook
>> events.DataReceived(tctx.server, b"Hello")
<< commands.SendData(tctx.server, b"hello")
)
def echo(playbook: tutils.playbook, tssl: SSLTest, conn: context.Connection) -> None:
tssl.obj.write(b"Hello World")
data = tutils.Placeholder()
assert (
playbook
>> events.DataReceived(conn, tssl.out.read())
<< commands.Hook("next_layer", tutils.Placeholder())
>> tutils.next_layer(tutils.EchoLayer)
<< commands.SendData(conn, data)
)
tssl.inc.write(data())
assert tssl.obj.read() == b"hello world"
def test_server_tls_no_conn(tctx):
"""
The server TLS layer is initiated, but there is no active connection yet, so nothing
should be done.
"""
layer = tls.ServerTLSLayer(tctx)
playbook = tutils.playbook(layer)
tctx.server.tls = True
# We did not have a server connection before, so let's do nothing.
assert (
playbook
<< None
)
def test_server_tls(tctx):
layer = tls.ServerTLSLayer(tctx)
playbook = tutils.playbook(layer)
tctx.server.connected = True
tctx.server.address = ("example.com", 443)
tctx.server.tls = True
tssl = SSLTest(server_side=True)
# send ClientHello
data = tutils.Placeholder()
assert (
playbook
<< commands.SendData(tctx.server, data)
)
# receive ServerHello, finish client handshake
tssl.inc.write(data())
with pytest.raises(ssl.SSLWantReadError):
tssl.obj.do_handshake()
data = tutils.Placeholder()
assert (
playbook
>> events.DataReceived(tctx.server, tssl.out.read())
<< commands.SendData(tctx.server, data)
)
tssl.inc.write(data())
# finish server handshake
tssl.obj.do_handshake()
assert (
playbook
>> events.DataReceived(tctx.server, tssl.out.read())
<< None
)
assert tctx.server.tls_established
assert tctx.server.sni == b"example.com"
# Echo
echo(playbook, tssl, tctx.server)
def _test_tls_client_server(tctx, alpn):
layer = tls.ClientTLSLayer(tctx) layer = tls.ClientTLSLayer(tctx)
playbook = tutils.playbook(layer) playbook = tutils.playbook(layer)
tctx.server.tls = True tctx.server.tls = True
@ -240,85 +112,215 @@ def _test_tls_client_server(tctx, alpn):
return playbook, tssl_client return playbook, tssl_client
def test_tls_client_server_no_server_conn(tctx): def echo(playbook: tutils.playbook, tssl: SSLTest, conn: context.Connection) -> None:
""" tssl.obj.write(b"Hello World")
Here we test the scenario where a server connection is _not_ required
to establish TLS with the client. After determining this when parsing the ClientHello,
we only establish a connection with the client. The server connection may ultimately
be established when OpenConnection is called.
"""
playbook, _ = _test_tls_client_server(tctx, None)
data = tutils.Placeholder() data = tutils.Placeholder()
assert ( assert (
playbook playbook
<< commands.SendData(tctx.client, data) >> events.DataReceived(conn, tssl.out.read())
<< commands.Hook("next_layer", tutils.Placeholder())
>> tutils.next_layer(tutils.EchoLayer)
<< commands.SendData(conn, data)
) )
assert data() tssl.inc.write(data())
assert playbook.layer._handle_event == playbook.layer.state_process assert tssl.obj.read() == b"hello world"
def test_tls_client_server_alpn(tctx): class TestServerTLS:
""" def test_no_tls(self, tctx: context.Context):
Here we test the scenario where a server connection is required (e.g. because of ALPN negotation) """Test TLS layer without TLS"""
to establish TLS with the client. layer = tls.ServerTLSLayer(tctx)
""" playbook = tutils.playbook(layer)
tssl_server = SSLTest(server_side=True, alpn=["foo", "bar"])
playbook, tssl_client = _test_tls_client_server(tctx, ["qux", "foo"]) # Handshake
assert (
playbook
>> events.DataReceived(tctx.client, b"Hello World")
<< commands.Hook("next_layer", tutils.Placeholder())
>> tutils.next_layer(tutils.EchoLayer)
<< commands.SendData(tctx.client, b"hello world")
)
# We should now get instructed to open a server connection. def test_no_connection(self, tctx):
assert ( """
playbook The server TLS layer is initiated, but there is no active connection yet, so nothing
<< commands.OpenConnection(tctx.server) should be done.
) """
tctx.server.connected = True layer = tls.ServerTLSLayer(tctx)
data = tutils.Placeholder() playbook = tutils.playbook(layer)
assert ( tctx.server.tls = True
playbook
>> events.OpenConnectionReply(-1, None)
<< commands.SendData(tctx.server, data)
)
assert playbook.layer._handle_event == playbook.layer.state_wait_for_server_tls
assert playbook.layer.child_layer.tls[tctx.server]
# Establish TLS with the server... # We did not have a server connection before, so let's do nothing.
tssl_server.inc.write(data()) assert (
with pytest.raises(ssl.SSLWantReadError): playbook
<< None
)
def test_simple(self, tctx):
layer = tls.ServerTLSLayer(tctx)
playbook = tutils.playbook(layer)
tctx.server.connected = True
tctx.server.address = ("example.com", 443)
tctx.server.tls = True
tssl = SSLTest(server_side=True)
# send ClientHello
data = tutils.Placeholder()
assert (
playbook
<< commands.SendData(tctx.server, data)
)
# receive ServerHello, finish client handshake
tssl.inc.write(data())
with pytest.raises(ssl.SSLWantReadError):
tssl.obj.do_handshake()
data = tutils.Placeholder()
assert (
playbook
>> events.DataReceived(tctx.server, tssl.out.read())
<< commands.SendData(tctx.server, data)
)
tssl.inc.write(data())
# finish server handshake
tssl.obj.do_handshake()
assert (
playbook
>> events.DataReceived(tctx.server, tssl.out.read())
<< None
)
assert tctx.server.tls_established
assert tctx.server.sni == b"example.com"
# Echo
echo(playbook, tssl, tctx.server)
class TestClientTLS:
def test_simple(self, tctx: context.Context):
"""Test TLS with client only"""
layer = tls.ClientTLSLayer(tctx)
playbook = tutils.playbook(layer)
tssl = SSLTest()
# Handshake
assert playbook
assert layer._handle_event == layer.state_wait_for_clienthello
def interact():
data = tutils.Placeholder()
assert (
playbook
>> events.DataReceived(tctx.client, tssl.out.read())
<< commands.SendData(tctx.client, data)
)
tssl.inc.write(data())
try:
tssl.obj.do_handshake()
except ssl.SSLWantReadError:
return False
else:
return True
# receive ClientHello, send ServerHello
with pytest.raises(ssl.SSLWantReadError):
tssl.obj.do_handshake()
assert not interact()
# Finish Handshake
assert interact()
tssl.obj.do_handshake()
assert layer._handle_event == layer.state_process
# Echo
echo(playbook, tssl, tctx.client)
assert (
playbook
>> events.DataReceived(tctx.server, b"Hello")
<< commands.SendData(tctx.server, b"hello")
)
def test_no_server_conn_required(self, tctx):
"""
Here we test the scenario where a server connection is _not_ required
to establish TLS with the client. After determining this when parsing the ClientHello,
we only establish a connection with the client. The server connection may ultimately
be established when OpenConnection is called.
"""
playbook, _ = _test_tls_client_server(tctx, None)
data = tutils.Placeholder()
assert (
playbook
<< commands.SendData(tctx.client, data)
)
assert data()
assert playbook.layer._handle_event == playbook.layer.state_process
def test_alpn(self, tctx):
"""
Here we test the scenario where a server connection is required (e.g. because of ALPN negotation)
to establish TLS with the client.
"""
tssl_server = SSLTest(server_side=True, alpn=["foo", "bar"])
playbook, tssl_client = _test_tls_client_server(tctx, ["qux", "foo"])
# We should now get instructed to open a server connection.
assert (
playbook
<< commands.OpenConnection(tctx.server)
)
tctx.server.connected = True
data = tutils.Placeholder()
assert (
playbook
>> events.OpenConnectionReply(-1, None)
<< commands.SendData(tctx.server, data)
)
assert playbook.layer._handle_event == playbook.layer.state_wait_for_server_tls
assert playbook.layer.child_layer.tls[tctx.server]
# Establish TLS with the server...
tssl_server.inc.write(data())
with pytest.raises(ssl.SSLWantReadError):
tssl_server.obj.do_handshake()
data = tutils.Placeholder()
assert (
playbook
>> events.DataReceived(tctx.server, tssl_server.out.read())
<< commands.SendData(tctx.server, data)
)
tssl_server.inc.write(data())
tssl_server.obj.do_handshake() tssl_server.obj.do_handshake()
data = tutils.Placeholder() data = tutils.Placeholder()
assert ( assert (
playbook playbook
>> events.DataReceived(tctx.server, tssl_server.out.read()) >> events.DataReceived(tctx.server, tssl_server.out.read())
<< commands.SendData(tctx.server, data) << commands.SendData(tctx.client, data)
) )
tssl_server.inc.write(data())
tssl_server.obj.do_handshake()
data = tutils.Placeholder()
assert (
playbook
>> events.DataReceived(tctx.server, tssl_server.out.read())
<< commands.SendData(tctx.client, data)
)
assert playbook.layer._handle_event == playbook.layer.state_process assert playbook.layer._handle_event == playbook.layer.state_process
assert tctx.server.tls_established assert tctx.server.tls_established
# Server TLS is established, we can now reply to the client handshake... # Server TLS is established, we can now reply to the client handshake...
tssl_client.inc.write(data()) tssl_client.inc.write(data())
with pytest.raises(ssl.SSLWantReadError): with pytest.raises(ssl.SSLWantReadError):
tssl_client.obj.do_handshake()
data = tutils.Placeholder()
assert (
playbook
>> events.DataReceived(tctx.client, tssl_client.out.read())
<< commands.SendData(tctx.client, data)
)
tssl_client.inc.write(data())
tssl_client.obj.do_handshake() tssl_client.obj.do_handshake()
data = tutils.Placeholder()
assert (
playbook
>> events.DataReceived(tctx.client, tssl_client.out.read())
<< commands.SendData(tctx.client, data)
)
tssl_client.inc.write(data())
tssl_client.obj.do_handshake()
# Both handshakes completed! # Both handshakes completed!
assert tctx.client.tls_established assert tctx.client.tls_established
assert tctx.server.tls_established assert tctx.server.tls_established
assert tssl_client.obj.selected_alpn_protocol() == "foo" assert tssl_client.obj.selected_alpn_protocol() == "foo"
assert tssl_server.obj.selected_alpn_protocol() == "foo" assert tssl_server.obj.selected_alpn_protocol() == "foo"

24
test/mitmproxy/proxy2/tutils.py
View File

@ -30,9 +30,9 @@ def _eq(
x, y = a[k], b[k] x, y = a[k], b[k]
# if there's a placeholder, make it x. # if there's a placeholder, make it x.
if isinstance(y, Placeholder): if isinstance(y, _Placeholder):
x, y = y, x x, y = y, x
if isinstance(x, Placeholder): if isinstance(x, _Placeholder):
if x.obj is None: if x.obj is None:
x.obj = y x.obj = y
x = x.obj x = x.obj
@ -56,7 +56,11 @@ def eq(
return _eq(a, b) return _eq(a, b)
class playbook: T = typing.TypeVar('T', bound=Layer)
# noinspection PyPep8Naming
class playbook(typing.Generic[T]):
""" """
Assert that a layer emits the expected commands in reaction to a given sequence of events. Assert that a layer emits the expected commands in reaction to a given sequence of events.
For example, the following code asserts that the TCP layer emits an OpenConnection command For example, the following code asserts that the TCP layer emits an OpenConnection command
@ -76,7 +80,7 @@ class playbook:
x2 = list(t.handle_event(events.OpenConnectionReply(x1[-1]))) x2 = list(t.handle_event(events.OpenConnectionReply(x1[-1])))
assert x2 == [] assert x2 == []
""" """
layer: Layer layer: T
"""The base layer""" """The base layer"""
expected: TPlaybook expected: TPlaybook
"""expected command/event sequence""" """expected command/event sequence"""
@ -89,9 +93,9 @@ class playbook:
def __init__( def __init__(
self, self,
layer: Layer, layer: T,
expected: typing.Optional[TPlaybook]=None, expected: typing.Optional[TPlaybook] = None,
ignore_log: bool=True ignore_log: bool = True
): ):
if expected is None: if expected is None:
expected = [ expected = [
@ -176,7 +180,7 @@ class playbook:
return copy.deepcopy(self) return copy.deepcopy(self)
class Placeholder: class _Placeholder:
""" """
Placeholder value in playbooks, so that objects (flows in particular) can be referenced before Placeholder value in playbooks, so that objects (flows in particular) can be referenced before
they are known. Example: they are known. Example:
@ -205,6 +209,10 @@ class Placeholder:
return f"Placeholder:{str(self.obj)}" return f"Placeholder:{str(self.obj)}"
def Placeholder() -> typing.Any:
return _Placeholder()
class EchoLayer(Layer): class EchoLayer(Layer):
"""Echo layer that sends all data back to the client in lowercase.""" """Echo layer that sends all data back to the client in lowercase."""