Grasscutters/mitmproxy
3
0
Fork 0
mirror of https://github.com/Grasscutters/mitmproxy.git synced 2024-11-25 18:03:50 +00:00
Code Issues Projects Releases Wiki Activity

Merge pull request #2350 from Kriechi/improve-tests

Browse Source 
fix various leaking fds/sockets
This commit is contained in:
Thomas Kriechbaumer 2017-05-24 14:53:45 +02:00 committed by GitHub
commit 8b31f5c324
26 changed files with 250 additions and 242 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
mitmproxy/addons/view.py
View File

@ -339,11 +339,12 @@ class View(collections.Sequence):
""" """
Load flows into the view, without processing them with addons. Load flows into the view, without processing them with addons.
""" """
for i in io.FlowReader(open(path, "rb")).stream(): with open(path, "rb") as f:
# Do this to get a new ID, so we can load the same file N times and for i in io.FlowReader(f).stream():
# get new flows each time. It would be more efficient to just have a # Do this to get a new ID, so we can load the same file N times and
# .newid() method or something. # get new flows each time. It would be more efficient to just have a
self.add([i.copy()]) # .newid() method or something.
self.add([i.copy()])
@command.command("view.go") @command.command("view.go")
def go(self, dst: int) -> None: def go(self, dst: int) -> None:

4
mitmproxy/net/socks.py
View File

@ -82,12 +82,12 @@ class ClientGreeting:
client_greeting = cls(ver, []) client_greeting = cls(ver, [])
if fail_early: if fail_early:
client_greeting.assert_socks5() client_greeting.assert_socks5()
client_greeting.methods.fromstring(f.safe_read(nmethods)) client_greeting.methods.frombytes(f.safe_read(nmethods))
return client_greeting return client_greeting
def to_file(self, f): def to_file(self, f):
f.write(struct.pack("!BB", self.ver, len(self.methods))) f.write(struct.pack("!BB", self.ver, len(self.methods)))
f.write(self.methods.tostring()) f.write(self.methods.tobytes())
class ServerGreeting: class ServerGreeting:

8
mitmproxy/net/tcp.py
View File

@ -569,7 +569,9 @@ class TCPClient(_Connection):
# Make sure to close the real socket, not the SSL proxy. # Make sure to close the real socket, not the SSL proxy.
# OpenSSL is really good at screwing up, i.e. when trying to recv from a failed connection, # OpenSSL is really good at screwing up, i.e. when trying to recv from a failed connection,
# it tries to renegotiate... # it tries to renegotiate...
if isinstance(self.connection, SSL.Connection): if not self.connection:
return
elif isinstance(self.connection, SSL.Connection):
close_socket(self.connection._socket) close_socket(self.connection._socket)
else: else:
close_socket(self.connection) close_socket(self.connection)
@ -674,6 +676,8 @@ class TCPClient(_Connection):
sock.setsockopt(socket.SOL_IP, socket.IP_TRANSPARENT, 1) # pragma: windows no cover pragma: osx no cover sock.setsockopt(socket.SOL_IP, socket.IP_TRANSPARENT, 1) # pragma: windows no cover pragma: osx no cover
except Exception as e: except Exception as e:
# socket.IP_TRANSPARENT might not be available on every OS and Python version # socket.IP_TRANSPARENT might not be available on every OS and Python version
if sock is not None:
sock.close()
raise exceptions.TcpException( raise exceptions.TcpException(
"Failed to spoof the source address: " + str(e) "Failed to spoof the source address: " + str(e)
) )
@ -864,6 +868,8 @@ class TCPServer:
self.socket.setsockopt(IPPROTO_IPV6, socket.IPV6_V6ONLY, 0) self.socket.setsockopt(IPPROTO_IPV6, socket.IPV6_V6ONLY, 0)
self.socket.bind(self.address) self.socket.bind(self.address)
except socket.error: except socket.error:
if self.socket:
self.socket.close()
self.socket = None self.socket = None
if not self.socket: if not self.socket:

4
mitmproxy/proxy/protocol/http_replay.py
View File

