Ejemplo n.º 1
0
 def test_main_py3_runs(
     self,
     mock_is_py2: mock.Mock,
     mock_print: mock.Mock,
 ) -> None:
     mock_is_py2.return_value = False
     FlagParser.initialize()
     mock_is_py2.assert_called()
     mock_print.assert_not_called()
Ejemplo n.º 2
0
 def test_main_py2_exit(
     self,
     mock_is_py2: mock.Mock,
     mock_print: mock.Mock,
 ) -> None:
     mock_is_py2.return_value = True
     with self.assertRaises(SystemExit) as e:
         FlagParser.initialize()
     mock_print.assert_called_with(PY2_DEPRECATION_MESSAGE)
     self.assertEqual(e.exception.code, 1)
     mock_is_py2.assert_called()
Ejemplo n.º 3
0
 def _setUp(self, mocker: MockerFixture) -> None:
     self.mock_fromfd = mocker.patch('socket.fromfd')
     self.mock_selector = mocker.patch('selectors.DefaultSelector')
     self.fileno = 10
     self._addr = ('127.0.0.1', 54382)
     self._conn = self.mock_fromfd.return_value
     # Setup a static directory
     self.static_server_dir = os.path.join(tempfile.gettempdir(), 'static')
     self.index_file_path = os.path.join(
         self.static_server_dir, 'index.html',
     )
     self.html_file_content = b'''<html><head></head><body><h1>Proxy.py Testing</h1></body></html>'''
     os.makedirs(self.static_server_dir, exist_ok=True)
     with open(self.index_file_path, 'wb') as f:
         f.write(self.html_file_content)
     #
     flags = FlagParser.initialize(
         enable_static_server=True,
         static_server_dir=self.static_server_dir,
         threaded=True,
     )
     flags.plugins = Plugins.load([
         bytes_(PLUGIN_HTTP_PROXY),
         bytes_(PLUGIN_WEB_SERVER),
     ])
     self.protocol_handler = HttpProtocolHandler(
         HttpClientConnection(self._conn, self._addr),
         flags=flags,
     )
     self.protocol_handler.initialize()
Ejemplo n.º 4
0
    def test_unix_path_listener(self, mock_socket: mock.Mock,
                                mock_remove: mock.Mock) -> None:
        sock = mock_socket.return_value
        sock_path = os.path.join(tempfile.gettempdir(), 'proxy.sock')
        flags = FlagParser.initialize(unix_socket_path=sock_path)
        listener = Listener(flags)
        listener.setup()

        mock_socket.assert_called_with(
            socket.AF_UNIX,
            socket.SOCK_STREAM,
        )
        self.assertEqual(sock.setsockopt.call_count, 2)
        self.assertEqual(
            sock.setsockopt.call_args_list[0][0],
            (socket.SOL_SOCKET, socket.SO_REUSEADDR, 1),
        )
        self.assertEqual(
            sock.setsockopt.call_args_list[1][0],
            (socket.IPPROTO_TCP, socket.TCP_NODELAY, 1),
        )
        sock.bind.assert_called_with(sock_path)
        sock.listen.assert_called_with(flags.backlog)
        sock.setblocking.assert_called_with(False)

        listener.shutdown()
        mock_remove.assert_called_once_with(sock_path)
        sock.close.assert_called_once()
Ejemplo n.º 5
0
 def test_multi_listener_on_ports(
         self,
         mock_unix_listener: mock.Mock,
         mock_tcp_listener: mock.Mock,
 ) -> None:
     flags = FlagParser.initialize(
         ['--ports', '9000', '--ports', '9001'],
         port=0,
     )
     with ListenerPool(flags=flags) as pool:
         mock_unix_listener.assert_not_called()
         self.assertEqual(len(pool.pool), 3)
         self.assertEqual(mock_tcp_listener.call_count, 3)
         self.assertEqual(
             mock_tcp_listener.call_args_list[0][1]['flags'],
             flags,
         )
         self.assertEqual(
             mock_tcp_listener.call_args_list[1][1]['flags'],
             flags,
         )
         self.assertEqual(
             mock_tcp_listener.call_args_list[1][1]['port'],
             9000,
         )
         self.assertEqual(
             mock_tcp_listener.call_args_list[2][1]['flags'],
             flags,
         )
         self.assertEqual(
             mock_tcp_listener.call_args_list[2][1]['port'],
             9001,
         )
Ejemplo n.º 6
0
    def _setUp(self, request: Any, mocker: MockerFixture) -> None:
        self.mock_fromfd = mocker.patch('socket.fromfd')
        self.mock_selector = mocker.patch('selectors.DefaultSelector')
        self.mock_server_conn = mocker.patch(
            'proxy.http.proxy.server.TcpServerConnection', )

        self.fileno = 10
        self._addr = ('127.0.0.1', 54382)
        adblock_json_path = Path(
            __file__,
        ).parent.parent.parent / "proxy" / "plugin" / "adblock.json"
        self.flags = FlagParser.initialize(
            input_args=[
                "--filtered-url-regex-config",
                str(adblock_json_path),
            ],
            threaded=True,
        )
        self.plugin = mock.MagicMock()

        plugin = get_plugin_by_test_name(request.param)

        self.flags.plugins = {
            b'HttpProtocolHandlerPlugin': [HttpProxyPlugin],
            b'HttpProxyBasePlugin': [plugin],
        }
        self._conn = self.mock_fromfd.return_value
        self.protocol_handler = HttpProtocolHandler(
            HttpClientConnection(self._conn, self._addr),
            flags=self.flags,
        )
        self.protocol_handler.initialize()
