# -*- coding: utf-8 -*- """ proxy.py ~~~~~~~~ ⚡⚡⚡ Fast, Lightweight, Pluggable, TLS interception capable proxy server focused on Network monitoring, controls & Application development, testing, debugging. :copyright: (c) 2013-present by Abhinav Singh and contributors. :license: BSD, see LICENSE for more details. """ import unittest import selectors from unittest import mock from proxy.common.constants import DEFAULT_HTTP_PORT from proxy.common.flags import Flags from proxy.core.connection import TcpClientConnection from proxy.http.proxy import HttpProxyPlugin from proxy.http.handler import HttpProtocolHandler from proxy.http.exception import HttpProtocolException from proxy.common.utils import build_http_request class TestHttpProxyPlugin(unittest.TestCase): @mock.patch('selectors.DefaultSelector') @mock.patch('socket.fromfd') def setUp(self, mock_fromfd: mock.Mock, mock_selector: mock.Mock) -> None: self.mock_fromfd = mock_fromfd self.mock_selector = mock_selector self.fileno = 10 self._addr = ('127.0.0.1', 54382) self.flags = Flags() self.plugin = mock.MagicMock() self.flags.plugins = { b'HttpProtocolHandlerPlugin': [HttpProxyPlugin], b'HttpProxyBasePlugin': [self.plugin] } self._conn = mock_fromfd.return_value self.protocol_handler = HttpProtocolHandler( TcpClientConnection(self._conn, self._addr), flags=self.flags) self.protocol_handler.initialize() def test_proxy_plugin_initialized(self) -> None: self.plugin.assert_called() @mock.patch('proxy.http.proxy.server.TcpServerConnection') def test_proxy_plugin_on_and_before_upstream_connection( self, mock_server_conn: mock.Mock) -> None: self.plugin.return_value.before_upstream_connection.side_effect = lambda r: r self.plugin.return_value.handle_client_request.side_effect = lambda r: r self._conn.recv.return_value = build_http_request( b'GET', b'http://upstream.host/not-found.html', headers={ b'Host': b'upstream.host' }) 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.protocol_handler.run_once() mock_server_conn.assert_called_with('upstream.host', DEFAULT_HTTP_PORT) self.plugin.return_value.before_upstream_connection.assert_called() self.plugin.return_value.handle_client_request.assert_called() @mock.patch('proxy.http.proxy.server.TcpServerConnection') def test_proxy_plugin_before_upstream_connection_can_teardown( self, mock_server_conn: mock.Mock) -> None: self.plugin.return_value.before_upstream_connection.side_effect = HttpProtocolException() self._conn.recv.return_value = build_http_request( b'GET', b'http://upstream.host/not-found.html', headers={ b'Host': b'upstream.host' }) 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.protocol_handler.run_once() self.plugin.return_value.before_upstream_connection.assert_called() mock_server_conn.assert_not_called()