@ -83,7 +83,11 @@ class RequestReplayThread(basethread.BaseThread):
server.wfile.write(http1.assemble_request(r)) server.wfile.write(http1.assemble_request(r))
server.wfile.flush() server.wfile.flush()
if self.f.server_conn:
self.f.server_conn.close()
self.f.server_conn = server self.f.server_conn = server
self.f.response = http.HTTPResponse.wrap( self.f.response = http.HTTPResponse.wrap(
http1.read_response( http1.read_response(
server.rfile, server.rfile,

2
mitmproxy/proxy/server.py
View File

@ -48,6 +48,8 @@ class ProxyServer(tcp.TCPServer):
if config.options.mode == "transparent": if config.options.mode == "transparent":
platform.init_transparent_mode() platform.init_transparent_mode()
except Exception as e: except Exception as e:
if self.socket:
self.socket.close()
raise exceptions.ServerException( raise exceptions.ServerException(
'Error starting proxy server: ' + repr(e) 'Error starting proxy server: ' + repr(e)
) from e ) from e

26
mitmproxy/tools/main.py
View File

@ -1,6 +1,8 @@
from __future__ import print_function # this is here for the version check to work on Python 2. from __future__ import print_function # this is here for the version check to work on Python 2.
import sys import sys
# This must be at the very top, before importing anything else that might break!
# Keep all other imports below with the 'noqa' magic comment.
if sys.version_info < (3, 5): if sys.version_info < (3, 5):
print("#" * 49, file=sys.stderr) print("#" * 49, file=sys.stderr)
print("# mitmproxy only supports Python 3.5 and above! #", file=sys.stderr) print("# mitmproxy only supports Python 3.5 and above! #", file=sys.stderr)
@ -13,8 +15,7 @@ from mitmproxy.tools import cmdline # noqa
from mitmproxy import exceptions # noqa from mitmproxy import exceptions # noqa
from mitmproxy import options # noqa from mitmproxy import options # noqa
from mitmproxy import optmanager # noqa from mitmproxy import optmanager # noqa
from mitmproxy.proxy import config # noqa from mitmproxy import proxy # noqa
from mitmproxy.proxy import server # noqa
from mitmproxy.utils import version_check # noqa from mitmproxy.utils import version_check # noqa
from mitmproxy.utils import debug # noqa from mitmproxy.utils import debug # noqa
@ -49,15 +50,7 @@ def process_options(parser, opts, args):
adict[n] = getattr(args, n) adict[n] = getattr(args, n)
opts.merge(adict) opts.merge(adict)
pconf = config.ProxyConfig(opts) return proxy.config.ProxyConfig(opts)
if opts.server:
try:
return server.ProxyServer(pconf)
except exceptions.ServerException as v:
print(str(v), file=sys.stderr)
sys.exit(1)
else:
return server.DummyServer(pconf)
def run(MasterKlass, args, extra=None): # pragma: no cover def run(MasterKlass, args, extra=None): # pragma: no cover
@ -74,7 +67,16 @@ def run(MasterKlass, args, extra=None): # pragma: no cover
master = None master = None
try: try:
unknown = optmanager.load_paths(opts, args.conf) unknown = optmanager.load_paths(opts, args.conf)
server = process_options(parser, opts, args) pconf = process_options(parser, opts, args)
if pconf.options.server:
try:
server = proxy.server.ProxyServer(pconf)
except exceptions.ServerException as v:
print(str(v), file=sys.stderr)
sys.exit(1)
else:
server = proxy.server.DummyServer(pconf)
master = MasterKlass(opts, server) master = MasterKlass(opts, server)
master.addons.trigger("configure", opts.keys()) master.addons.trigger("configure", opts.keys())
master.addons.trigger("tick") master.addons.trigger("tick")

4
pathod/language/generators.py
View File

@ -91,3 +91,7 @@ class FileGenerator:
def __repr__(self): def __repr__(self):
return "<%s" % self.path return "<%s" % self.path
def close(self):
self.map.close()
self.fp.close()

29
setup.py
View File

@ -61,8 +61,9 @@ setup(
# It is not considered best practice to use install_requires to pin dependencies to specific versions. # It is not considered best practice to use install_requires to pin dependencies to specific versions.
install_requires=[ install_requires=[
"blinker>=1.4, <1.5", "blinker>=1.4, <1.5",
"click>=6.2, <7", "brotlipy>=0.5.1, <0.7",
"certifi>=2015.11.20.1", # no semver here - this should always be on the last release! "certifi>=2015.11.20.1", # no semver here - this should always be on the last release!
"click>=6.2, <7",
"construct>=2.8, <2.9", "construct>=2.8, <2.9",
"cryptography>=1.4, <1.9", "cryptography>=1.4, <1.9",
"cssutils>=1.0.1, <1.1", "cssutils>=1.0.1, <1.1",
@ -79,37 +80,29 @@ setup(
"pyperclip>=1.5.22, <1.6", "pyperclip>=1.5.22, <1.6",
"requests>=2.9.1, <3", "requests>=2.9.1, <3",
"ruamel.yaml>=0.13.2, <0.15", "ruamel.yaml>=0.13.2, <0.15",
"sortedcontainers>=1.5.4, <1.6",
"tornado>=4.3, <4.6", "tornado>=4.3, <4.6",
"urwid>=1.3.1, <1.4", "urwid>=1.3.1, <1.4",
"brotlipy>=0.5.1, <0.7",
"sortedcontainers>=1.5.4, <1.6",
# transitive from cryptography, we just blacklist here.
# https://github.com/pypa/setuptools/issues/861
"setuptools>=11.3, !=29.0.0",
], ],
extras_require={ extras_require={
':sys_platform == "win32"': [ ':sys_platform == "win32"': [
"pydivert>=2.0.3, <2.1", "pydivert>=2.0.3, <2.1",
], ],
':sys_platform != "win32"': [
],
'dev': [ 'dev': [
"Flask>=0.10.1, <0.13",
"flake8>=3.2.1, <3.4", "flake8>=3.2.1, <3.4",
"Flask>=0.10.1, <0.13",
"mypy>=0.501, <0.502", "mypy>=0.501, <0.502",
"rstcheck>=2.2, <4.0",
"tox>=2.3, <3",
"pytest>=3, <3.1",
"pytest-cov>=2.2.1, <3", "pytest-cov>=2.2.1, <3",
"pytest-faulthandler>=1.3.0, <2",
"pytest-timeout>=1.0.0, <2", "pytest-timeout>=1.0.0, <2",
"pytest-xdist>=1.14, <2", "pytest-xdist>=1.14, <2",
"pytest-faulthandler>=1.3.0, <2", "pytest>=3.1, <4",
"sphinx>=1.3.5, <1.7", "rstcheck>=2.2, <4.0",
"sphinx-autobuild>=0.5.2, <0.7",
"sphinxcontrib-documentedlist>=0.5.0, <0.7",
"sphinx_rtd_theme>=0.1.9, <0.3", "sphinx_rtd_theme>=0.1.9, <0.3",
], "sphinx-autobuild>=0.5.2, <0.7",
'contentviews': [ "sphinx>=1.3.5, <1.7",
"sphinxcontrib-documentedlist>=0.5.0, <0.7",
"tox>=2.3, <3",
], ],
'examples': [ 'examples': [
"beautifulsoup4>=4.4.1, <4.7", "beautifulsoup4>=4.4.1, <4.7",

20
test/conftest.py
View File

@ -1,8 +1,6 @@
import os import os
import pytest import pytest
import OpenSSL import OpenSSL
import functools
from contextlib import contextmanager
import mitmproxy.net.tcp import mitmproxy.net.tcp
@ -32,21 +30,3 @@ skip_appveyor = pytest.mark.skipif(
def disable_alpn(monkeypatch): def disable_alpn(monkeypatch):
monkeypatch.setattr(mitmproxy.net.tcp, 'HAS_ALPN', False) monkeypatch.setattr(mitmproxy.net.tcp, 'HAS_ALPN', False)
monkeypatch.setattr(OpenSSL.SSL._lib, 'Cryptography_HAS_ALPN', False) monkeypatch.setattr(OpenSSL.SSL._lib, 'Cryptography_HAS_ALPN', False)
################################################################################
# TODO: remove this wrapper when pytest 3.1.0 is released
original_pytest_raises = pytest.raises
@contextmanager
@functools.wraps(original_pytest_raises)
def raises(exc, *args, **kwargs):
with original_pytest_raises(exc, *args, **kwargs) as exc_info:
yield
if 'match' in kwargs:
assert exc_info.match(kwargs['match'])
pytest.raises = raises
################################################################################

7
test/mitmproxy/addons/test_clientplayback.py
View File

@ -10,9 +10,10 @@ from mitmproxy.test import taddons
def tdump(path, flows): def tdump(path, flows):
w = io.FlowWriter(open(path, "wb")) with open(path, "wb") as f:
for i in flows: w = io.FlowWriter(f)
w.add(i) for i in flows:
w.add(i)
class MockThread(): class MockThread():

5
test/mitmproxy/addons/test_save.py
View File

@ -26,8 +26,9 @@ def test_configure(tmpdir):
def rd(p): def rd(p):
x = io.FlowReader(open(p, "rb")) with open(p, "rb") as f:
return list(x.stream()) x = io.FlowReader(f)
return list(x.stream())
def test_tcp(tmpdir): def test_tcp(tmpdir):

7
test/mitmproxy/addons/test_serverplayback.py
View File

@ -11,9 +11,10 @@ from mitmproxy import io
def tdump(path, flows): def tdump(path, flows):
w = io.FlowWriter(open(path, "wb")) with open(path, "wb") as f:
for i in flows: w = io.FlowWriter(f)
w.add(i) for i in flows:
w.add(i)
def test_load_file(tmpdir): def test_load_file(tmpdir):

7
test/mitmproxy/addons/test_view.py
View File

@ -132,9 +132,10 @@ def test_filter():
def tdump(path, flows): def tdump(path, flows):
w = io.FlowWriter(open(path, "wb")) with open(path, "wb") as f:
for i in flows: w = io.FlowWriter(f)
w.add(i) for i in flows:
w.add(i)
def test_create(): def test_create():

4
test/mitmproxy/contentviews/test_protobuf.py
View File

@ -17,7 +17,9 @@ def test_view_protobuf_request():
m.configure_mock(**attrs) m.configure_mock(**attrs)
n.return_value = m n.return_value = m
content_type, output = v(open(p, "rb").read()) with open(p, "rb") as f:
data = f.read()
content_type, output = v(data)
assert content_type == "Protobuf" assert content_type == "Protobuf"
assert output[0] == [('text', b'1: "3bbc333c-e61c-433b-819a-0b9a8cc103b8"')] assert output[0] == [('text', b'1: "3bbc333c-e61c-433b-819a-0b9a8cc103b8"')]

10
test/mitmproxy/net/test_tcp.py
View File

@ -34,7 +34,7 @@ class ClientCipherListHandler(tcp.BaseHandler):
sni = None sni = None
def handle(self): def handle(self):
self.wfile.write("%s" % self.connection.get_cipher_list()) self.wfile.write(str(self.connection.get_cipher_list()).encode())
self.wfile.flush() self.wfile.flush()
@ -391,14 +391,15 @@ class TestSNI(tservers.ServerTestBase):
class TestServerCipherList(tservers.ServerTestBase): class TestServerCipherList(tservers.ServerTestBase):
handler = ClientCipherListHandler handler = ClientCipherListHandler
ssl = dict( ssl = dict(
cipher_list='AES256-GCM-SHA384' cipher_list=b'AES256-GCM-SHA384'
) )
def test_echo(self): def test_echo(self):
c = tcp.TCPClient(("127.0.0.1", self.port)) c = tcp.TCPClient(("127.0.0.1", self.port))
with c.connect(): with c.connect():
c.convert_to_ssl(sni="foo.com") c.convert_to_ssl(sni="foo.com")
assert c.rfile.readline() == b"['AES256-GCM-SHA384']" expected = b"['AES256-GCM-SHA384']"
assert c.rfile.read(len(expected) + 2) == expected
class TestServerCurrentCipher(tservers.ServerTestBase): class TestServerCurrentCipher(tservers.ServerTestBase):
@ -424,7 +425,7 @@ class TestServerCurrentCipher(tservers.ServerTestBase):
class TestServerCipherListError(tservers.ServerTestBase): class TestServerCipherListError(tservers.ServerTestBase):
handler = ClientCipherListHandler handler = ClientCipherListHandler
ssl = dict( ssl = dict(
cipher_list='bogus' cipher_list=b'bogus'
) )
def test_echo(self): def test_echo(self):
@ -632,6 +633,7 @@ class TestTCPServer:
with s.handler_counter: with s.handler_counter:
with pytest.raises(exceptions.Timeout): with pytest.raises(exceptions.Timeout):
s.wait_for_silence() s.wait_for_silence()
s.shutdown()
class TestFileLike: class TestFileLike:

14
test/mitmproxy/net/tservers.py
View File

@ -16,9 +16,6 @@ class _ServerThread(threading.Thread):
def run(self): def run(self):
self.server.serve_forever() self.server.serve_forever()
def shutdown(self):
self.server.shutdown()
class _TServer(tcp.TCPServer): class _TServer(tcp.TCPServer):
@ -54,9 +51,9 @@ class _TServer(tcp.TCPServer):
raw_key = self.ssl.get( raw_key = self.ssl.get(
"key", "key",
tutils.test_data.path("mitmproxy/net/data/server.key")) tutils.test_data.path("mitmproxy/net/data/server.key"))
key = OpenSSL.crypto.load_privatekey( with open(raw_key) as f:
OpenSSL.crypto.FILETYPE_PEM, raw_key = f.read()
open(raw_key, "rb").read()) key = OpenSSL.crypto.load_privatekey(OpenSSL.crypto.FILETYPE_PEM, raw_key)
if self.ssl.get("v3_only", False): if self.ssl.get("v3_only", False):
method = OpenSSL.SSL.SSLv3_METHOD method = OpenSSL.SSL.SSLv3_METHOD
options = OpenSSL.SSL.OP_NO_SSLv2 | OpenSSL.SSL.OP_NO_TLSv1 options = OpenSSL.SSL.OP_NO_SSLv2 | OpenSSL.SSL.OP_NO_TLSv1
@ -64,7 +61,8 @@ class _TServer(tcp.TCPServer):
method = OpenSSL.SSL.SSLv23_METHOD method = OpenSSL.SSL.SSLv23_METHOD
options = None options = None
h.convert_to_ssl( h.convert_to_ssl(
cert, key, cert,
key,
method=method, method=method,
options=options, options=options,
handle_sni=getattr(h, "handle_sni", None), handle_sni=getattr(h, "handle_sni", None),
@ -103,7 +101,7 @@ class ServerTestBase:
@classmethod @classmethod
def teardown_class(cls): def teardown_class(cls):
cls.server.shutdown() cls.server.server.shutdown()
def teardown(self): def teardown(self):
self.server.server.wait_for_silence() self.server.server.wait_for_silence()

5
test/mitmproxy/platform/test_pf.py
View File

@ -9,10 +9,11 @@ class TestLookup:
def test_simple(self): def test_simple(self):
if sys.platform == "freebsd10": if sys.platform == "freebsd10":
p = tutils.test_data.path("mitmproxy/data/pf02") p = tutils.test_data.path("mitmproxy/data/pf02")
d = open(p, "rb").read()
else: else:
p = tutils.test_data.path("mitmproxy/data/pf01") p = tutils.test_data.path("mitmproxy/data/pf01")
d = open(p, "rb").read() with open(p, "rb") as f:
d = f.read()
assert pf.lookup("192.168.1.111", 40000, d) == ("5.5.5.5", 80) assert pf.lookup("192.168.1.111", 40000, d) == ("5.5.5.5", 80)
with pytest.raises(Exception, match="Could not resolve original destination"): with pytest.raises(Exception, match="Could not resolve original destination"):
pf.lookup("192.168.1.112", 40000, d) pf.lookup("192.168.1.112", 40000, d)

1
test/mitmproxy/proxy/protocol/test_http1.py
View File

@ -65,6 +65,7 @@ class TestExpectHeader(tservers.HTTPProxyTest):
assert resp.status_code == 200 assert resp.status_code == 200
client.finish() client.finish()
client.close()
class TestHeadContentLength(tservers.HTTPProxyTest): class TestHeadContentLength(tservers.HTTPProxyTest):

172
test/mitmproxy/proxy/protocol/test_http2.py
View File

@ -118,12 +118,16 @@ class _Http2TestBase:
self.master.reset([]) self.master.reset([])
self.server.server.handle_server_event = self.handle_server_event self.server.server.handle_server_event = self.handle_server_event
def _setup_connection(self): def teardown(self):
client = mitmproxy.net.tcp.TCPClient(("127.0.0.1", self.proxy.port)) if self.client:
client.connect() self.client.close()
def setup_connection(self):
self.client = mitmproxy.net.tcp.TCPClient(("127.0.0.1", self.proxy.port))
self.client.connect()
# send CONNECT request # send CONNECT request
client.wfile.write(http1.assemble_request(mitmproxy.net.http.Request( self.client.wfile.write(http1.assemble_request(mitmproxy.net.http.Request(
'authority', 'authority',
b'CONNECT', b'CONNECT',
b'', b'',
@ -134,13 +138,13 @@ class _Http2TestBase:
[(b'host', b'localhost:%d' % self.server.server.address[1])], [(b'host', b'localhost:%d' % self.server.server.address[1])],
b'', b'',
))) )))
client.wfile.flush() self.client.wfile.flush()
# read CONNECT response # read CONNECT response
while client.rfile.readline() != b"\r\n": while self.client.rfile.readline() != b"\r\n":
pass pass
client.convert_to_ssl(alpn_protos=[b'h2']) self.client.convert_to_ssl(alpn_protos=[b'h2'])
config = h2.config.H2Configuration( config = h2.config.H2Configuration(
client_side=True, client_side=True,
@ -148,10 +152,10 @@ class _Http2TestBase:
validate_inbound_headers=False) validate_inbound_headers=False)
h2_conn = h2.connection.H2Connection(config) h2_conn = h2.connection.H2Connection(config)
h2_conn.initiate_connection() h2_conn.initiate_connection()
client.wfile.write(h2_conn.data_to_send()) self.client.wfile.write(h2_conn.data_to_send())
client.wfile.flush() self.client.wfile.flush()
return client, h2_conn return h2_conn
def _send_request(self, def _send_request(self,
wfile, wfile,
@ -205,8 +209,8 @@ class TestSimple(_Http2Test):
if isinstance(event, h2.events.ConnectionTerminated): if isinstance(event, h2.events.ConnectionTerminated):
return False return False
elif isinstance(event, h2.events.RequestReceived): elif isinstance(event, h2.events.RequestReceived):
assert (b'client-foo', b'client-bar-1') in event.headers assert (b'self.client-foo', b'self.client-bar-1') in event.headers
assert (b'client-foo', b'client-bar-2') in event.headers assert (b'self.client-foo', b'self.client-bar-2') in event.headers
elif isinstance(event, h2.events.StreamEnded): elif isinstance(event, h2.events.StreamEnded):
import warnings import warnings
with warnings.catch_warnings(): with warnings.catch_warnings():
@ -233,32 +237,32 @@ class TestSimple(_Http2Test):
def test_simple(self): def test_simple(self):
response_body_buffer = b'' response_body_buffer = b''
client, h2_conn = self._setup_connection() h2_conn = self.setup_connection()
self._send_request( self._send_request(
client.wfile, self.client.wfile,
h2_conn, h2_conn,
headers=[ headers=[
(':authority', "127.0.0.1:{}".format(self.server.server.address[1])), (':authority', "127.0.0.1:{}".format(self.server.server.address[1])),
(':method', 'GET'), (':method', 'GET'),
(':scheme', 'https'), (':scheme', 'https'),
(':path', '/'), (':path', '/'),
('ClIeNt-FoO', 'client-bar-1'), ('self.client-FoO', 'self.client-bar-1'),
('ClIeNt-FoO', 'client-bar-2'), ('self.client-FoO', 'self.client-bar-2'),
], ],
body=b'request body') body=b'request body')
done = False done = False
while not done: while not done:
try: try:
raw = b''.join(http2.read_raw_frame(client.rfile)) raw = b''.join(http2.read_raw_frame(self.client.rfile))
events = h2_conn.receive_data(raw) events = h2_conn.receive_data(raw)
except exceptions.HttpException: except exceptions.HttpException:
print(traceback.format_exc()) print(traceback.format_exc())
assert False assert False
client.wfile.write(h2_conn.data_to_send()) self.client.wfile.write(h2_conn.data_to_send())
client.wfile.flush() self.client.wfile.flush()
for event in events: for event in events:
if isinstance(event, h2.events.DataReceived): if isinstance(event, h2.events.DataReceived):
@ -267,8 +271,8 @@ class TestSimple(_Http2Test):
done = True done = True
h2_conn.close_connection() h2_conn.close_connection()
client.wfile.write(h2_conn.data_to_send()) self.client.wfile.write(h2_conn.data_to_send())
client.wfile.flush() self.client.wfile.flush()
assert len(self.master.state.flows) == 1 assert len(self.master.state.flows) == 1
assert self.master.state.flows[0].response.status_code == 200 assert self.master.state.flows[0].response.status_code == 200
@ -317,10 +321,10 @@ class TestRequestWithPriority(_Http2Test):
def test_request_with_priority(self, http2_priority_enabled, priority, expected_priority): def test_request_with_priority(self, http2_priority_enabled, priority, expected_priority):
self.config.options.http2_priority = http2_priority_enabled self.config.options.http2_priority = http2_priority_enabled
client, h2_conn = self._setup_connection() h2_conn = self.setup_connection()
self._send_request( self._send_request(
client.wfile, self.client.wfile,
h2_conn, h2_conn,
headers=[ headers=[
(':authority', "127.0.0.1:{}".format(self.server.server.address[1])), (':authority', "127.0.0.1:{}".format(self.server.server.address[1])),
@ -336,22 +340,22 @@ class TestRequestWithPriority(_Http2Test):
done = False done = False
while not done: while not done:
try: try:
raw = b''.join(http2.read_raw_frame(client.rfile)) raw = b''.join(http2.read_raw_frame(self.client.rfile))
events = h2_conn.receive_data(raw) events = h2_conn.receive_data(raw)
except exceptions.HttpException: except exceptions.HttpException:
print(traceback.format_exc()) print(traceback.format_exc())
assert False assert False
client.wfile.write(h2_conn.data_to_send()) self.client.wfile.write(h2_conn.data_to_send())
client.wfile.flush() self.client.wfile.flush()
for event in events: for event in events:
if isinstance(event, h2.events.StreamEnded): if isinstance(event, h2.events.StreamEnded):
done = True done = True
h2_conn.close_connection() h2_conn.close_connection()
client.wfile.write(h2_conn.data_to_send()) self.client.wfile.write(h2_conn.data_to_send())
client.wfile.flush() self.client.wfile.flush()
assert len(self.master.state.flows) == 1 assert len(self.master.state.flows) == 1
@ -397,15 +401,15 @@ class TestPriority(_Http2Test):
self.config.options.http2_priority = http2_priority_enabled self.config.options.http2_priority = http2_priority_enabled
self.__class__.priority_data = [] self.__class__.priority_data = []
client, h2_conn = self._setup_connection() h2_conn = self.setup_connection()
if prioritize_before: if prioritize_before:
h2_conn.prioritize(1, exclusive=priority[0], depends_on=priority[1], weight=priority[2]) h2_conn.prioritize(1, exclusive=priority[0], depends_on=priority[1], weight=priority[2])
client.wfile.write(h2_conn.data_to_send()) self.client.wfile.write(h2_conn.data_to_send())
client.wfile.flush() self.client.wfile.flush()
self._send_request( self._send_request(
client.wfile, self.client.wfile,
h2_conn, h2_conn,
headers=[ headers=[
(':authority', "127.0.0.1:{}".format(self.server.server.address[1])), (':authority', "127.0.0.1:{}".format(self.server.server.address[1])),
@ -419,28 +423,28 @@ class TestPriority(_Http2Test):
if not prioritize_before: if not prioritize_before:
h2_conn.prioritize(1, exclusive=priority[0], depends_on=priority[1], weight=priority[2]) h2_conn.prioritize(1, exclusive=priority[0], depends_on=priority[1], weight=priority[2])
h2_conn.end_stream(1) h2_conn.end_stream(1)
client.wfile.write(h2_conn.data_to_send()) self.client.wfile.write(h2_conn.data_to_send())
client.wfile.flush() self.client.wfile.flush()
done = False done = False
while not done: while not done:
try: try:
raw = b''.join(http2.read_raw_frame(client.rfile)) raw = b''.join(http2.read_raw_frame(self.client.rfile))
events = h2_conn.receive_data(raw) events = h2_conn.receive_data(raw)
except exceptions.HttpException: except exceptions.HttpException:
print(traceback.format_exc()) print(traceback.format_exc())
assert False assert False
client.wfile.write(h2_conn.data_to_send()) self.client.wfile.write(h2_conn.data_to_send())
client.wfile.flush() self.client.wfile.flush()
for event in events: for event in events:
if isinstance(event, h2.events.StreamEnded): if isinstance(event, h2.events.StreamEnded):
done = True done = True
h2_conn.close_connection() h2_conn.close_connection()
client.wfile.write(h2_conn.data_to_send()) self.client.wfile.write(h2_conn.data_to_send())
client.wfile.flush() self.client.wfile.flush()
assert len(self.master.state.flows) == 1 assert len(self.master.state.flows) == 1
assert self.priority_data == expected_priority assert self.priority_data == expected_priority
@ -460,10 +464,10 @@ class TestStreamResetFromServer(_Http2Test):
return True return True
def test_request_with_priority(self): def test_request_with_priority(self):
client, h2_conn = self._setup_connection() h2_conn = self.setup_connection()
self._send_request( self._send_request(
client.wfile, self.client.wfile,
h2_conn, h2_conn,
headers=[ headers=[
(':authority', "127.0.0.1:{}".format(self.server.server.address[1])), (':authority', "127.0.0.1:{}".format(self.server.server.address[1])),
@ -476,22 +480,22 @@ class TestStreamResetFromServer(_Http2Test):
done = False done = False
while not done: while not done:
try: try:
raw = b''.join(http2.read_raw_frame(client.rfile)) raw = b''.join(http2.read_raw_frame(self.client.rfile))
events = h2_conn.receive_data(raw) events = h2_conn.receive_data(raw)
except exceptions.HttpException: except exceptions.HttpException:
print(traceback.format_exc()) print(traceback.format_exc())
assert False assert False
client.wfile.write(h2_conn.data_to_send()) self.client.wfile.write(h2_conn.data_to_send())
client.wfile.flush() self.client.wfile.flush()
for event in events: for event in events:
if isinstance(event, h2.events.StreamReset): if isinstance(event, h2.events.StreamReset):
done = True done = True
h2_conn.close_connection() h2_conn.close_connection()
client.wfile.write(h2_conn.data_to_send()) self.client.wfile.write(h2_conn.data_to_send())
client.wfile.flush() self.client.wfile.flush()
assert len(self.master.state.flows) == 1 assert len(self.master.state.flows) == 1
assert self.master.state.flows[0].response is None assert self.master.state.flows[0].response is None
@ -510,10 +514,10 @@ class TestBodySizeLimit(_Http2Test):
self.config.options.body_size_limit = "20" self.config.options.body_size_limit = "20"
self.config.options._processed["body_size_limit"] = 20 self.config.options._processed["body_size_limit"] = 20
client, h2_conn = self._setup_connection() h2_conn = self.setup_connection()
self._send_request( self._send_request(
client.wfile, self.client.wfile,
h2_conn, h2_conn,
headers=[ headers=[
(':authority', "127.0.0.1:{}".format(self.server.server.address[1])), (':authority', "127.0.0.1:{}".format(self.server.server.address[1])),
@ -527,22 +531,22 @@ class TestBodySizeLimit(_Http2Test):
done = False done = False
while not done: while not done:
try: try:
raw = b''.join(http2.read_raw_frame(client.rfile)) raw = b''.join(http2.read_raw_frame(self.client.rfile))
events = h2_conn.receive_data(raw) events = h2_conn.receive_data(raw)
except exceptions.HttpException: except exceptions.HttpException:
print(traceback.format_exc()) print(traceback.format_exc())
assert False assert False
client.wfile.write(h2_conn.data_to_send()) self.client.wfile.write(h2_conn.data_to_send())
client.wfile.flush() self.client.wfile.flush()
for event in events: for event in events:
if isinstance(event, h2.events.StreamReset): if isinstance(event, h2.events.StreamReset):
done = True done = True
h2_conn.close_connection() h2_conn.close_connection()
client.wfile.write(h2_conn.data_to_send()) self.client.wfile.write(h2_conn.data_to_send())
client.wfile.flush() self.client.wfile.flush()
assert len(self.master.state.flows) == 0 assert len(self.master.state.flows) == 0
@ -609,9 +613,9 @@ class TestPushPromise(_Http2Test):
return True return True
def test_push_promise(self): def test_push_promise(self):
client, h2_conn = self._setup_connection() h2_conn = self.setup_connection()
self._send_request(client.wfile, h2_conn, stream_id=1, headers=[ self._send_request(self.client.wfile, h2_conn, stream_id=1, headers=[
(':authority', "127.0.0.1:{}".format(self.server.server.address[1])), (':authority', "127.0.0.1:{}".format(self.server.server.address[1])),
(':method', 'GET'), (':method', 'GET'),
(':scheme', 'https'), (':scheme', 'https'),
@ -625,15 +629,15 @@ class TestPushPromise(_Http2Test):
responses = 0 responses = 0
while not done: while not done:
try: try:
raw = b''.join(http2.read_raw_frame(client.rfile)) raw = b''.join(http2.read_raw_frame(self.client.rfile))
events = h2_conn.receive_data(raw) events = h2_conn.receive_data(raw)
except exceptions.HttpException: except exceptions.HttpException:
print(traceback.format_exc()) print(traceback.format_exc())
assert False assert False
except: except:
break break
client.wfile.write(h2_conn.data_to_send()) self.client.wfile.write(h2_conn.data_to_send())
client.wfile.flush() self.client.wfile.flush()
for event in events: for event in events:
if isinstance(event, h2.events.StreamEnded): if isinstance(event, h2.events.StreamEnded):
@ -649,8 +653,8 @@ class TestPushPromise(_Http2Test):
done = True done = True
h2_conn.close_connection() h2_conn.close_connection()
client.wfile.write(h2_conn.data_to_send()) self.client.wfile.write(h2_conn.data_to_send())
client.wfile.flush() self.client.wfile.flush()
assert ended_streams == 3 assert ended_streams == 3
assert pushed_streams == 2 assert pushed_streams == 2
@ -665,9 +669,9 @@ class TestPushPromise(_Http2Test):
assert len(pushed_flows) == 2 assert len(pushed_flows) == 2
def test_push_promise_reset(self): def test_push_promise_reset(self):
client, h2_conn = self._setup_connection() h2_conn = self.setup_connection()
self._send_request(client.wfile, h2_conn, stream_id=1, headers=[ self._send_request(self.client.wfile, h2_conn, stream_id=1, headers=[
(':authority', "127.0.0.1:{}".format(self.server.server.address[1])), (':authority', "127.0.0.1:{}".format(self.server.server.address[1])),
(':method', 'GET'), (':method', 'GET'),
(':scheme', 'https'), (':scheme', 'https'),
@ -681,14 +685,14 @@ class TestPushPromise(_Http2Test):
responses = 0 responses = 0
while not done: while not done:
try: try:
raw = b''.join(http2.read_raw_frame(client.rfile)) raw = b''.join(http2.read_raw_frame(self.client.rfile))
events = h2_conn.receive_data(raw) events = h2_conn.receive_data(raw)
except exceptions.HttpException: except exceptions.HttpException:
print(traceback.format_exc()) print(traceback.format_exc())
assert False assert False
client.wfile.write(h2_conn.data_to_send()) self.client.wfile.write(h2_conn.data_to_send())
client.wfile.flush() self.client.wfile.flush()
for event in events: for event in events:
if isinstance(event, h2.events.StreamEnded) and event.stream_id == 1: if isinstance(event, h2.events.StreamEnded) and event.stream_id == 1:
@ -696,8 +700,8 @@ class TestPushPromise(_Http2Test):
elif isinstance(event, h2.events.PushedStreamReceived): elif isinstance(event, h2.events.PushedStreamReceived):
pushed_streams += 1 pushed_streams += 1
h2_conn.reset_stream(event.pushed_stream_id, error_code=0x8) h2_conn.reset_stream(event.pushed_stream_id, error_code=0x8)
client.wfile.write(h2_conn.data_to_send()) self.client.wfile.write(h2_conn.data_to_send())
client.wfile.flush() self.client.wfile.flush()
elif isinstance(event, h2.events.ResponseReceived): elif isinstance(event, h2.events.ResponseReceived):
responses += 1 responses += 1
if isinstance(event, h2.events.ConnectionTerminated): if isinstance(event, h2.events.ConnectionTerminated):
@ -707,8 +711,8 @@ class TestPushPromise(_Http2Test):
done = True done = True
h2_conn.close_connection() h2_conn.close_connection()
client.wfile.write(h2_conn.data_to_send()) self.client.wfile.write(h2_conn.data_to_send())
client.wfile.flush() self.client.wfile.flush()
bodies = [flow.response.content for flow in self.master.state.flows if flow.response] bodies = [flow.response.content for flow in self.master.state.flows if flow.response]
assert len(bodies) >= 1 assert len(bodies) >= 1
@ -728,9 +732,9 @@ class TestConnectionLost(_Http2Test):
return False return False
def test_connection_lost(self): def test_connection_lost(self):
client, h2_conn = self._setup_connection() h2_conn = self.setup_connection()
self._send_request(client.wfile, h2_conn, stream_id=1, headers=[ self._send_request(self.client.wfile, h2_conn, stream_id=1, headers=[
(':authority', "127.0.0.1:{}".format(self.server.server.address[1])), (':authority', "127.0.0.1:{}".format(self.server.server.address[1])),
(':method', 'GET'), (':method', 'GET'),
(':scheme', 'https'), (':scheme', 'https'),
@ -741,7 +745,7 @@ class TestConnectionLost(_Http2Test):
done = False done = False
while not done: while not done:
try: try:
raw = b''.join(http2.read_raw_frame(client.rfile)) raw = b''.join(http2.read_raw_frame(self.client.rfile))
h2_conn.receive_data(raw) h2_conn.receive_data(raw)
except exceptions.HttpException: except exceptions.HttpException:
print(traceback.format_exc()) print(traceback.format_exc())
@ -749,8 +753,8 @@ class TestConnectionLost(_Http2Test):
except: except:
break break
try: try:
client.wfile.write(h2_conn.data_to_send()) self.client.wfile.write(h2_conn.data_to_send())
client.wfile.flush() self.client.wfile.flush()
except: except:
break break
@ -782,12 +786,12 @@ class TestMaxConcurrentStreams(_Http2Test):
return True return True
def test_max_concurrent_streams(self): def test_max_concurrent_streams(self):
client, h2_conn = self._setup_connection() h2_conn = self.setup_connection()
new_streams = [1, 3, 5, 7, 9, 11] new_streams = [1, 3, 5, 7, 9, 11]
for stream_id in new_streams: for stream_id in new_streams:
# this will exceed MAX_CONCURRENT_STREAMS on the server connection # this will exceed MAX_CONCURRENT_STREAMS on the server connection
# and cause mitmproxy to throttle stream creation to the server # and cause mitmproxy to throttle stream creation to the server
self._send_request(client.wfile, h2_conn, stream_id=stream_id, headers=[ self._send_request(self.client.wfile, h2_conn, stream_id=stream_id, headers=[
(':authority', "127.0.0.1:{}".format(self.server.server.address[1])), (':authority', "127.0.0.1:{}".format(self.server.server.address[1])),
(':method', 'GET'), (':method', 'GET'),
(':scheme', 'https'), (':scheme', 'https'),
@ -798,20 +802,20 @@ class TestMaxConcurrentStreams(_Http2Test):
ended_streams = 0 ended_streams = 0
while ended_streams != len(new_streams): while ended_streams != len(new_streams):
try: try:
header, body = http2.read_raw_frame(client.rfile) header, body = http2.read_raw_frame(self.client.rfile)
events = h2_conn.receive_data(b''.join([header, body])) events = h2_conn.receive_data(b''.join([header, body]))
except: except:
break break
client.wfile.write(h2_conn.data_to_send()) self.client.wfile.write(h2_conn.data_to_send())
client.wfile.flush() self.client.wfile.flush()
for event in events: for event in events:
if isinstance(event, h2.events.StreamEnded): if isinstance(event, h2.events.StreamEnded):
ended_streams += 1 ended_streams += 1
h2_conn.close_connection() h2_conn.close_connection()
client.wfile.write(h2_conn.data_to_send()) self.client.wfile.write(h2_conn.data_to_send())
client.wfile.flush() self.client.wfile.flush()
assert len(self.master.state.flows) == len(new_streams) assert len(self.master.state.flows) == len(new_streams)
for flow in self.master.state.flows: for flow in self.master.state.flows:
@ -831,9 +835,9 @@ class TestConnectionTerminated(_Http2Test):
return True return True
def test_connection_terminated(self): def test_connection_terminated(self):
client, h2_conn = self._setup_connection() h2_conn = self.setup_connection()
self._send_request(client.wfile, h2_conn, headers=[ self._send_request(self.client.wfile, h2_conn, headers=[
(':authority', "127.0.0.1:{}".format(self.server.server.address[1])), (':authority', "127.0.0.1:{}".format(self.server.server.address[1])),
(':method', 'GET'), (':method', 'GET'),
(':scheme', 'https'), (':scheme', 'https'),
@ -844,7 +848,7 @@ class TestConnectionTerminated(_Http2Test):
connection_terminated_event = None connection_terminated_event = None
while not done: while not done:
try: try:
raw = b''.join(http2.read_raw_frame(client.rfile)) raw = b''.join(http2.read_raw_frame(self.client.rfile))
events = h2_conn.receive_data(raw) events = h2_conn.receive_data(raw)
for event in events: for event in events:
if isinstance(event, h2.events.ConnectionTerminated): if isinstance(event, h2.events.ConnectionTerminated):

122
test/mitmproxy/proxy/protocol/test_websocket.py
View File

@ -79,9 +79,13 @@ class _WebSocketTestBase:
self.master.reset([]) self.master.reset([])
self.server.server.handle_websockets = self.handle_websockets self.server.server.handle_websockets = self.handle_websockets
def _setup_connection(self): def teardown(self):
client = tcp.TCPClient(("127.0.0.1", self.proxy.port)) if self.client:
client.connect() self.client.close()
def setup_connection(self):
self.client = tcp.TCPClient(("127.0.0.1", self.proxy.port))
self.client.connect()
request = http.Request( request = http.Request(
"authority", "authority",
@ -92,14 +96,14 @@ class _WebSocketTestBase:
"", "",
"HTTP/1.1", "HTTP/1.1",
content=b'') content=b'')
client.wfile.write(http.http1.assemble_request(request)) self.client.wfile.write(http.http1.assemble_request(request))
client.wfile.flush() self.client.wfile.flush()
response = http.http1.read_response(client.rfile, request) response = http.http1.read_response(self.client.rfile, request)
if self.ssl: if self.ssl:
client.convert_to_ssl() self.client.convert_to_ssl()
assert client.ssl_established assert self.client.ssl_established
request = http.Request( request = http.Request(
"relative", "relative",
@ -116,14 +120,12 @@ class _WebSocketTestBase:
sec_websocket_key="1234", sec_websocket_key="1234",
), ),
content=b'') content=b'')
client.wfile.write(http.http1.assemble_request(request)) self.client.wfile.write(http.http1.assemble_request(request))
client.wfile.flush() self.client.wfile.flush()
response = http.http1.read_response(client.rfile, request) response = http.http1.read_response(self.client.rfile, request)
assert websockets.check_handshake(response.headers) assert websockets.check_handshake(response.headers)
return client
class _WebSocketTest(_WebSocketTestBase, _WebSocketServerBase): class _WebSocketTest(_WebSocketTestBase, _WebSocketServerBase):
@ -154,25 +156,25 @@ class TestSimple(_WebSocketTest):
wfile.flush() wfile.flush()
def test_simple(self): def test_simple(self):
client = self._setup_connection() self.setup_connection()
frame = websockets.Frame.from_file(client.rfile) frame = websockets.Frame.from_file(self.client.rfile)
assert frame.payload == b'server-foobar' assert frame.payload == b'server-foobar'
client.wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.TEXT, payload=b'client-foobar'))) self.client.wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.TEXT, payload=b'self.client-foobar')))
client.wfile.flush() self.client.wfile.flush()
frame = websockets.Frame.from_file(client.rfile) frame = websockets.Frame.from_file(self.client.rfile)
assert frame.payload == b'client-foobar' assert frame.payload == b'self.client-foobar'
client.wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.BINARY, payload=b'\xde\xad\xbe\xef'))) self.client.wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.BINARY, payload=b'\xde\xad\xbe\xef')))
client.wfile.flush() self.client.wfile.flush()
frame = websockets.Frame.from_file(client.rfile) frame = websockets.Frame.from_file(self.client.rfile)
assert frame.payload == b'\xde\xad\xbe\xef' assert frame.payload == b'\xde\xad\xbe\xef'
client.wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.CLOSE))) self.client.wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.CLOSE)))
client.wfile.flush() self.client.wfile.flush()
assert len(self.master.state.flows) == 2 assert len(self.master.state.flows) == 2
assert isinstance(self.master.state.flows[0], HTTPFlow) assert isinstance(self.master.state.flows[0], HTTPFlow)
@ -180,9 +182,9 @@ class TestSimple(_WebSocketTest):
assert len(self.master.state.flows[1].messages) == 5 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].content == b'server-foobar'
assert self.master.state.flows[1].messages[0].type == websockets.OPCODE.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].content == b'self.client-foobar'
assert self.master.state.flows[1].messages[1].type == websockets.OPCODE.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].content == b'self.client-foobar'
assert self.master.state.flows[1].messages[2].type == websockets.OPCODE.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].content == b'\xde\xad\xbe\xef'
assert self.master.state.flows[1].messages[3].type == websockets.OPCODE.BINARY assert self.master.state.flows[1].messages[3].type == websockets.OPCODE.BINARY
@ -203,19 +205,19 @@ class TestSimpleTLS(_WebSocketTest):
wfile.flush() wfile.flush()
def test_simple_tls(self): def test_simple_tls(self):
client = self._setup_connection() self.setup_connection()
frame = websockets.Frame.from_file(client.rfile) frame = websockets.Frame.from_file(self.client.rfile)
assert frame.payload == b'server-foobar' assert frame.payload == b'server-foobar'
client.wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.TEXT, payload=b'client-foobar'))) self.client.wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.TEXT, payload=b'self.client-foobar')))
client.wfile.flush() self.client.wfile.flush()
frame = websockets.Frame.from_file(client.rfile) frame = websockets.Frame.from_file(self.client.rfile)
assert frame.payload == b'client-foobar' assert frame.payload == b'self.client-foobar'
client.wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.CLOSE))) self.client.wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.CLOSE)))
client.wfile.flush() self.client.wfile.flush()
class TestPing(_WebSocketTest): class TestPing(_WebSocketTest):
@ -233,16 +235,16 @@ class TestPing(_WebSocketTest):
wfile.flush() wfile.flush()
def test_ping(self): def test_ping(self):
client = self._setup_connection() self.setup_connection()
frame = websockets.Frame.from_file(client.rfile) frame = websockets.Frame.from_file(self.client.rfile)
assert frame.header.opcode == websockets.OPCODE.PING assert frame.header.opcode == websockets.OPCODE.PING
assert frame.payload == b'foobar' assert frame.payload == b'foobar'
client.wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.PONG, payload=frame.payload))) self.client.wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.PONG, payload=frame.payload)))
client.wfile.flush() self.client.wfile.flush()
frame = websockets.Frame.from_file(client.rfile) frame = websockets.Frame.from_file(self.client.rfile)
assert frame.header.opcode == websockets.OPCODE.TEXT assert frame.header.opcode == websockets.OPCODE.TEXT
assert frame.payload == b'pong-received' assert frame.payload == b'pong-received'
@ -259,12 +261,12 @@ class TestPong(_WebSocketTest):
wfile.flush() wfile.flush()
def test_pong(self): def test_pong(self):
client = self._setup_connection() self.setup_connection()
client.wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.PING, payload=b'foobar'))) self.client.wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.PING, payload=b'foobar')))
client.wfile.flush() self.client.wfile.flush()
frame = websockets.Frame.from_file(client.rfile) frame = websockets.Frame.from_file(self.client.rfile)
assert frame.header.opcode == websockets.OPCODE.PONG assert frame.header.opcode == websockets.OPCODE.PONG
assert frame.payload == b'foobar' assert frame.payload == b'foobar'
@ -282,34 +284,34 @@ class TestClose(_WebSocketTest):
websockets.Frame.from_file(rfile) websockets.Frame.from_file(rfile)
def test_close(self): def test_close(self):
client = self._setup_connection() self.setup_connection()
client.wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.CLOSE))) self.client.wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.CLOSE)))
client.wfile.flush() self.client.wfile.flush()
websockets.Frame.from_file(client.rfile) websockets.Frame.from_file(self.client.rfile)
with pytest.raises(exceptions.TcpDisconnect): with pytest.raises(exceptions.TcpDisconnect):
websockets.Frame.from_file(client.rfile) websockets.Frame.from_file(self.client.rfile)
def test_close_payload_1(self): def test_close_payload_1(self):
client = self._setup_connection() self.setup_connection()
client.wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.CLOSE, payload=b'\00\42'))) self.client.wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.CLOSE, payload=b'\00\42')))
client.wfile.flush() self.client.wfile.flush()
websockets.Frame.from_file(client.rfile) websockets.Frame.from_file(self.client.rfile)
with pytest.raises(exceptions.TcpDisconnect): with pytest.raises(exceptions.TcpDisconnect):
websockets.Frame.from_file(client.rfile) websockets.Frame.from_file(self.client.rfile)
def test_close_payload_2(self): def test_close_payload_2(self):
client = self._setup_connection() self.setup_connection()
client.wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.CLOSE, payload=b'\00\42foobar'))) self.client.wfile.write(bytes(websockets.Frame(fin=1, opcode=websockets.OPCODE.CLOSE, payload=b'\00\42foobar')))
client.wfile.flush() self.client.wfile.flush()
websockets.Frame.from_file(client.rfile) websockets.Frame.from_file(self.client.rfile)
with pytest.raises(exceptions.TcpDisconnect): with pytest.raises(exceptions.TcpDisconnect):
websockets.Frame.from_file(client.rfile) websockets.Frame.from_file(self.client.rfile)
class TestInvalidFrame(_WebSocketTest): class TestInvalidFrame(_WebSocketTest):
@ -320,9 +322,9 @@ class TestInvalidFrame(_WebSocketTest):
wfile.flush() wfile.flush()
def test_invalid_frame(self): def test_invalid_frame(self):
client = self._setup_connection() self.setup_connection()
# with pytest.raises(exceptions.TcpDisconnect): # with pytest.raises(exceptions.TcpDisconnect):
frame = websockets.Frame.from_file(client.rfile) frame = websockets.Frame.from_file(self.client.rfile)
assert frame.header.opcode == 15 assert frame.header.opcode == 15
assert frame.payload == b'foobar' assert frame.payload == b'foobar'