Ejemplo n.º 7
0
 def test_setup_and_teardown(self, mock_socket: mock.Mock) -> None:
     sock = mock_socket.return_value
     flags = FlagParser.initialize(port=0)
     with TcpSocketListener(flags=flags) as listener:
         mock_socket.assert_called_with(
             socket.AF_INET6
             if flags.hostname.version == 6 else socket.AF_INET,
             socket.SOCK_STREAM,
         )
         self.assertEqual(sock.setsockopt.call_count, 2)
         self.assertEqual(
             sock.setsockopt.call_args_list[0][0],
             (socket.SOL_SOCKET, socket.SO_REUSEADDR, 1),
         )
         self.assertEqual(
             sock.setsockopt.call_args_list[1][0],
             (socket.IPPROTO_TCP, socket.TCP_NODELAY, 1),
         )
         sock.bind.assert_called_with((str(flags.hostname), 0), )
         sock.listen.assert_called_with(flags.backlog)
         sock.setblocking.assert_called_with(False)
         self.assertEqual(
             listener.fileno(),
             mock_socket.return_value.fileno.return_value,
         )
     sock.close.assert_called_once()
Ejemplo n.º 8
0
 async def test_proxy_authentication_failed(self) -> None:
     self._conn = self.mock_fromfd.return_value
     mock_selector_for_client_read(self)
     flags = FlagParser.initialize(
         auth_code=base64.b64encode(b'user:pass'),
         threaded=True,
     )
     flags.plugins = Plugins.load([
         bytes_(PLUGIN_HTTP_PROXY),
         bytes_(PLUGIN_WEB_SERVER),
         bytes_(PLUGIN_PROXY_AUTH),
     ])
     self.protocol_handler = HttpProtocolHandler(
         HttpClientConnection(self._conn, self._addr),
         flags=flags,
     )
     self.protocol_handler.initialize()
     self._conn.recv.return_value = CRLF.join([
         b'GET http://abhinavsingh.com HTTP/1.1',
         b'Host: abhinavsingh.com',
         CRLF,
     ])
     await self.protocol_handler._run_once()
     self.assertEqual(
         self.protocol_handler.work.buffer[0],
         PROXY_AUTH_FAILED_RESPONSE_PKT,
     )
Ejemplo n.º 9
0
 def _setUp(self, request: Any, mocker: MockerFixture) -> None:
     self.mock_fromfd = mocker.patch('socket.fromfd')
     self.mock_selector = mocker.patch('selectors.DefaultSelector')
     self.fileno = 10
     self._addr = ('127.0.0.1', 54382)
     self._conn = self.mock_fromfd.return_value
     self.pac_file = request.param
     if isinstance(self.pac_file, str):
         with open(self.pac_file, 'rb') as f:
             self.expected_response = f.read()
     else:
         self.expected_response = PAC_FILE_CONTENT
     self.flags = FlagParser.initialize(
         pac_file=self.pac_file, threaded=True,
     )
     self.flags.plugins = Plugins.load([
         bytes_(PLUGIN_HTTP_PROXY),
         bytes_(PLUGIN_WEB_SERVER),
         bytes_(PLUGIN_PAC_FILE),
     ])
     self.protocol_handler = HttpProtocolHandler(
         HttpClientConnection(self._conn, self._addr),
         flags=self.flags,
     )
     self.protocol_handler.initialize()
     self._conn.recv.return_value = CRLF.join([
         b'GET / HTTP/1.1',
         CRLF,
     ])
     mock_selector_for_client_read(self)
Ejemplo n.º 10
0
 def test_load_plugin_from_class(self) -> None:
     self.flags = FlagParser.initialize(
         [],
         plugins=[
             CacheResponsesPlugin,
         ],
     )
     self.assert_plugins({'HttpProxyBasePlugin': [CacheResponsesPlugin]})
Ejemplo n.º 11
0
def main() -> None:
    try:
        flags = FlagParser.initialize(
            sys.argv[2:] + ['--disable-http-proxy'],
            work_klass='proxy.core.work.task.TaskHandler',
        )
        globals()['start_%s' % sys.argv[1]](flags)
    except KeyboardInterrupt:
        pass
Ejemplo n.º 12
0
 def test_unique_plugin_from_args(self) -> None:
     self.flags = FlagParser.initialize([
         '--plugins',
         PLUGIN_HTTP_PROXY,
     ])
     self.assert_plugins({
         'HttpProtocolHandlerPlugin': [
             HttpProxyPlugin,
         ],
     })
Ejemplo n.º 13
0
 def test_load_plugins_from_args(self) -> None:
     self.flags = FlagParser.initialize([
         '--plugins',
         'proxy.plugin.CacheResponsesPlugin,proxy.plugin.FilterByUpstreamHostPlugin',
     ])
     self.assert_plugins({
         'HttpProxyBasePlugin': [
             CacheResponsesPlugin,
             FilterByUpstreamHostPlugin,
         ],
     })
Ejemplo n.º 14
0
 def test_unique_plugin_from_bytes(self) -> None:
     self.flags = FlagParser.initialize(
         [],
         plugins=[
             bytes_(PLUGIN_HTTP_PROXY),
         ],
     )
     self.assert_plugins({
         'HttpProtocolHandlerPlugin': [
             HttpProxyPlugin,
         ],
     })
