From c6c98a0866bfe5ef7c5b33c8a0960608db050986 Mon Sep 17 00:00:00 2001 From: Abhinav Singh Date: Sat, 12 Oct 2019 22:55:29 -0700 Subject: [PATCH 1/6] Add tests for plugin_examples.* to ensure we never break functionality --- Makefile | 2 +- proxy.py | 2 +- tests.py | 102 +++++++++++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 104 insertions(+), 2 deletions(-) diff --git a/Makefile b/Makefile index 5dc56f9406..bb44f08a1d 100644 --- a/Makefile +++ b/Makefile @@ -39,7 +39,7 @@ release: package twine upload dist/* coverage: - coverage3 run --source=proxy tests.py + coverage3 run --source=proxy,plugin_examples tests.py coverage3 html open htmlcov/index.html diff --git a/proxy.py b/proxy.py index fc2064c5b7..c24a47d169 100755 --- a/proxy.py +++ b/proxy.py @@ -1444,7 +1444,7 @@ def on_request_complete(self) -> Union[socket.socket, bool]: # https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Connection # connection headers are meant for communication between client and # first intercepting proxy. - self.request.add_headers([(b'Via', b'1.1 proxy.py v%s' % version)]) + self.request.add_headers([(b'Via', b'1.1 %s' % PROXY_AGENT_HEADER_VALUE)]) # Disable args.disable_headers before dispatching to upstream self.server.queue( self.request.build( diff --git a/tests.py b/tests.py index e009f3da3f..518d6c118d 100644 --- a/tests.py +++ b/tests.py @@ -24,6 +24,7 @@ from unittest import mock import proxy +import plugin_examples if os.name != 'nt': import resource @@ -1701,6 +1702,107 @@ def mock_connection() -> Any: self.assertEqual(self.proxy_plugin.return_value.client._conn, self.mock_ssl_wrap.return_value) +class TestPluginExamples(unittest.TestCase): + + def setUp(self): + self.fileno = 10 + self._addr = ('127.0.0.1', 54382) + self.config = proxy.ProtocolConfig() + self.plugin = mock.MagicMock() + + @mock.patch('proxy.TcpServerConnection') + @mock.patch('selectors.DefaultSelector') + @mock.patch('socket.fromfd') + def test_post_data_modified_plugin( + self, + mock_fromfd: mock.Mock, + mock_selector: mock.Mock, + mock_server_conn: mock.Mock) -> None: + self.mock_fromfd = mock_fromfd + self.mock_selector = mock_selector + self.mock_server_conn = mock_server_conn + + self.config.plugins = { + b'ProtocolHandlerPlugin': [proxy.HttpProxyPlugin], + b'HttpProxyBasePlugin': [plugin_examples.ModifyPostDataPlugin] + } + self._conn = mock_fromfd.return_value + self.proxy = proxy.ProtocolHandler( + self.fileno, self._addr, config=self.config) + self.proxy.initialize() + + original = b'{"key": "value"}' + modified = b'{"key": "modified"}' + + self._conn.recv.return_value = proxy.build_http_request( + b'POST', b'http://httpbin.org/post', + headers={ + b'Host': b'httpbin.org', + b'Content-Type': b'application/x-www-form-urlencoded', + b'Content-Length': proxy.bytes_(len(original)), + }, + body=original + ) + self.mock_selector.return_value.select.side_effect = [ + [(selectors.SelectorKey( + fileobj=self._conn, + fd=self._conn.fileno, + events=selectors.EVENT_READ, + data=None), selectors.EVENT_READ)], ] + + self.proxy.run_once() + mock_server_conn.assert_called_with('httpbin.org', 80) + mock_server_conn.return_value.queue.assert_called_with( + proxy.build_http_request( + b'POST', b'/post', + headers={ + b'Host': b'httpbin.org', + b'Content-Length': proxy.bytes_(len(modified)), + b'Content-Type': b'application/json', + b'Via': b'1.1 %s' % proxy.PROXY_AGENT_HEADER_VALUE, + }, + body=modified + ) + ) + + @mock.patch('proxy.TcpServerConnection') + @mock.patch('selectors.DefaultSelector') + @mock.patch('socket.fromfd') + def test_proposed_rest_api_plugin( + self, + mock_fromfd: mock.Mock, + mock_selector: mock.Mock, + mock_server_conn: mock.Mock) -> None: + self.mock_fromfd = mock_fromfd + self.mock_selector = mock_selector + self.mock_server_conn = mock_server_conn + + self.config.plugins = { + b'ProtocolHandlerPlugin': [proxy.HttpProxyPlugin], + b'HttpProxyBasePlugin': [plugin_examples.ProposedRestApiPlugin] + } + self._conn = mock_fromfd.return_value + self.proxy = proxy.ProtocolHandler( + self.fileno, self._addr, config=self.config) + self.proxy.initialize() + + self._conn.recv.return_value = proxy.build_http_request( + b'GET', b'http://%s%s' % (plugin_examples.ProposedRestApiPlugin.API_SERVER, b'/v1/users/'), + headers={ + b'Host': plugin_examples.ProposedRestApiPlugin.API_SERVER, + } + ) + self.mock_selector.return_value.select.side_effect = [ + [(selectors.SelectorKey( + fileobj=self._conn, + fd=self._conn.fileno, + events=selectors.EVENT_READ, + data=None), selectors.EVENT_READ)], ] + self.proxy.run_once() + + mock_server_conn.assert_not_called() + + class TestHttpRequestRejected(unittest.TestCase): def setUp(self) -> None: From afade55d2dbbe99872b83bda4f95b2dd04d7d30b Mon Sep 17 00:00:00 2001 From: Abhinav Singh Date: Sat, 12 Oct 2019 23:55:53 -0700 Subject: [PATCH 2/6] Add tests for plugin_examples.* --- plugin_examples.py | 11 ++-- proxy.py | 10 +++- tests.py | 138 ++++++++++++++++++++++++++++++++++----------- 3 files changed, 117 insertions(+), 42 deletions(-) diff --git a/plugin_examples.py b/plugin_examples.py index 5a1f123e1c..584730ce5d 100644 --- a/plugin_examples.py +++ b/plugin_examples.py @@ -107,15 +107,16 @@ def on_upstream_connection_close(self) -> None: class RedirectToCustomServerPlugin(proxy.HttpProxyBasePlugin): """Modifies client request to redirect all incoming requests to a fixed server address.""" - UPSTREAM_SERVER = b'http://localhost:8899' + UPSTREAM_SERVER = b'http://localhost:8899/' def before_upstream_connection(self, request: proxy.HttpParser) -> Optional[proxy.HttpParser]: # Redirect all non-https requests to inbuilt WebServer. if request.method != proxy.httpMethods.CONNECT: - request.url = urlparse.urlsplit(self.UPSTREAM_SERVER) - # This command will re-parse modified url and - # update host, port, path fields - request.set_line_attributes() + request.set_url(self.UPSTREAM_SERVER) + # Update Host header too, otherwise upstream can reject our request + if request.has_header(b'Host'): + request.del_header(b'Host') + request.add_header(b'Host', urlparse.urlsplit(self.UPSTREAM_SERVER).netloc) return request def handle_client_request(self, request: proxy.HttpParser) -> Optional[proxy.HttpParser]: diff --git a/proxy.py b/proxy.py index c24a47d169..c663ae4e62 100755 --- a/proxy.py +++ b/proxy.py @@ -598,6 +598,10 @@ def del_headers(self, headers: List[bytes]) -> None: for key in headers: self.del_header(key.lower()) + def set_url(self, url: bytes) -> None: + self.url = urlparse.urlsplit(url) + self.set_line_attributes() + def set_line_attributes(self) -> None: if self.type == httpParserTypes.REQUEST_PARSER: if self.method == httpMethods.CONNECT and self.url: @@ -700,13 +704,12 @@ def process_line(self, raw: bytes) -> None: line = raw.split(WHITESPACE) if self.type == httpParserTypes.REQUEST_PARSER: self.method = line[0].upper() - self.url = urlparse.urlsplit(line[1]) + self.set_url(line[1]) self.version = line[2] else: self.version = line[0] self.code = line[1] self.reason = WHITESPACE.join(line[2:]) - self.set_line_attributes() def process_header(self, raw: bytes) -> None: parts = raw.split(COLON) @@ -2095,7 +2098,8 @@ def initialize(self) -> None: """Optionally upgrades connection to HTTPS, set conn in non-blocking mode and initializes plugins.""" conn = self.optionally_wrap_socket(self.client.connection) conn.setblocking(False) - self.client = TcpClientConnection(conn=conn, addr=self.addr) + if self.config.encryption_enabled(): + self.client = TcpClientConnection(conn=conn, addr=self.addr) if b'ProtocolHandlerPlugin' in self.config.plugins: for klass in self.config.plugins[b'ProtocolHandlerPlugin']: instance = klass(self.config, self.client, self.request) diff --git a/tests.py b/tests.py index 518d6c118d..2d79bc9170 100644 --- a/tests.py +++ b/tests.py @@ -10,6 +10,7 @@ import base64 import errno import ipaddress +import json import logging import multiprocessing import os @@ -20,11 +21,12 @@ import unittest import uuid from contextlib import closing -from typing import Dict, Optional, Tuple, Union, Any, cast +from typing import Dict, Optional, Tuple, Union, Any, cast, Type from unittest import mock +from urllib import parse as urlparse -import proxy import plugin_examples +import proxy if os.name != 'nt': import resource @@ -1702,35 +1704,46 @@ def mock_connection() -> Any: self.assertEqual(self.proxy_plugin.return_value.client._conn, self.mock_ssl_wrap.return_value) -class TestPluginExamples(unittest.TestCase): +class TestHttpProxyPluginExamples(unittest.TestCase): - def setUp(self): + @mock.patch('selectors.DefaultSelector') + @mock.patch('socket.fromfd') + def setUp(self, + mock_fromfd: mock.Mock, + mock_selector: mock.Mock) -> None: self.fileno = 10 self._addr = ('127.0.0.1', 54382) self.config = proxy.ProtocolConfig() self.plugin = mock.MagicMock() - @mock.patch('proxy.TcpServerConnection') - @mock.patch('selectors.DefaultSelector') - @mock.patch('socket.fromfd') - def test_post_data_modified_plugin( - self, - mock_fromfd: mock.Mock, - mock_selector: mock.Mock, - mock_server_conn: mock.Mock) -> None: self.mock_fromfd = mock_fromfd self.mock_selector = mock_selector - self.mock_server_conn = mock_server_conn + + plugin: Type[proxy.HttpProxyBasePlugin] = plugin_examples.ModifyPostDataPlugin + if self._testMethodName == 'test_modify_post_data_plugin': + plugin = plugin_examples.ModifyPostDataPlugin + elif self._testMethodName == 'test_proposed_rest_api_plugin': + plugin = plugin_examples.ProposedRestApiPlugin + elif self._testMethodName == 'test_redirect_to_custom_server_plugin': + plugin = plugin_examples.RedirectToCustomServerPlugin + elif self._testMethodName == 'test_filter_by_upstream_host_plugin': + plugin = plugin_examples.FilterByUpstreamHostPlugin + elif self._testMethodName == 'test_cache_responses_plugin': + plugin = plugin_examples.CacheResponsesPlugin + elif self._testMethodName == 'test_man_in_the_middle_plugin': + plugin = plugin_examples.ManInTheMiddlePlugin self.config.plugins = { b'ProtocolHandlerPlugin': [proxy.HttpProxyPlugin], - b'HttpProxyBasePlugin': [plugin_examples.ModifyPostDataPlugin] + b'HttpProxyBasePlugin': [plugin], } self._conn = mock_fromfd.return_value self.proxy = proxy.ProtocolHandler( self.fileno, self._addr, config=self.config) self.proxy.initialize() + @mock.patch('proxy.TcpServerConnection') + def test_modify_post_data_plugin(self, mock_server_conn: mock.Mock) -> None: original = b'{"key": "value"}' modified = b'{"key": "modified"}' @@ -1766,28 +1779,11 @@ def test_post_data_modified_plugin( ) @mock.patch('proxy.TcpServerConnection') - @mock.patch('selectors.DefaultSelector') - @mock.patch('socket.fromfd') def test_proposed_rest_api_plugin( - self, - mock_fromfd: mock.Mock, - mock_selector: mock.Mock, - mock_server_conn: mock.Mock) -> None: - self.mock_fromfd = mock_fromfd - self.mock_selector = mock_selector - self.mock_server_conn = mock_server_conn - - self.config.plugins = { - b'ProtocolHandlerPlugin': [proxy.HttpProxyPlugin], - b'HttpProxyBasePlugin': [plugin_examples.ProposedRestApiPlugin] - } - self._conn = mock_fromfd.return_value - self.proxy = proxy.ProtocolHandler( - self.fileno, self._addr, config=self.config) - self.proxy.initialize() - + self, mock_server_conn: mock.Mock) -> None: + path = b'/v1/users/' self._conn.recv.return_value = proxy.build_http_request( - b'GET', b'http://%s%s' % (plugin_examples.ProposedRestApiPlugin.API_SERVER, b'/v1/users/'), + b'GET', b'http://%s%s' % (plugin_examples.ProposedRestApiPlugin.API_SERVER, path), headers={ b'Host': plugin_examples.ProposedRestApiPlugin.API_SERVER, } @@ -1801,6 +1797,80 @@ def test_proposed_rest_api_plugin( self.proxy.run_once() mock_server_conn.assert_not_called() + self.assertEqual( + self.proxy.client.buffer, + proxy.build_http_response( + proxy.httpStatusCodes.OK, reason=b'OK', + headers={b'Content-Type': b'application/json'}, + body=proxy.bytes_(json.dumps(plugin_examples.ProposedRestApiPlugin.REST_API_SPEC[path])) + )) + + @mock.patch('proxy.TcpServerConnection') + def test_redirect_to_custom_server_plugin( + self, mock_server_conn: mock.Mock) -> None: + request = proxy.build_http_request( + b'GET', b'http://example.org/get', + headers={ + b'Host': b'example.org', + } + ) + self._conn.recv.return_value = request + self.mock_selector.return_value.select.side_effect = [ + [(selectors.SelectorKey( + fileobj=self._conn, + fd=self._conn.fileno, + events=selectors.EVENT_READ, + data=None), selectors.EVENT_READ)], ] + self.proxy.run_once() + + upstream = urlparse.urlsplit( + plugin_examples.RedirectToCustomServerPlugin.UPSTREAM_SERVER) + mock_server_conn.assert_called_with('localhost', 8899) + mock_server_conn.return_value.queue.assert_called_with( + proxy.build_http_request( + b'GET', upstream.path, + headers={ + b'Host': upstream.netloc, + b'Via': b'1.1 %s' % proxy.PROXY_AGENT_HEADER_VALUE, + } + ) + ) + + @mock.patch('proxy.TcpServerConnection') + def test_filter_by_upstream_host_plugin( + self, mock_server_conn: mock.Mock) -> None: + request = proxy.build_http_request( + b'GET', b'http://google.com/', + headers={ + b'Host': b'google.com', + } + ) + self._conn.recv.return_value = request + self.mock_selector.return_value.select.side_effect = [ + [(selectors.SelectorKey( + fileobj=self._conn, + fd=self._conn.fileno, + events=selectors.EVENT_READ, + data=None), selectors.EVENT_READ)], ] + self.proxy.run_once() + + mock_server_conn.assert_not_called() + self.assertEqual( + self.proxy.client.buffer, + proxy.build_http_response( + proxy.httpStatusCodes.I_AM_A_TEAPOT, + reason=b'I\'m a tea pot', + headers={ + proxy.PROXY_AGENT_HEADER_KEY: proxy.PROXY_AGENT_HEADER_VALUE + }, + ) + ) + + def test_cache_responses_plugin(self) -> None: + pass + + def test_man_in_the_middle_plugin(self) -> None: + pass class TestHttpRequestRejected(unittest.TestCase): From d92778a174fb6cf4ff79aaf643f76d3e2dcf00a8 Mon Sep 17 00:00:00 2001 From: Abhinav Singh Date: Sun, 13 Oct 2019 00:57:18 -0700 Subject: [PATCH 3/6] Test man in the middle --- proxy.py | 2 +- tests.py | 83 ++++++++++++++++++++++++++++++++++++++++++++++++++++---- 2 files changed, 79 insertions(+), 6 deletions(-) diff --git a/proxy.py b/proxy.py index c663ae4e62..7d989c11c8 100755 --- a/proxy.py +++ b/proxy.py @@ -1229,7 +1229,7 @@ def get_descriptors( def write_to_descriptors(self, w: List[Union[int, _HasFileno]]) -> bool: if self.request.has_upstream_server() and \ self.server and not self.server.closed and \ - self.server.buffer_size() > 0 and \ + self.server.has_buffer() and \ self.server.connection in w: logger.debug('Server is write ready, flushing buffer') try: diff --git a/tests.py b/tests.py index 2d79bc9170..1922549443 100644 --- a/tests.py +++ b/tests.py @@ -1051,8 +1051,11 @@ def assert_tunnel_response( def test_http_tunnel(self, mock_server_connection: mock.Mock) -> None: server = mock_server_connection.return_value server.connect.return_value = True - server.buffer_size.return_value = 0 - server.has_buffer.side_effect = [False, False, False, True] + + def has_buffer() -> bool: + return server.queue.called + + server.has_buffer.side_effect = has_buffer self.mock_selector.return_value.select.side_effect = [ [(selectors.SelectorKey( fileobj=self._conn, @@ -1866,11 +1869,81 @@ def test_filter_by_upstream_host_plugin( ) ) - def test_cache_responses_plugin(self) -> None: + @mock.patch('proxy.TcpServerConnection') + def test_cache_responses_plugin( + self, mock_server_conn: mock.Mock) -> None: pass - def test_man_in_the_middle_plugin(self) -> None: - pass + @mock.patch('proxy.TcpServerConnection') + def test_man_in_the_middle_plugin( + self, mock_server_conn: mock.Mock) -> None: + request = proxy.build_http_request( + b'GET', b'http://super.secure/', + headers={ + b'Host': b'super.secure', + } + ) + self._conn.recv.return_value = request + + server = mock_server_conn.return_value + server.connect.return_value = True + + def has_buffer() -> None: + return server.queue.called + + def closed() -> bool: + return not server.connect.called + + server.has_buffer.side_effect = has_buffer + type(server).closed = mock.PropertyMock(side_effect=closed) + + self.mock_selector.return_value.select.side_effect = [ + [(selectors.SelectorKey( + fileobj=self._conn, + fd=self._conn.fileno, + events=selectors.EVENT_READ, + data=None), selectors.EVENT_READ)], + [(selectors.SelectorKey( + fileobj=server.connection, + fd=server.connection.fileno, + events=selectors.EVENT_WRITE, + data=None), selectors.EVENT_WRITE)], + [(selectors.SelectorKey( + fileobj=server.connection, + fd=server.connection.fileno, + events=selectors.EVENT_READ, + data=None), selectors.EVENT_READ)], ] + + # Client read + self.proxy.run_once() + mock_server_conn.assert_called_with('super.secure', 80) + server.connect.assert_called_once() + queued_request = \ + proxy.build_http_request( + b'GET', b'/', + headers={ + b'Host': b'super.secure', + b'Via': b'1.1 %s' % proxy.PROXY_AGENT_HEADER_VALUE + } + ) + server.queue.assert_called_once_with(queued_request) + + # Server write + self.proxy.run_once() + server.flush.assert_called_once() + + # Server read + server.recv.return_value = \ + proxy.build_http_response( + proxy.httpStatusCodes.OK, + reason=b'OK', body=b'Original Response From Upstream') + self.proxy.run_once() + self.assertEqual( + self.proxy.client.buffer, + proxy.build_http_response( + proxy.httpStatusCodes.OK, + reason=b'OK', body=b'Hello from man in the middle') + ) class TestHttpRequestRejected(unittest.TestCase): From bc8ac4340f476d9891f719269bd0be821af6521f Mon Sep 17 00:00:00 2001 From: Abhinav Singh Date: Sun, 13 Oct 2019 00:59:23 -0700 Subject: [PATCH 4/6] Lint fixes --- tests.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests.py b/tests.py index 1922549443..a53afd006f 100644 --- a/tests.py +++ b/tests.py @@ -1053,7 +1053,7 @@ def test_http_tunnel(self, mock_server_connection: mock.Mock) -> None: server.connect.return_value = True def has_buffer() -> bool: - return server.queue.called + return cast(bool, server.queue.called) server.has_buffer.side_effect = has_buffer self.mock_selector.return_value.select.side_effect = [ @@ -1888,8 +1888,8 @@ def test_man_in_the_middle_plugin( server = mock_server_conn.return_value server.connect.return_value = True - def has_buffer() -> None: - return server.queue.called + def has_buffer() -> bool: + return cast(bool, server.queue.called) def closed() -> bool: return not server.connect.called From 49ef1f647b4aa1f72e9c91719f8ae25801c567be Mon Sep 17 00:00:00 2001 From: Abhinav Singh Date: Sun, 13 Oct 2019 02:35:25 -0700 Subject: [PATCH 5/6] Checkin --- plugin_examples.py | 3 +- proxy.py | 3 +- tests.py | 376 +++++++++++++++++++++++++++++++++++++++------ 3 files changed, 332 insertions(+), 50 deletions(-) diff --git a/plugin_examples.py b/plugin_examples.py index 584730ce5d..289e0db58e 100644 --- a/plugin_examples.py +++ b/plugin_examples.py @@ -49,7 +49,8 @@ class ProposedRestApiPlugin(proxy.HttpProxyBasePlugin): Used to test and develop client side applications without need of an actual upstream REST API server. - Returns proposed REST API mock responses to the client.""" + Returns proposed REST API mock responses to the client + without establishing upstream connection.""" API_SERVER = b'api.example.com' diff --git a/proxy.py b/proxy.py index 7d989c11c8..bd72c91d13 100755 --- a/proxy.py +++ b/proxy.py @@ -1381,6 +1381,8 @@ def on_request_complete(self) -> Union[socket.socket, bool]: if not self.request.has_upstream_server(): return False + self.authenticate() + # Note: can raise HttpRequestRejected exception # Invoke plugin.before_upstream_connection for plugin in self.plugins.values(): @@ -1389,7 +1391,6 @@ def on_request_complete(self) -> Union[socket.socket, bool]: return False self.request = r - self.authenticate() self.connect_upstream() for plugin in self.plugins.values(): diff --git a/tests.py b/tests.py index a53afd006f..3a3be14190 100644 --- a/tests.py +++ b/tests.py @@ -47,6 +47,23 @@ def get_available_port() -> int: return int(port) +def get_plugin_by_test_name(test_name: str) -> Type[proxy.HttpProxyBasePlugin]: + plugin: Type[proxy.HttpProxyBasePlugin] = plugin_examples.ModifyPostDataPlugin + if test_name == 'test_modify_post_data_plugin': + plugin = plugin_examples.ModifyPostDataPlugin + elif test_name == 'test_proposed_rest_api_plugin': + plugin = plugin_examples.ProposedRestApiPlugin + elif test_name == 'test_redirect_to_custom_server_plugin': + plugin = plugin_examples.RedirectToCustomServerPlugin + elif test_name == 'test_filter_by_upstream_host_plugin': + plugin = plugin_examples.FilterByUpstreamHostPlugin + elif test_name == 'test_cache_responses_plugin': + plugin = plugin_examples.CacheResponsesPlugin + elif test_name == 'test_man_in_the_middle_plugin': + plugin = plugin_examples.ManInTheMiddlePlugin + return plugin + + class TestTextBytes(unittest.TestCase): def test_text(self) -> None: @@ -1576,6 +1593,233 @@ def test_proxy_plugin_before_upstream_connection_can_teardown( mock_server_conn.assert_not_called() +class TestHttpProxyPluginExamples(unittest.TestCase): + + @mock.patch('selectors.DefaultSelector') + @mock.patch('socket.fromfd') + def setUp(self, + mock_fromfd: mock.Mock, + mock_selector: mock.Mock) -> None: + self.fileno = 10 + self._addr = ('127.0.0.1', 54382) + self.config = proxy.ProtocolConfig() + self.plugin = mock.MagicMock() + + self.mock_fromfd = mock_fromfd + self.mock_selector = mock_selector + + plugin = get_plugin_by_test_name(self._testMethodName) + + self.config.plugins = { + b'ProtocolHandlerPlugin': [proxy.HttpProxyPlugin], + b'HttpProxyBasePlugin': [plugin], + } + self._conn = mock_fromfd.return_value + self.proxy = proxy.ProtocolHandler( + self.fileno, self._addr, config=self.config) + self.proxy.initialize() + + @mock.patch('proxy.TcpServerConnection') + def test_modify_post_data_plugin(self, mock_server_conn: mock.Mock) -> None: + original = b'{"key": "value"}' + modified = b'{"key": "modified"}' + + self._conn.recv.return_value = proxy.build_http_request( + b'POST', b'http://httpbin.org/post', + headers={ + b'Host': b'httpbin.org', + b'Content-Type': b'application/x-www-form-urlencoded', + b'Content-Length': proxy.bytes_(len(original)), + }, + body=original + ) + self.mock_selector.return_value.select.side_effect = [ + [(selectors.SelectorKey( + fileobj=self._conn, + fd=self._conn.fileno, + events=selectors.EVENT_READ, + data=None), selectors.EVENT_READ)], ] + + self.proxy.run_once() + mock_server_conn.assert_called_with('httpbin.org', 80) + mock_server_conn.return_value.queue.assert_called_with( + proxy.build_http_request( + b'POST', b'/post', + headers={ + b'Host': b'httpbin.org', + b'Content-Length': proxy.bytes_(len(modified)), + b'Content-Type': b'application/json', + b'Via': b'1.1 %s' % proxy.PROXY_AGENT_HEADER_VALUE, + }, + body=modified + ) + ) + + @mock.patch('proxy.TcpServerConnection') + def test_proposed_rest_api_plugin( + self, mock_server_conn: mock.Mock) -> None: + path = b'/v1/users/' + self._conn.recv.return_value = proxy.build_http_request( + b'GET', b'http://%s%s' % (plugin_examples.ProposedRestApiPlugin.API_SERVER, path), + headers={ + b'Host': plugin_examples.ProposedRestApiPlugin.API_SERVER, + } + ) + self.mock_selector.return_value.select.side_effect = [ + [(selectors.SelectorKey( + fileobj=self._conn, + fd=self._conn.fileno, + events=selectors.EVENT_READ, + data=None), selectors.EVENT_READ)], ] + self.proxy.run_once() + + mock_server_conn.assert_not_called() + self.assertEqual( + self.proxy.client.buffer, + proxy.build_http_response( + proxy.httpStatusCodes.OK, reason=b'OK', + headers={b'Content-Type': b'application/json'}, + body=proxy.bytes_(json.dumps(plugin_examples.ProposedRestApiPlugin.REST_API_SPEC[path])) + )) + + @mock.patch('proxy.TcpServerConnection') + def test_redirect_to_custom_server_plugin( + self, mock_server_conn: mock.Mock) -> None: + request = proxy.build_http_request( + b'GET', b'http://example.org/get', + headers={ + b'Host': b'example.org', + } + ) + self._conn.recv.return_value = request + self.mock_selector.return_value.select.side_effect = [ + [(selectors.SelectorKey( + fileobj=self._conn, + fd=self._conn.fileno, + events=selectors.EVENT_READ, + data=None), selectors.EVENT_READ)], ] + self.proxy.run_once() + + upstream = urlparse.urlsplit( + plugin_examples.RedirectToCustomServerPlugin.UPSTREAM_SERVER) + mock_server_conn.assert_called_with('localhost', 8899) + mock_server_conn.return_value.queue.assert_called_with( + proxy.build_http_request( + b'GET', upstream.path, + headers={ + b'Host': upstream.netloc, + b'Via': b'1.1 %s' % proxy.PROXY_AGENT_HEADER_VALUE, + } + ) + ) + + @mock.patch('proxy.TcpServerConnection') + def test_filter_by_upstream_host_plugin( + self, mock_server_conn: mock.Mock) -> None: + request = proxy.build_http_request( + b'GET', b'http://google.com/', + headers={ + b'Host': b'google.com', + } + ) + self._conn.recv.return_value = request + self.mock_selector.return_value.select.side_effect = [ + [(selectors.SelectorKey( + fileobj=self._conn, + fd=self._conn.fileno, + events=selectors.EVENT_READ, + data=None), selectors.EVENT_READ)], ] + self.proxy.run_once() + + mock_server_conn.assert_not_called() + self.assertEqual( + self.proxy.client.buffer, + proxy.build_http_response( + proxy.httpStatusCodes.I_AM_A_TEAPOT, + reason=b'I\'m a tea pot', + headers={ + proxy.PROXY_AGENT_HEADER_KEY: proxy.PROXY_AGENT_HEADER_VALUE + }, + ) + ) + + @mock.patch('proxy.TcpServerConnection') + def test_cache_responses_plugin( + self, mock_server_conn: mock.Mock) -> None: + pass + + @mock.patch('proxy.TcpServerConnection') + def test_man_in_the_middle_plugin( + self, mock_server_conn: mock.Mock) -> None: + request = proxy.build_http_request( + b'GET', b'http://super.secure/', + headers={ + b'Host': b'super.secure', + } + ) + self._conn.recv.return_value = request + + server = mock_server_conn.return_value + server.connect.return_value = True + + def has_buffer() -> bool: + return cast(bool, server.queue.called) + + def closed() -> bool: + return not server.connect.called + + server.has_buffer.side_effect = has_buffer + type(server).closed = mock.PropertyMock(side_effect=closed) + + self.mock_selector.return_value.select.side_effect = [ + [(selectors.SelectorKey( + fileobj=self._conn, + fd=self._conn.fileno, + events=selectors.EVENT_READ, + data=None), selectors.EVENT_READ)], + [(selectors.SelectorKey( + fileobj=server.connection, + fd=server.connection.fileno, + events=selectors.EVENT_WRITE, + data=None), selectors.EVENT_WRITE)], + [(selectors.SelectorKey( + fileobj=server.connection, + fd=server.connection.fileno, + events=selectors.EVENT_READ, + data=None), selectors.EVENT_READ)], ] + + # Client read + self.proxy.run_once() + mock_server_conn.assert_called_with('super.secure', 80) + server.connect.assert_called_once() + queued_request = \ + proxy.build_http_request( + b'GET', b'/', + headers={ + b'Host': b'super.secure', + b'Via': b'1.1 %s' % proxy.PROXY_AGENT_HEADER_VALUE + } + ) + server.queue.assert_called_once_with(queued_request) + + # Server write + self.proxy.run_once() + server.flush.assert_called_once() + + # Server read + server.recv.return_value = \ + proxy.build_http_response( + proxy.httpStatusCodes.OK, + reason=b'OK', body=b'Original Response From Upstream') + self.proxy.run_once() + self.assertEqual( + self.proxy.client.buffer, + proxy.build_http_response( + proxy.httpStatusCodes.OK, + reason=b'OK', body=b'Hello from man in the middle') + ) + + class TestHttpProxyTlsInterception(unittest.TestCase): @mock.patch('ssl.wrap_socket') @@ -1707,98 +1951,134 @@ def mock_connection() -> Any: self.assertEqual(self.proxy_plugin.return_value.client._conn, self.mock_ssl_wrap.return_value) -class TestHttpProxyPluginExamples(unittest.TestCase): +class TestHttpProxyPluginExamplesWithTlsInterception(unittest.TestCase): + @mock.patch('ssl.wrap_socket') + @mock.patch('ssl.create_default_context') + @mock.patch('proxy.TcpServerConnection') + @mock.patch('subprocess.Popen') @mock.patch('selectors.DefaultSelector') @mock.patch('socket.fromfd') def setUp(self, mock_fromfd: mock.Mock, - mock_selector: mock.Mock) -> None: + mock_selector: mock.Mock, + mock_popen: mock.Mock, + mock_server_conn: mock.Mock, + mock_ssl_context: mock.Mock, + mock_ssl_wrap: mock.Mock) -> None: + self.mock_fromfd = mock_fromfd + self.mock_selector = mock_selector + self.mock_popen = mock_popen + self.mock_server_conn = mock_server_conn + self.mock_ssl_context = mock_ssl_context + self.mock_ssl_wrap = mock_ssl_wrap + self.fileno = 10 self._addr = ('127.0.0.1', 54382) - self.config = proxy.ProtocolConfig() + self.config = proxy.ProtocolConfig( + ca_cert_file='ca-cert.pem', + ca_key_file='ca-key.pem', + ca_signing_key_file='ca-signing-key.pem',) self.plugin = mock.MagicMock() - self.mock_fromfd = mock_fromfd - self.mock_selector = mock_selector - - plugin: Type[proxy.HttpProxyBasePlugin] = plugin_examples.ModifyPostDataPlugin - if self._testMethodName == 'test_modify_post_data_plugin': - plugin = plugin_examples.ModifyPostDataPlugin - elif self._testMethodName == 'test_proposed_rest_api_plugin': - plugin = plugin_examples.ProposedRestApiPlugin - elif self._testMethodName == 'test_redirect_to_custom_server_plugin': - plugin = plugin_examples.RedirectToCustomServerPlugin - elif self._testMethodName == 'test_filter_by_upstream_host_plugin': - plugin = plugin_examples.FilterByUpstreamHostPlugin - elif self._testMethodName == 'test_cache_responses_plugin': - plugin = plugin_examples.CacheResponsesPlugin - elif self._testMethodName == 'test_man_in_the_middle_plugin': - plugin = plugin_examples.ManInTheMiddlePlugin + plugin = get_plugin_by_test_name(self._testMethodName) self.config.plugins = { b'ProtocolHandlerPlugin': [proxy.HttpProxyPlugin], b'HttpProxyBasePlugin': [plugin], } - self._conn = mock_fromfd.return_value + self._conn = mock.MagicMock(spec=socket.socket) + mock_fromfd.return_value = self._conn self.proxy = proxy.ProtocolHandler( self.fileno, self._addr, config=self.config) self.proxy.initialize() - @mock.patch('proxy.TcpServerConnection') - def test_modify_post_data_plugin(self, mock_server_conn: mock.Mock) -> None: - original = b'{"key": "value"}' - modified = b'{"key": "modified"}' + self.server = self.mock_server_conn.return_value + + self.server_ssl_connection = mock.MagicMock(spec=ssl.SSLSocket) + self.mock_ssl_context.return_value.wrap_socket.return_value = self.server_ssl_connection + self.client_ssl_connection = mock.MagicMock(spec=ssl.SSLSocket) + self.mock_ssl_wrap.return_value = self.client_ssl_connection + + def has_buffer() -> bool: + return cast(bool, self.server.queue.called) + + def closed() -> bool: + return not self.server.connect.called + + def mock_connection() -> Any: + if self.mock_ssl_context.return_value.wrap_socket.called: + return self.server_ssl_connection + return self._conn + + self.server.has_buffer.side_effect = has_buffer + type(self.server).closed = mock.PropertyMock(side_effect=closed) + type(self.server).connection = mock.PropertyMock(side_effect=mock_connection) - self._conn.recv.return_value = proxy.build_http_request( - b'POST', b'http://httpbin.org/post', - headers={ - b'Host': b'httpbin.org', - b'Content-Type': b'application/x-www-form-urlencoded', - b'Content-Length': proxy.bytes_(len(original)), - }, - body=original - ) self.mock_selector.return_value.select.side_effect = [ [(selectors.SelectorKey( fileobj=self._conn, fd=self._conn.fileno, events=selectors.EVENT_READ, + data=None), selectors.EVENT_READ)], + [(selectors.SelectorKey( + fileobj=self.client_ssl_connection, + fd=self.client_ssl_connection.fileno, + events=selectors.EVENT_READ, data=None), selectors.EVENT_READ)], ] + # Connect + self._conn.recv.return_value = proxy.build_http_request( + proxy.httpMethods.CONNECT, b'uni.corn:443' + ) self.proxy.run_once() - mock_server_conn.assert_called_with('httpbin.org', 80) - mock_server_conn.return_value.queue.assert_called_with( + + self.mock_popen.assert_called() + self.mock_server_conn.assert_called_once_with('uni.corn', 443) + self.server.connect.assert_called() + self.assertEqual(self.proxy.client.connection, self.client_ssl_connection) + self.assertEqual(self.server.connection, self.server_ssl_connection) + self._conn.send.assert_called_with( + proxy.HttpProxyPlugin.PROXY_TUNNEL_ESTABLISHED_RESPONSE_PKT + ) + + def test_modify_post_data_plugin(self) -> None: + original = b'{"key": "value"}' + modified = b'{"key": "modified"}' + self.client_ssl_connection.recv.return_value = proxy.build_http_request( + b'POST', b'/', + headers={ + b'Host': b'uni.corn', + b'Content-Type': b'application/x-www-form-urlencoded', + b'Content-Length': proxy.bytes_(len(original)), + }, + body=original + ) + self.proxy.run_once() + self.server.queue.assert_called_with( proxy.build_http_request( - b'POST', b'/post', + b'POST', b'/', headers={ - b'Host': b'httpbin.org', + b'Host': b'uni.corn', b'Content-Length': proxy.bytes_(len(modified)), b'Content-Type': b'application/json', - b'Via': b'1.1 %s' % proxy.PROXY_AGENT_HEADER_VALUE, }, body=modified ) ) + ''' @mock.patch('proxy.TcpServerConnection') def test_proposed_rest_api_plugin( self, mock_server_conn: mock.Mock) -> None: path = b'/v1/users/' - self._conn.recv.return_value = proxy.build_http_request( - b'GET', b'http://%s%s' % (plugin_examples.ProposedRestApiPlugin.API_SERVER, path), + self.client_ssl_connection.recv.return_value = proxy.build_http_request( + b'GET', path, headers={ b'Host': plugin_examples.ProposedRestApiPlugin.API_SERVER, } ) - self.mock_selector.return_value.select.side_effect = [ - [(selectors.SelectorKey( - fileobj=self._conn, - fd=self._conn.fileno, - events=selectors.EVENT_READ, - data=None), selectors.EVENT_READ)], ] self.proxy.run_once() - mock_server_conn.assert_not_called() self.assertEqual( self.proxy.client.buffer, @@ -1944,7 +2224,7 @@ def closed() -> bool: proxy.httpStatusCodes.OK, reason=b'OK', body=b'Hello from man in the middle') ) - + ''' class TestHttpRequestRejected(unittest.TestCase): From fa7d8bbd0e2bcd7d3c3aa41f3f4eb8c96e5fd6a7 Mon Sep 17 00:00:00 2001 From: Abhinav Singh Date: Sun, 13 Oct 2019 02:52:11 -0700 Subject: [PATCH 6/6] Add tests for plugin examples with TLS encryption enabled --- plugin_examples.py | 14 +++-- proxy.py | 7 ++- tests.py | 150 +++++++-------------------------------------- 3 files changed, 37 insertions(+), 134 deletions(-) diff --git a/plugin_examples.py b/plugin_examples.py index 289e0db58e..7ae2d73d31 100644 --- a/plugin_examples.py +++ b/plugin_examples.py @@ -50,7 +50,11 @@ class ProposedRestApiPlugin(proxy.HttpProxyBasePlugin): without need of an actual upstream REST API server. Returns proposed REST API mock responses to the client - without establishing upstream connection.""" + without establishing upstream connection. + + Note: This plugin won't work if your client is making + HTTPS connection to api.example.com. + """ API_SERVER = b'api.example.com' @@ -77,6 +81,11 @@ class ProposedRestApiPlugin(proxy.HttpProxyBasePlugin): } def before_upstream_connection(self, request: proxy.HttpParser) -> Optional[proxy.HttpParser]: + # Return None to disable establishing connection to upstream + # Most likely our api.example.com won't even exist under development scenario + return None + + def handle_client_request(self, request: proxy.HttpParser) -> Optional[proxy.HttpParser]: if request.host != self.API_SERVER: return request assert request.path @@ -95,9 +104,6 @@ def before_upstream_connection(self, request: proxy.HttpParser) -> Optional[prox )) return None - def handle_client_request(self, request: proxy.HttpParser) -> Optional[proxy.HttpParser]: - return request - def handle_upstream_chunk(self, chunk: bytes) -> bytes: return chunk diff --git a/proxy.py b/proxy.py index bd72c91d13..b0e2a05a13 100755 --- a/proxy.py +++ b/proxy.py @@ -1385,13 +1385,16 @@ def on_request_complete(self) -> Union[socket.socket, bool]: # Note: can raise HttpRequestRejected exception # Invoke plugin.before_upstream_connection + do_connect = True for plugin in self.plugins.values(): r = plugin.before_upstream_connection(self.request) if r is None: - return False + do_connect = False + break self.request = r - self.connect_upstream() + if do_connect: + self.connect_upstream() for plugin in self.plugins.values(): assert self.request is not None diff --git a/tests.py b/tests.py index 3a3be14190..8111f0f718 100644 --- a/tests.py +++ b/tests.py @@ -2025,9 +2025,23 @@ def mock_connection() -> Any: fileobj=self.client_ssl_connection, fd=self.client_ssl_connection.fileno, events=selectors.EVENT_READ, + data=None), selectors.EVENT_READ)], + [(selectors.SelectorKey( + fileobj=self.server_ssl_connection, + fd=self.server_ssl_connection.fileno, + events=selectors.EVENT_WRITE, + data=None), selectors.EVENT_WRITE)], + [(selectors.SelectorKey( + fileobj=self.server_ssl_connection, + fd=self.server_ssl_connection.fileno, + events=selectors.EVENT_READ, data=None), selectors.EVENT_READ)], ] # Connect + def send(raw: bytes) -> int: + return len(raw) + + self._conn.send.side_effect = send self._conn.recv.return_value = proxy.build_http_request( proxy.httpMethods.CONNECT, b'uni.corn:443' ) @@ -2041,6 +2055,7 @@ def mock_connection() -> Any: self._conn.send.assert_called_with( proxy.HttpProxyPlugin.PROXY_TUNNEL_ESTABLISHED_RESPONSE_PKT ) + self.assertEqual(self.proxy.client.buffer, b'') def test_modify_post_data_plugin(self) -> None: original = b'{"key": "value"}' @@ -2067,88 +2082,6 @@ def test_modify_post_data_plugin(self) -> None: ) ) - ''' - @mock.patch('proxy.TcpServerConnection') - def test_proposed_rest_api_plugin( - self, mock_server_conn: mock.Mock) -> None: - path = b'/v1/users/' - self.client_ssl_connection.recv.return_value = proxy.build_http_request( - b'GET', path, - headers={ - b'Host': plugin_examples.ProposedRestApiPlugin.API_SERVER, - } - ) - self.proxy.run_once() - mock_server_conn.assert_not_called() - self.assertEqual( - self.proxy.client.buffer, - proxy.build_http_response( - proxy.httpStatusCodes.OK, reason=b'OK', - headers={b'Content-Type': b'application/json'}, - body=proxy.bytes_(json.dumps(plugin_examples.ProposedRestApiPlugin.REST_API_SPEC[path])) - )) - - @mock.patch('proxy.TcpServerConnection') - def test_redirect_to_custom_server_plugin( - self, mock_server_conn: mock.Mock) -> None: - request = proxy.build_http_request( - b'GET', b'http://example.org/get', - headers={ - b'Host': b'example.org', - } - ) - self._conn.recv.return_value = request - self.mock_selector.return_value.select.side_effect = [ - [(selectors.SelectorKey( - fileobj=self._conn, - fd=self._conn.fileno, - events=selectors.EVENT_READ, - data=None), selectors.EVENT_READ)], ] - self.proxy.run_once() - - upstream = urlparse.urlsplit( - plugin_examples.RedirectToCustomServerPlugin.UPSTREAM_SERVER) - mock_server_conn.assert_called_with('localhost', 8899) - mock_server_conn.return_value.queue.assert_called_with( - proxy.build_http_request( - b'GET', upstream.path, - headers={ - b'Host': upstream.netloc, - b'Via': b'1.1 %s' % proxy.PROXY_AGENT_HEADER_VALUE, - } - ) - ) - - @mock.patch('proxy.TcpServerConnection') - def test_filter_by_upstream_host_plugin( - self, mock_server_conn: mock.Mock) -> None: - request = proxy.build_http_request( - b'GET', b'http://google.com/', - headers={ - b'Host': b'google.com', - } - ) - self._conn.recv.return_value = request - self.mock_selector.return_value.select.side_effect = [ - [(selectors.SelectorKey( - fileobj=self._conn, - fd=self._conn.fileno, - events=selectors.EVENT_READ, - data=None), selectors.EVENT_READ)], ] - self.proxy.run_once() - - mock_server_conn.assert_not_called() - self.assertEqual( - self.proxy.client.buffer, - proxy.build_http_response( - proxy.httpStatusCodes.I_AM_A_TEAPOT, - reason=b'I\'m a tea pot', - headers={ - proxy.PROXY_AGENT_HEADER_KEY: proxy.PROXY_AGENT_HEADER_VALUE - }, - ) - ) - @mock.patch('proxy.TcpServerConnection') def test_cache_responses_plugin( self, mock_server_conn: mock.Mock) -> None: @@ -2158,62 +2091,23 @@ def test_cache_responses_plugin( def test_man_in_the_middle_plugin( self, mock_server_conn: mock.Mock) -> None: request = proxy.build_http_request( - b'GET', b'http://super.secure/', + b'GET', b'/', headers={ - b'Host': b'super.secure', + b'Host': b'uni.corn', } ) - self._conn.recv.return_value = request - - server = mock_server_conn.return_value - server.connect.return_value = True - - def has_buffer() -> bool: - return cast(bool, server.queue.called) - - def closed() -> bool: - return not server.connect.called - - server.has_buffer.side_effect = has_buffer - type(server).closed = mock.PropertyMock(side_effect=closed) - - self.mock_selector.return_value.select.side_effect = [ - [(selectors.SelectorKey( - fileobj=self._conn, - fd=self._conn.fileno, - events=selectors.EVENT_READ, - data=None), selectors.EVENT_READ)], - [(selectors.SelectorKey( - fileobj=server.connection, - fd=server.connection.fileno, - events=selectors.EVENT_WRITE, - data=None), selectors.EVENT_WRITE)], - [(selectors.SelectorKey( - fileobj=server.connection, - fd=server.connection.fileno, - events=selectors.EVENT_READ, - data=None), selectors.EVENT_READ)], ] + self.client_ssl_connection.recv.return_value = request # Client read self.proxy.run_once() - mock_server_conn.assert_called_with('super.secure', 80) - server.connect.assert_called_once() - queued_request = \ - proxy.build_http_request( - b'GET', b'/', - headers={ - b'Host': b'super.secure', - b'Via': b'1.1 %s' % proxy.PROXY_AGENT_HEADER_VALUE - } - ) - server.queue.assert_called_once_with(queued_request) + self.server.queue.assert_called_once_with(request) # Server write self.proxy.run_once() - server.flush.assert_called_once() + self.server.flush.assert_called_once() # Server read - server.recv.return_value = \ + self.server.recv.return_value = \ proxy.build_http_response( proxy.httpStatusCodes.OK, reason=b'OK', body=b'Original Response From Upstream') @@ -2224,7 +2118,7 @@ def closed() -> bool: proxy.httpStatusCodes.OK, reason=b'OK', body=b'Hello from man in the middle') ) - ''' + class TestHttpRequestRejected(unittest.TestCase):