13
test/mitmproxy/test_connections.py
View File

@ -1,5 +1,4 @@
import socket import socket
import os
import threading import threading
import ssl import ssl
import OpenSSL import OpenSSL
@ -140,16 +139,17 @@ class TestServerConnection:
assert d.last_log() assert d.last_log()
c.finish() c.finish()
c.close()
d.shutdown() d.shutdown()
def test_terminate_error(self): def test_terminate_error(self):
d = test.Daemon() d = test.Daemon()
c = connections.ServerConnection((d.IFACE, d.port)) c = connections.ServerConnection((d.IFACE, d.port))
c.connect() c.connect()
c.close()
c.connection = mock.Mock() c.connection = mock.Mock()
c.connection.recv = mock.Mock(return_value=False) c.connection.recv = mock.Mock(return_value=False)
c.connection.flush = mock.Mock(side_effect=exceptions.TcpDisconnect) c.connection.flush = mock.Mock(side_effect=exceptions.TcpDisconnect)
c.finish()
d.shutdown() d.shutdown()
def test_sni(self): def test_sni(self):
@ -194,22 +194,25 @@ class TestClientConnectionTLS:
s = socket.create_connection(address) s = socket.create_connection(address)
s = ctx.wrap_socket(s, server_hostname=sni) s = ctx.wrap_socket(s, server_hostname=sni)
s.send(b'foobar') s.send(b'foobar')
s.shutdown(socket.SHUT_RDWR) s.close()
threading.Thread(target=client_run).start() threading.Thread(target=client_run).start()
connection, client_address = sock.accept() connection, client_address = sock.accept()
c = connections.ClientConnection(connection, client_address, None) c = connections.ClientConnection(connection, client_address, None)
cert = tutils.test_data.path("mitmproxy/net/data/server.crt") cert = tutils.test_data.path("mitmproxy/net/data/server.crt")
with open(tutils.test_data.path("mitmproxy/net/data/server.key")) as f:
raw_key = f.read()
key = OpenSSL.crypto.load_privatekey( key = OpenSSL.crypto.load_privatekey(
OpenSSL.crypto.FILETYPE_PEM, OpenSSL.crypto.FILETYPE_PEM,
open(tutils.test_data.path("mitmproxy/net/data/server.key"), "rb").read()) raw_key)
c.convert_to_ssl(cert, key) c.convert_to_ssl(cert, key)
assert c.connected() assert c.connected()
assert c.sni == sni assert c.sni == sni
assert c.tls_established assert c.tls_established
assert c.rfile.read(6) == b'foobar' assert c.rfile.read(6) == b'foobar'
c.finish() c.finish()
sock.close()
class TestServerConnectionTLS(tservers.ServerTestBase): class TestServerConnectionTLS(tservers.ServerTestBase):
@ -222,7 +225,7 @@ class TestServerConnectionTLS(tservers.ServerTestBase):
@pytest.mark.parametrize("clientcert", [ @pytest.mark.parametrize("clientcert", [
None, None,
tutils.test_data.path("mitmproxy/data/clientcert"), tutils.test_data.path("mitmproxy/data/clientcert"),
os.path.join(tutils.test_data.path("mitmproxy/data/clientcert"), "client.pem"), tutils.test_data.path("mitmproxy/data/clientcert/client.pem"),
]) ])
def test_tls(self, clientcert): def test_tls(self, clientcert):
c = connections.ServerConnection(("127.0.0.1", self.port)) c = connections.ServerConnection(("127.0.0.1", self.port))