Ejemplo n.º 15
0
 def test_setup_and_teardown(
         self,
         mock_unix_listener: mock.Mock,
         mock_tcp_listener: mock.Mock,
 ) -> None:
     flags = FlagParser.initialize(port=0)
     with ListenerPool(flags=flags) as pool:
         mock_tcp_listener.assert_called_once_with(flags=flags)
         mock_unix_listener.assert_not_called()
         mock_tcp_listener.return_value.setup.assert_called_once()
         self.assertEqual(pool.pool[0], mock_tcp_listener.return_value)
     mock_tcp_listener.return_value.shutdown.assert_called_once()
Ejemplo n.º 16
0
 def test_unique_plugin_from_class(self) -> None:
     self.flags = FlagParser.initialize(
         [],
         plugins=[
             HttpProxyPlugin,
         ],
     )
     self.assert_plugins({
         'HttpProtocolHandlerPlugin': [
             HttpProxyPlugin,
         ],
     })
Ejemplo n.º 17
0
 def test_unix_socket_listener(
         self,
         mock_unix_listener: mock.Mock,
         mock_tcp_listener: mock.Mock,
 ) -> None:
     flags = FlagParser.initialize(
         unix_socket_path=os.path.join(tempfile.gettempdir(), 'proxy.sock'),
     )
     with ListenerPool(flags=flags) as pool:
         mock_unix_listener.assert_called_once_with(flags=flags)
         mock_tcp_listener.assert_not_called()
         mock_unix_listener.return_value.setup.assert_called_once()
         self.assertEqual(pool.pool[0], mock_unix_listener.return_value)
     mock_unix_listener.return_value.shutdown.assert_called_once()
Ejemplo n.º 18
0
 def test_load_plugins_from_bytes_and_class(self) -> None:
     self.flags = FlagParser.initialize(
         [],
         plugins=[
             CacheResponsesPlugin,
             b'proxy.plugin.FilterByUpstreamHostPlugin',
         ],
     )
     self.assert_plugins({
         'HttpProxyBasePlugin': [
             CacheResponsesPlugin,
             FilterByUpstreamHostPlugin,
         ],
     })
Ejemplo n.º 19
0
    def _setUp(self, mocker: MockerFixture) -> None:
        self.mock_fromfd = mocker.patch('socket.fromfd')
        self.mock_selector = mocker.patch('selectors.DefaultSelector')

        self.fileno = 10
        self._addr = ('127.0.0.1', 54382)
        self._conn = self.mock_fromfd.return_value

        self.flags = FlagParser.initialize(threaded=True)

        self.handler = SocksProtocolHandler(
            SocksClientConnection(conn=self._conn, addr=self._addr),
            flags=self.flags,
        )
        self.handler.initialize()
Ejemplo n.º 20
0
    async def test_authenticated_proxy_http_get(self) -> None:
        self._conn = self.mock_fromfd.return_value
        mock_selector_for_client_read(self)

        server = self.mock_server_connection.return_value
        server.connect.return_value = True
        server.buffer_size.return_value = 0

        flags = FlagParser.initialize(
            auth_code=base64.b64encode(b'user:pass'),
            threaded=True,
        )
        flags.plugins = Plugins.load([
            bytes_(PLUGIN_HTTP_PROXY),
            bytes_(PLUGIN_WEB_SERVER),
        ])

        self.protocol_handler = HttpProtocolHandler(
            HttpClientConnection(self._conn, self._addr),
            flags=flags,
        )
        self.protocol_handler.initialize()
        assert self.http_server_port is not None

        self._conn.recv.return_value = b'GET http://localhost:%d HTTP/1.1' % self.http_server_port
        await self.protocol_handler._run_once()
        self.assertEqual(
            self.protocol_handler.request.state,
            httpParserStates.INITIALIZED,
        )

        self._conn.recv.return_value = CRLF
        await self.protocol_handler._run_once()
        self.assertEqual(
            self.protocol_handler.request.state,
            httpParserStates.LINE_RCVD,
        )

        assert self.http_server_port is not None
        self._conn.recv.return_value = CRLF.join([
            b'User-Agent: proxy.py/%s' % bytes_(__version__),
            b'Host: localhost:%d' % self.http_server_port,
            b'Accept: */*',
            httpHeaders.PROXY_CONNECTION + b': Keep-Alive',
            httpHeaders.PROXY_AUTHORIZATION + b': Basic dXNlcjpwYXNz',
            CRLF,
        ])
        await self.assert_data_queued(server)
Ejemplo n.º 21
0
 def _setUp(self, mocker: MockerFixture) -> None:
     self.mock_fromfd = mocker.patch('socket.fromfd')
     self.mock_selector = mocker.patch('selectors.DefaultSelector')
     self.fileno = 10
     self._addr = ('127.0.0.1', 54382)
     self._conn = self.mock_fromfd.return_value
     self.flags = FlagParser.initialize(threaded=True)
     self.flags.plugins = Plugins.load([
         bytes_(PLUGIN_HTTP_PROXY),
         bytes_(PLUGIN_WEB_SERVER),
     ])
     self.protocol_handler = HttpProtocolHandler(
         HttpClientConnection(self._conn, self._addr),
         flags=self.flags,
     )
     self.protocol_handler.initialize()
Ejemplo n.º 22
0
def test_on_client_connection_called_on_teardown(mocker: MockerFixture) -> None:
    plugin = mocker.MagicMock()
    mock_fromfd = mocker.patch('socket.fromfd')
    flags = FlagParser.initialize(threaded=True)
    flags.plugins = {b'HttpProtocolHandlerPlugin': [plugin]}
    _conn = mock_fromfd.return_value
    _addr = ('127.0.0.1', 54382)
    protocol_handler = HttpProtocolHandler(
        HttpClientConnection(_conn, _addr),
        flags=flags,
    )
    protocol_handler.initialize()
    plugin.assert_not_called()
    mock_run_once = mocker.patch.object(protocol_handler, '_run_once')
    mock_run_once.return_value = True
    protocol_handler.run()
    assert _conn.closed
Ejemplo n.º 23
0
 def setUp(self) -> None:
     self.acceptor_id = 1
     self.pipe = multiprocessing.Pipe()
     self.work_klass = mock.MagicMock()
     self.flags = FlagParser.initialize(
         threaded=True,
         work_klass=self.work_klass,
         local_executor=0,
     )
     self.acceptor = Acceptor(
         idd=self.acceptor_id,
         fd_queue=self.pipe[1],
         flags=self.flags,
         lock=multiprocessing.Lock(),
         executor_queues=[],
         executor_pids=[],
         executor_locks=[],
     )
Ejemplo n.º 24
0
    def _setUp(self, mocker: MockerFixture) -> None:
        self.mock_fromfd = mocker.patch('socket.fromfd')
        self.mock_selector = mocker.patch('selectors.DefaultSelector')
        self.mock_server_conn = mocker.patch(
            'proxy.http.proxy.server.TcpServerConnection', )

        self.fileno = 10
        self._addr = ('127.0.0.1', 54382)
        self.flags = FlagParser.initialize(
            ["--basic-auth", "user:pass"],
            threaded=True,
        )
        self._conn = self.mock_fromfd.return_value
        self.protocol_handler = HttpProtocolHandler(
            HttpClientConnection(self._conn, self._addr),
            flags=self.flags,
        )
        self.protocol_handler.initialize()
Ejemplo n.º 25
0
    def _setUp(self, mocker: MockerFixture) -> None:
        self.mock_server_conn = mocker.patch(
            'proxy.http.proxy.server.TcpServerConnection',
        )
        self.mock_selector = mocker.patch('selectors.DefaultSelector')
        self.mock_fromfd = mocker.patch('socket.fromfd')

        self.fileno = 10
        self._addr = ('127.0.0.1', 54382)
        self.flags = FlagParser.initialize(threaded=True)
        self.plugin = mocker.MagicMock()
        self.flags.plugins = {
            b'HttpProtocolHandlerPlugin': [HttpProxyPlugin],
            b'HttpProxyBasePlugin': [self.plugin],
        }
        self._conn = self.mock_fromfd.return_value
        self.protocol_handler = HttpProtocolHandler(
            HttpClientConnection(self._conn, self._addr),
            flags=self.flags,
        )
        self.protocol_handler.initialize()
Ejemplo n.º 26
0
    async def test_authenticated_proxy_http_tunnel(self) -> None:
        server = self.mock_server_connection.return_value
        server.connect.return_value = True
        server.buffer_size.return_value = 0
        self._conn = self.mock_fromfd.return_value
        self.mock_selector_for_client_read_and_server_write(server)

        flags = FlagParser.initialize(
            auth_code=base64.b64encode(b'user:pass'),
            threaded=True,
        )
        flags.plugins = Plugins.load([
            bytes_(PLUGIN_HTTP_PROXY),
            bytes_(PLUGIN_WEB_SERVER),
        ])

        self.protocol_handler = HttpProtocolHandler(
            HttpClientConnection(self._conn, self._addr),
            flags=flags,
        )
        self.protocol_handler.initialize()

        assert self.http_server_port is not None
        self._conn.recv.return_value = CRLF.join([
            b'CONNECT localhost:%d HTTP/1.1' % self.http_server_port,
            b'Host: localhost:%d' % self.http_server_port,
            b'User-Agent: proxy.py/%s' % bytes_(__version__),
            httpHeaders.PROXY_CONNECTION + b': Keep-Alive',
            httpHeaders.PROXY_AUTHORIZATION + b': Basic dXNlcjpwYXNz',
            CRLF,
        ])
        await self.assert_tunnel_response(server)
        self.protocol_handler.work.flush()
        await self.assert_data_queued_to_server(server)

        await self.protocol_handler._run_once()
        server.flush.assert_called_once()
Ejemplo n.º 27
0
 async def test_default_web_server_returns_404(self) -> None:
     self._conn = self.mock_fromfd.return_value
     self.mock_selector.return_value.select.return_value = [
         (
             selectors.SelectorKey(
                 fileobj=self._conn.fileno(),
                 fd=self._conn.fileno(),
                 events=selectors.EVENT_READ,
                 data=None,
             ),
             selectors.EVENT_READ,
         ),
     ]
     flags = FlagParser.initialize(threaded=True)
     flags.plugins = Plugins.load([
         bytes_(PLUGIN_HTTP_PROXY),
         bytes_(PLUGIN_WEB_SERVER),
     ])
     self.protocol_handler = HttpProtocolHandler(
         HttpClientConnection(self._conn, self._addr),
         flags=flags,
     )
     self.protocol_handler.initialize()
     self._conn.recv.return_value = CRLF.join([
         b'GET /hello HTTP/1.1',
         CRLF,
     ])
     await self.protocol_handler._run_once()
     self.assertEqual(
         self.protocol_handler.request.state,
         httpParserStates.COMPLETE,
     )
     self.assertEqual(
         self.protocol_handler.work.buffer[0],
         NOT_FOUND_RESPONSE_PKT,
     )