3
test/mitmproxy/test_proxy.py
View File

@ -32,8 +32,7 @@ class TestProcessProxyOptions:
opts = options.Options() opts = options.Options()
cmdline.common_options(parser, opts) cmdline.common_options(parser, opts)
args = parser.parse_args(args=args) args = parser.parse_args(args=args)
main.process_options(parser, opts, args) pconf = main.process_options(parser, opts, args)
pconf = config.ProxyConfig(opts)
return parser, pconf return parser, pconf
def assert_noerr(self, *args): def assert_noerr(self, *args):

2
test/mitmproxy/tservers.py
View File

@ -133,7 +133,7 @@ class ProxyTestBase:
@classmethod @classmethod
def teardown_class(cls): def teardown_class(cls):
# perf: we want to run tests in parallell # perf: we want to run tests in parallel
# should this ever cause an error, travis should catch it. # should this ever cause an error, travis should catch it.
# shutil.rmtree(cls.cadir) # shutil.rmtree(cls.cadir)
cls.proxy.shutdown() cls.proxy.shutdown()

6
test/pathod/language/test_base.py
View File

@ -202,12 +202,14 @@ class TestMisc:
e.parseString("m@1") e.parseString("m@1")
s = base.Settings(staticdir=str(tmpdir)) s = base.Settings(staticdir=str(tmpdir))
tmpdir.join("path").write_binary(b"a" * 20, ensure=True) with open(str(tmpdir.join("path")), 'wb') as f:
f.write(b"a" * 20)
v = e.parseString("m<path")[0] v = e.parseString("m<path")[0]
with pytest.raises(Exception, match="Invalid value length"): with pytest.raises(Exception, match="Invalid value length"):
v.values(s) v.values(s)
tmpdir.join("path2").write_binary(b"a" * 4, ensure=True) with open(str(tmpdir.join("path2")), 'wb') as f:
f.write(b"a" * 4)
v = e.parseString("m<path2")[0] v = e.parseString("m<path2")[0]
assert v.values(s) assert v.values(s)

4
test/pathod/language/test_generators.py
View File

@ -23,9 +23,7 @@ def test_filegenerator(tmpdir):
assert len(g[1:10]) == 9 assert len(g[1:10]) == 9
assert len(g[10000:10001]) == 0 assert len(g[10000:10001]) == 0
assert repr(g) assert repr(g)
# remove all references to FileGenerator instance to close the file g.close()
# handle.
del g
def test_transform_generator(): def test_transform_generator():

2
test/pathod/protocols/test_http2.py
View File

@ -202,7 +202,7 @@ class TestApplySettings(net_tservers.ServerTestBase):
def handle(self): def handle(self):
# check settings acknowledgement # check settings acknowledgement
assert self.rfile.read(9) == codecs.decode('000000040100000000', 'hex_codec') assert self.rfile.read(9) == codecs.decode('000000040100000000', 'hex_codec')
self.wfile.write("OK") self.wfile.write(b"OK")
self.wfile.flush() self.wfile.flush()
self.rfile.safe_read(9) # just to keep the connection alive a bit longer self.rfile.safe_read(9) # just to keep the connection alive a bit longer