Ejemplo n.º 28
0
    async def test_e2e(self, mocker: MockerFixture) -> None:
        host, port = uuid.uuid4().hex, 443
        netloc = '{0}:{1}'.format(host, port)

        self.mock_fromfd = mocker.patch('socket.fromfd')
        self.mock_selector = mocker.patch('selectors.DefaultSelector')
        self.mock_sign_csr = mocker.patch('proxy.http.proxy.server.sign_csr')
        self.mock_gen_csr = mocker.patch('proxy.http.proxy.server.gen_csr')
        self.mock_gen_public_key = mocker.patch(
            'proxy.http.proxy.server.gen_public_key', )
        self.mock_server_conn = mocker.patch(
            'proxy.http.proxy.server.TcpServerConnection', )
        self.mock_sign_csr.return_value = True
        self.mock_gen_csr.return_value = True
        self.mock_gen_public_key.return_value = True

        # Used for server side wrapping
        self.mock_ssl_context = mocker.patch('ssl.create_default_context')
        upstream_tls_sock = mock.MagicMock(spec=ssl.SSLSocket)
        self.mock_ssl_context.return_value.wrap_socket.return_value = upstream_tls_sock

        # Used for client wrapping
        self.mock_ssl_wrap = mocker.patch('ssl.wrap_socket')
        client_tls_sock = mock.MagicMock(spec=ssl.SSLSocket)
        self.mock_ssl_wrap.return_value = client_tls_sock

        plain_connection = mock.MagicMock(spec=socket.socket)

        def mock_connection() -> Any:
            if self.mock_ssl_context.return_value.wrap_socket.called:
                return upstream_tls_sock
            return plain_connection

        # Do not mock the original wrap method
        self.mock_server_conn.return_value.wrap.side_effect = \
            lambda x, y, as_non_blocking: TcpServerConnection.wrap(
                self.mock_server_conn.return_value, x, y, as_non_blocking=as_non_blocking,
            )

        type(self.mock_server_conn.return_value).connection = \
            mock.PropertyMock(side_effect=mock_connection)

        type(self.mock_server_conn.return_value).closed = \
            mock.PropertyMock(return_value=False)

        self.fileno = 10
        self._addr = ('127.0.0.1', 54382)
        self.flags = FlagParser.initialize(
            ca_cert_file='ca-cert.pem',
            ca_key_file='ca-key.pem',
            ca_signing_key_file='ca-signing-key.pem',
            threaded=True,
        )
        self.assertTrue(tls_interception_enabled(self.flags))
        # In this test we enable a mock http protocol handler plugin
        # and a mock http proxy plugin.  Internally, http protocol
        # handler will only initialize proxy plugin as we'll never
        # make any other request.
        self.plugin = mock.MagicMock()
        self.proxy_plugin = mock.MagicMock()
        self.flags.plugins = {
            b'HttpProtocolHandlerPlugin': [self.plugin, HttpProxyPlugin],
            b'HttpProxyBasePlugin': [self.proxy_plugin],
        }
        self._conn = self.mock_fromfd.return_value
        self.protocol_handler = HttpProtocolHandler(
            HttpClientConnection(self._conn, self._addr),
            flags=self.flags,
        )
        self.protocol_handler.initialize()

        self.plugin.assert_not_called()
        self.proxy_plugin.assert_not_called()

        # Mock a CONNECT request followed by a GET request
        # from client connection
        headers = {
            b'Host': bytes_(netloc),
        }
        connect_request = build_http_request(
            httpMethods.CONNECT,
            bytes_(netloc),
            headers=headers,
            no_ua=True,
        )
        self._conn.recv.return_value = connect_request
        get_request = build_http_request(
            httpMethods.GET,
            b'/',
            headers=headers,
        )
        client_tls_sock.recv.return_value = get_request

        T = TypeVar('T')  # noqa: N806

        async def asyncReturn(val: T) -> T:
            return val

        # Prepare mocked HttpProxyBasePlugin
        # 1. Mock descriptor mixin methods
        #
        # NOTE: We need multiple async result otherwise
        # we will end up with cannot await on already
        # awaited coroutine.
        self.proxy_plugin.return_value.get_descriptors.side_effect = \
            [asyncReturn(([], [])), asyncReturn(([], []))]
        self.proxy_plugin.return_value.write_to_descriptors.side_effect = \
            [asyncReturn(False), asyncReturn(False)]
        self.proxy_plugin.return_value.read_from_descriptors.side_effect = \
            [asyncReturn(False), asyncReturn(False)]
        # 2. Mock plugin lifecycle methods
        self.proxy_plugin.return_value.resolve_dns.return_value = None, None
        self.proxy_plugin.return_value.before_upstream_connection.side_effect = lambda r: r
        self.proxy_plugin.return_value.handle_client_data.side_effect = lambda r: r
        self.proxy_plugin.return_value.handle_client_request.side_effect = lambda r: r
        self.proxy_plugin.return_value.handle_upstream_chunk.side_effect = lambda r: r
        self.proxy_plugin.return_value.on_upstream_connection_close.return_value = None
        self.proxy_plugin.return_value.on_access_log.side_effect = lambda r: r
        self.proxy_plugin.return_value.do_intercept.return_value = True

        self.mock_selector.return_value.select.side_effect = [
            # Trigger read on plain socket
            [(
                selectors.SelectorKey(
                    fileobj=self._conn.fileno(),
                    fd=self._conn.fileno(),
                    events=selectors.EVENT_READ,
                    data=None,
                ),
                selectors.EVENT_READ,
            )],
            # Trigger read on encrypted socket
            [(
                selectors.SelectorKey(
                    fileobj=client_tls_sock.fileno(),
                    fd=client_tls_sock.fileno(),
                    events=selectors.EVENT_READ,
                    data=None,
                ),
                selectors.EVENT_READ,
            )],
        ]

        await self.protocol_handler._run_once()

        # Assert correct plugin was initialized
        self.plugin.assert_not_called()
        self.proxy_plugin.assert_called_once()
        self.assertEqual(self.proxy_plugin.call_args[0][1], self.flags)
        # Actual call arg must be `_conn` object
        # but because internally the reference is updated
        # we assert it against `mock_ssl_wrap` which is
        # called during proxy plugin initialization
        # for interception
        self.assertEqual(
            self.proxy_plugin.call_args[0][2].connection,
            client_tls_sock,
        )

        # Invoked lifecycle callbacks
        self.proxy_plugin.return_value.resolve_dns.assert_called_once_with(
            host,
            port,
        )
        self.proxy_plugin.return_value.before_upstream_connection.assert_called(
        )
        self.proxy_plugin.return_value.handle_client_request.assert_called_once(
        )
        self.proxy_plugin.return_value.do_intercept.assert_called_once()
        # All the invoked lifecycle callbacks will receive the CONNECT request
        # packet with / as the path
        callback_request: HttpParser = \
            self.proxy_plugin.return_value.before_upstream_connection.call_args_list[0][0][0]
        callback_request1: HttpParser = \
            self.proxy_plugin.return_value.handle_client_request.call_args_list[0][0][0]
        callback_request2: HttpParser = \
            self.proxy_plugin.return_value.do_intercept.call_args_list[0][0][0]
        self.assertEqual(callback_request.host, bytes_(host))
        self.assertEqual(callback_request.port, 443)
        self.assertEqual(callback_request.header(b'Host'), headers[b'Host'])
        assert callback_request._url
        self.assertEqual(callback_request._url.remainder, None)
        self.assertEqual(callback_request.method, httpMethods.CONNECT)
        self.assertEqual(callback_request.is_https_tunnel, True)
        self.assertEqual(callback_request.build(), callback_request1.build())
        self.assertEqual(callback_request.build(), callback_request2.build())
        # Lifecycle callbacks not invoked
        self.proxy_plugin.return_value.handle_client_data.assert_not_called()
        self.proxy_plugin.return_value.handle_upstream_chunk.assert_not_called(
        )
        self.proxy_plugin.return_value.on_upstream_connection_close.assert_not_called(
        )
        self.proxy_plugin.return_value.on_access_log.assert_not_called()

        self.mock_server_conn.assert_called_with(host, port)
        self.mock_server_conn.return_value.connection.setblocking.assert_called_with(
            False, )

        self.mock_ssl_context.assert_called_with(
            ssl.Purpose.SERVER_AUTH,
            cafile=str(DEFAULT_CA_FILE),
        )
        self.assertEqual(plain_connection.setblocking.call_count, 2)
        self.mock_ssl_context.return_value.wrap_socket.assert_called_with(
            plain_connection,
            server_hostname=host,
        )
        self.assertEqual(self.mock_sign_csr.call_count, 1)
        self.assertEqual(self.mock_gen_csr.call_count, 1)
        self.assertEqual(self.mock_gen_public_key.call_count, 1)
        self.assertEqual(upstream_tls_sock.setblocking.call_count, 1)
        self.assertEqual(
            self.mock_server_conn.return_value._conn,
            upstream_tls_sock,
        )
        self._conn.send.assert_called_with(
            PROXY_TUNNEL_ESTABLISHED_RESPONSE_PKT, )
        assert self.flags.ca_cert_dir is not None
        self.mock_ssl_wrap.assert_called_with(
            self._conn,
            server_side=True,
            keyfile=self.flags.ca_signing_key_file,
            certfile=HttpProxyPlugin.generated_cert_file_path(
                self.flags.ca_cert_dir,
                host,
            ),
            ssl_version=ssl.PROTOCOL_TLS,
        )
        self.assertEqual(self._conn.setblocking.call_count, 2)
        self.assertEqual(
            self.protocol_handler.work.connection,
            client_tls_sock,
        )

        # Assert connection references for all other plugins is updated
        self.assertEqual(
            self.proxy_plugin.return_value.client._conn,
            client_tls_sock,
        )

        # Now process the GET request
        await self.protocol_handler._run_once()
        self.plugin.assert_not_called()
        self.proxy_plugin.assert_called_once()

        # Lifecycle callbacks still not invoked
        self.proxy_plugin.return_value.handle_client_data.assert_not_called()
        self.proxy_plugin.return_value.handle_upstream_chunk.assert_not_called(
        )
        self.proxy_plugin.return_value.on_upstream_connection_close.assert_not_called(
        )
        self.proxy_plugin.return_value.on_access_log.assert_not_called()
        # Only handle client request lifecycle must be called again
        self.proxy_plugin.return_value.resolve_dns.assert_called_once_with(
            host,
            port,
        )
        self.proxy_plugin.return_value.before_upstream_connection.assert_called(
        )
        self.assertEqual(
            self.proxy_plugin.return_value.handle_client_request.call_count,
            2,
        )
        self.proxy_plugin.return_value.do_intercept.assert_called_once()

        callback_request = \
            self.proxy_plugin.return_value.handle_client_request.call_args_list[1][0][0]
        self.assertEqual(callback_request.host, None)
        self.assertEqual(callback_request.port, 80)
        self.assertEqual(callback_request.header(b'Host'), headers[b'Host'])
        assert callback_request._url
        self.assertEqual(callback_request._url.remainder, b'/')
        self.assertEqual(callback_request.method, httpMethods.GET)
        self.assertEqual(callback_request.is_https_tunnel, False)
Ejemplo n.º 29
0
    async def test_e2e(self, mocker: MockerFixture) -> None:
        host, port = uuid.uuid4().hex, 443
        netloc = '{0}:{1}'.format(host, port)

        self.mock_fromfd = mocker.patch('socket.fromfd')
        self.mock_selector = mocker.patch('selectors.DefaultSelector')
        self.mock_sign_csr = mocker.patch('proxy.http.proxy.server.sign_csr')
        self.mock_gen_csr = mocker.patch('proxy.http.proxy.server.gen_csr')
        self.mock_gen_public_key = mocker.patch(
            'proxy.http.proxy.server.gen_public_key', )
        self.mock_server_conn = mocker.patch(
            'proxy.http.proxy.server.TcpServerConnection', )
        self.mock_ssl_context = mocker.patch('ssl.create_default_context')
        self.mock_ssl_wrap = mocker.patch('ssl.wrap_socket')

        self.mock_sign_csr.return_value = True
        self.mock_gen_csr.return_value = True
        self.mock_gen_public_key.return_value = True

        ssl_connection = mock.MagicMock(spec=ssl.SSLSocket)
        self.mock_ssl_context.return_value.wrap_socket.return_value = ssl_connection
        self.mock_ssl_wrap.return_value = mock.MagicMock(spec=ssl.SSLSocket)
        plain_connection = mock.MagicMock(spec=socket.socket)

        def mock_connection() -> Any:
            if self.mock_ssl_context.return_value.wrap_socket.called:
                return ssl_connection
            return plain_connection

        # Do not mock the original wrap method
        self.mock_server_conn.return_value.wrap.side_effect = \
            lambda x, y: TcpServerConnection.wrap(
                self.mock_server_conn.return_value, x, y,
            )

        type(self.mock_server_conn.return_value).connection = \
            mock.PropertyMock(side_effect=mock_connection)

        self.fileno = 10
        self._addr = ('127.0.0.1', 54382)
        self.flags = FlagParser.initialize(
            ca_cert_file='ca-cert.pem',
            ca_key_file='ca-key.pem',
            ca_signing_key_file='ca-signing-key.pem',
            threaded=True,
        )
        self.plugin = mock.MagicMock()
        self.proxy_plugin = mock.MagicMock()
        self.flags.plugins = {
            b'HttpProtocolHandlerPlugin': [self.plugin, HttpProxyPlugin],
            b'HttpProxyBasePlugin': [self.proxy_plugin],
        }
        self._conn = self.mock_fromfd.return_value
        self.protocol_handler = HttpProtocolHandler(
            HttpClientConnection(self._conn, self._addr),
            flags=self.flags,
        )
        self.protocol_handler.initialize()

        self.plugin.assert_not_called()
        self.proxy_plugin.assert_not_called()

        connect_request = build_http_request(
            httpMethods.CONNECT,
            bytes_(netloc),
            headers={
                b'Host': bytes_(netloc),
            },
        )
        self._conn.recv.return_value = connect_request

        async def asyncReturnBool(val: bool) -> bool:
            return val

        # Prepare mocked HttpProtocolHandlerPlugin
        # self.plugin.return_value.get_descriptors.return_value = ([], [])
        # self.plugin.return_value.write_to_descriptors.return_value = asyncReturnBool(False)
        # self.plugin.return_value.read_from_descriptors.return_value = asyncReturnBool(False)
        # self.plugin.return_value.on_client_data.side_effect = lambda raw: raw
        # self.plugin.return_value.on_request_complete.return_value = False
        # self.plugin.return_value.on_response_chunk.side_effect = lambda chunk: chunk
        # self.plugin.return_value.on_client_connection_close.return_value = None

        # Prepare mocked HttpProxyBasePlugin
        self.proxy_plugin.return_value.write_to_descriptors.return_value = \
            asyncReturnBool(False)
        self.proxy_plugin.return_value.read_from_descriptors.return_value = \
            asyncReturnBool(False)
        self.proxy_plugin.return_value.before_upstream_connection.side_effect = lambda r: r
        self.proxy_plugin.return_value.handle_client_request.side_effect = lambda r: r
        self.proxy_plugin.return_value.resolve_dns.return_value = None, None

        self.mock_selector.return_value.select.side_effect = [
            [(
                selectors.SelectorKey(
                    fileobj=self._conn.fileno(),
                    fd=self._conn.fileno(),
                    events=selectors.EVENT_READ,
                    data=None,
                ),
                selectors.EVENT_READ,
            )],
        ]

        await self.protocol_handler._run_once()

        # Assert correct plugin was initialized
        self.plugin.assert_not_called()
        self.proxy_plugin.assert_called_once()
        self.assertEqual(self.proxy_plugin.call_args[0][1], self.flags)
        # Actual call arg must be `_conn` object
        # but because internally the reference is updated
        # we assert it against `mock_ssl_wrap` which is
        # called during proxy plugin initialization
        # for interception
        self.assertEqual(
            self.proxy_plugin.call_args[0][2].connection,
            self.mock_ssl_wrap.return_value,
        )

        # Assert our mocked plugins invocations
        # self.plugin.return_value.get_descriptors.assert_called()
        # self.plugin.return_value.write_to_descriptors.assert_called_with([])
        # # on_client_data is only called after initial request has completed
        # self.plugin.return_value.on_client_data.assert_not_called()
        # self.plugin.return_value.on_request_complete.assert_called()
        # self.plugin.return_value.read_from_descriptors.assert_called_with([
        #     self._conn.fileno(),
        # ])
        self.proxy_plugin.return_value.before_upstream_connection.assert_called(
        )
        self.proxy_plugin.return_value.handle_client_request.assert_called()

        self.mock_server_conn.assert_called_with(host, port)
        self.mock_server_conn.return_value.connection.setblocking.assert_called_with(
            False, )

        self.mock_ssl_context.assert_called_with(
            ssl.Purpose.SERVER_AUTH,
            cafile=str(DEFAULT_CA_FILE),
        )
        # self.assertEqual(self.mock_ssl_context.return_value.options,
        #                  ssl.OP_NO_SSLv2 | ssl.OP_NO_SSLv3 | ssl.OP_NO_TLSv1 |
        #                  ssl.OP_NO_TLSv1_1)
        self.assertEqual(plain_connection.setblocking.call_count, 2)
        self.mock_ssl_context.return_value.wrap_socket.assert_called_with(
            plain_connection,
            server_hostname=host,
        )
        self.assertEqual(self.mock_sign_csr.call_count, 1)
        self.assertEqual(self.mock_gen_csr.call_count, 1)
        self.assertEqual(self.mock_gen_public_key.call_count, 1)
        self.assertEqual(ssl_connection.setblocking.call_count, 1)
        self.assertEqual(
            self.mock_server_conn.return_value._conn,
            ssl_connection,
        )
        self._conn.send.assert_called_with(
            PROXY_TUNNEL_ESTABLISHED_RESPONSE_PKT, )
        assert self.flags.ca_cert_dir is not None
        self.mock_ssl_wrap.assert_called_with(
            self._conn,
            server_side=True,
            keyfile=self.flags.ca_signing_key_file,
            certfile=HttpProxyPlugin.generated_cert_file_path(
                self.flags.ca_cert_dir,
                host,
            ),
            ssl_version=ssl.PROTOCOL_TLS,
        )
        self.assertEqual(self._conn.setblocking.call_count, 2)
        self.assertEqual(
            self.protocol_handler.work.connection,
            self.mock_ssl_wrap.return_value,
        )

        # Assert connection references for all other plugins is updated
        # self.assertEqual(
        #     self.plugin.return_value.client._conn,
        #     self.mock_ssl_wrap.return_value,
        # )
        self.assertEqual(
            self.proxy_plugin.return_value.client._conn,
            self.mock_ssl_wrap.return_value,
        )
Ejemplo n.º 30
0
    def test_setup_and_shutdown(
        self,
        mock_listener: mock.Mock,
        mock_acceptor: mock.Mock,
        mock_pipe: mock.Mock,
        mock_send_handle: mock.Mock,
    ) -> None:
        acceptor1 = mock.MagicMock()
        acceptor2 = mock.MagicMock()
        mock_acceptor.side_effect = [acceptor1, acceptor2]

        num_acceptors = 2
        flags = FlagParser.initialize(
            num_acceptors=num_acceptors,
            threaded=True,
        )
        self.assertEqual(flags.num_acceptors, num_acceptors)

        pool = AcceptorPool(
            flags=flags,
            listener=mock_listener.return_value,
            executor_queues=[],
            executor_pids=[],
            executor_locks=[],
        )
        pool.setup()

        self.assertEqual(mock_pipe.call_count, num_acceptors)
        self.assertEqual(mock_acceptor.call_count, num_acceptors)
        mock_send_handle.assert_called()
        self.assertEqual(mock_send_handle.call_count, num_acceptors)

        self.assertEqual(
            mock_acceptor.call_args_list[0][1]['idd'],
            0,
        )
        self.assertEqual(
            mock_acceptor.call_args_list[0][1]['fd_queue'],
            mock_pipe.return_value[1],
        )
        self.assertEqual(
            mock_acceptor.call_args_list[0][1]['flags'],
            flags,
        )
        self.assertEqual(
            mock_acceptor.call_args_list[0][1]['event_queue'],
            None,
        )
        # executor_queues=[],
        # executor_pids=[]
        self.assertEqual(
            mock_acceptor.call_args_list[1][1]['idd'],
            1,
        )

        acceptor1.start.assert_called_once()
        acceptor2.start.assert_called_once()
        mock_listener.return_value.fileno.assert_called_once()

        acceptor1.join.assert_not_called()
        acceptor2.join.assert_not_called()

        pool.shutdown()
        acceptor1.join.assert_called_once()
        acceptor2.join.assert_called_once()