import asyncio
from contextlib import asynccontextmanager

import pytest

from mitmproxy import exceptions
from mitmproxy.addons.proxyserver import Proxyserver
from mitmproxy.connection import Address
from mitmproxy.proxy import layers, server_hooks
from mitmproxy.proxy.layers.http import HTTPMode
from mitmproxy.test import taddons, tflow
from mitmproxy.test.tflow import tclient_conn, tserver_conn


class HelperAddon:
    def __init__(self):
        self.flows = []
        self.layers = [
            lambda ctx: layers.modes.HttpProxy(ctx),
            lambda ctx: layers.HttpLayer(ctx, HTTPMode.regular),
            lambda ctx: layers.TCPLayer(ctx),
        ]

    def request(self, f):
        self.flows.append(f)

    def tcp_start(self, f):
        self.flows.append(f)

    def next_layer(self, nl):
        nl.layer = self.layers.pop(0)(nl.context)


@asynccontextmanager
async def tcp_server(handle_conn) -> Address:
    server = await asyncio.start_server(handle_conn, '127.0.0.1', 0)
    await server.start_serving()
    try:
        yield server.sockets[0].getsockname()
    finally:
        server.close()


@pytest.mark.asyncio
async def test_start_stop():
    async def server_handler(reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
        assert await reader.readuntil(b"\r\n\r\n") == b"GET /hello HTTP/1.1\r\n\r\n"
        writer.write(b"HTTP/1.1 204 No Content\r\n\r\n")
        await writer.drain()
        writer.close()

    ps = Proxyserver()
    with taddons.context(ps) as tctx:
        state = HelperAddon()
        tctx.master.addons.add(state)
        async with tcp_server(server_handler) as addr:
            tctx.configure(ps, listen_host="127.0.0.1", listen_port=0)
            assert not ps.server
            ps.running()
            await tctx.master.await_log("Proxy server listening", level="info")
            assert ps.server

            proxy_addr = ps.server.sockets[0].getsockname()[:2]
            reader, writer = await asyncio.open_connection(*proxy_addr)
            req = f"GET http://{addr[0]}:{addr[1]}/hello HTTP/1.1\r\n\r\n"
            writer.write(req.encode())
            assert await reader.readuntil(b"\r\n\r\n") == b"HTTP/1.1 204 No Content\r\n\r\n"
            assert repr(ps) == "ProxyServer(running, 1 active conns)"

            tctx.configure(ps, server=False)
            await tctx.master.await_log("Stopping server", level="info")
            assert not ps.server
            assert state.flows
            assert state.flows[0].request.path == "/hello"
            assert state.flows[0].response.status_code == 204

            # Waiting here until everything is really torn down... takes some effort.
            conn_handler = list(ps._connections.values())[0]
            client_handler = conn_handler.transports[conn_handler.client].handler
            writer.close()
            await writer.wait_closed()
            try:
                await client_handler
            except asyncio.CancelledError:
                pass
            for _ in range(5):
                # Get all other scheduled coroutines to run.
                await asyncio.sleep(0)
            assert repr(ps) == "ProxyServer(stopped, 0 active conns)"


@pytest.mark.asyncio
async def test_inject() -> None:
    async def server_handler(reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
        while s := await reader.read(1):
            writer.write(s.upper())

    ps = Proxyserver()
    with taddons.context(ps) as tctx:
        state = HelperAddon()
        tctx.master.addons.add(state)
        async with tcp_server(server_handler) as addr:
            tctx.configure(ps, listen_host="127.0.0.1", listen_port=0)
            ps.running()
            await tctx.master.await_log("Proxy server listening", level="info")
            proxy_addr = ps.server.sockets[0].getsockname()[:2]
            reader, writer = await asyncio.open_connection(*proxy_addr)

            req = f"CONNECT {addr[0]}:{addr[1]} HTTP/1.1\r\n\r\n"
            writer.write(req.encode())
            assert await reader.readuntil(b"\r\n\r\n") == b"HTTP/1.1 200 Connection established\r\n\r\n"

            writer.write(b"a")
            assert await reader.read(1) == b"A"
            ps.inject_tcp(state.flows[0], False, b"b")
            assert await reader.read(1) == b"B"
            ps.inject_tcp(state.flows[0], True, b"c")
            assert await reader.read(1) == b"c"


@pytest.mark.asyncio
async def test_inject_fail() -> None:
    ps = Proxyserver()
    with taddons.context(ps) as tctx:
        ps.inject_websocket(
            tflow.tflow(),
            True,
            b"test"
        )
        await tctx.master.await_log("Cannot inject WebSocket messages into non-WebSocket flows.", level="warn")
        ps.inject_tcp(
            tflow.tflow(),
            True,
            b"test"
        )
        await tctx.master.await_log("Cannot inject TCP messages into non-TCP flows.", level="warn")

        ps.inject_websocket(
            tflow.twebsocketflow(),
            True,
            b"test"
        )
        await tctx.master.await_log("Flow is not from a live connection.", level="warn")
        ps.inject_websocket(
            tflow.ttcpflow(),
            True,
            b"test"
        )
        await tctx.master.await_log("Flow is not from a live connection.", level="warn")


@pytest.mark.asyncio
async def test_warn_no_nextlayer():
    """
    Test that we log an error if the proxy server is started without NextLayer addon.
    That is a mean trap to fall into when writing end-to-end tests.
    """
    ps = Proxyserver()
    with taddons.context(ps) as tctx:
        tctx.configure(ps, listen_host="127.0.0.1", listen_port=0)
        ps.running()
        await tctx.master.await_log("Proxy server listening at", level="info")
        assert tctx.master.has_log("Warning: Running proxyserver without nextlayer addon!", level="warn")
        await ps.shutdown_server()


def test_self_connect():
    server = tserver_conn()
    client = tclient_conn()
    server.address = ("localhost", 8080)
    ps = Proxyserver()
    with taddons.context(ps) as tctx:
        # not calling .running() here to avoid unnecessary socket
        ps.options = tctx.options
        ps.server_connect(
            server_hooks.ServerConnectionHookData(server, client)
        )
        assert server.error == "Stopped mitmproxy from recursively connecting to itself."


def test_options():
    ps = Proxyserver()
    with taddons.context(ps) as tctx:
        with pytest.raises(exceptions.OptionsError):
            tctx.configure(ps, body_size_limit="invalid")
        tctx.configure(ps, body_size_limit="1m")

        with pytest.raises(exceptions.OptionsError):
            tctx.configure(ps, stream_large_bodies="invalid")
        tctx.configure(ps, stream_large_bodies="1m")

[((1471, 1484), 'mitmproxy.addons.proxyserver.Proxyserver', 'Proxyserver', ([], {}), '()\n', (1482, 1484), False, 'from mitmproxy.addons.proxyserver import Proxyserver\n'), ((3450, 3463), 'mitmproxy.addons.proxyserver.Proxyserver', 'Proxyserver', ([], {}), '()\n', (3461, 3463), False, 'from mitmproxy.addons.proxyserver import Proxyserver\n'), ((4503, 4516), 'mitmproxy.addons.proxyserver.Proxyserver', 'Proxyserver', ([], {}), '()\n', (4514, 4516), False, 'from mitmproxy.addons.proxyserver import Proxyserver\n'), ((5606, 5619), 'mitmproxy.addons.proxyserver.Proxyserver', 'Proxyserver', ([], {}), '()\n', (5617, 5619), False, 'from mitmproxy.addons.proxyserver import Proxyserver\n'), ((6006, 6020), 'mitmproxy.test.tflow.tserver_conn', 'tserver_conn', ([], {}), '()\n', (6018, 6020), False, 'from mitmproxy.test.tflow import tclient_conn, tserver_conn\n'), ((6034, 6048), 'mitmproxy.test.tflow.tclient_conn', 'tclient_conn', ([], {}), '()\n', (6046, 6048), False, 'from mitmproxy.test.tflow import tclient_conn, tserver_conn\n'), ((6099, 6112), 'mitmproxy.addons.proxyserver.Proxyserver', 'Proxyserver', ([], {}), '()\n', (6110, 6112), False, 'from mitmproxy.addons.proxyserver import Proxyserver\n'), ((6475, 6488), 'mitmproxy.addons.proxyserver.Proxyserver', 'Proxyserver', ([], {}), '()\n', (6486, 6488), False, 'from mitmproxy.addons.proxyserver import Proxyserver\n'), ((951, 1000), 'asyncio.start_server', 'asyncio.start_server', (['handle_conn', '"""127.0.0.1"""', '(0)'], {}), "(handle_conn, '127.0.0.1', 0)\n", (971, 1000), False, 'import asyncio\n'), ((1494, 1513), 'mitmproxy.test.taddons.context', 'taddons.context', (['ps'], {}), '(ps)\n', (1509, 1513), False, 'from mitmproxy.test import taddons, tflow\n'), ((3473, 3492), 'mitmproxy.test.taddons.context', 'taddons.context', (['ps'], {}), '(ps)\n', (3488, 3492), False, 'from mitmproxy.test import taddons, tflow\n'), ((4526, 4545), 'mitmproxy.test.taddons.context', 'taddons.context', (['ps'], {}), '(ps)\n', (4541, 4545), False, 'from mitmproxy.test import taddons, tflow\n'), ((5629, 5648), 'mitmproxy.test.taddons.context', 'taddons.context', (['ps'], {}), '(ps)\n', (5644, 5648), False, 'from mitmproxy.test import taddons, tflow\n'), ((6122, 6141), 'mitmproxy.test.taddons.context', 'taddons.context', (['ps'], {}), '(ps)\n', (6137, 6141), False, 'from mitmproxy.test import taddons, tflow\n'), ((6498, 6517), 'mitmproxy.test.taddons.context', 'taddons.context', (['ps'], {}), '(ps)\n', (6513, 6517), False, 'from mitmproxy.test import taddons, tflow\n'), ((4596, 4609), 'mitmproxy.test.tflow.tflow', 'tflow.tflow', ([], {}), '()\n', (4607, 4609), False, 'from mitmproxy.test import taddons, tflow\n'), ((4806, 4819), 'mitmproxy.test.tflow.tflow', 'tflow.tflow', ([], {}), '()\n', (4817, 4819), False, 'from mitmproxy.test import taddons, tflow\n'), ((5011, 5033), 'mitmproxy.test.tflow.twebsocketflow', 'tflow.twebsocketflow', ([], {}), '()\n', (5031, 5033), False, 'from mitmproxy.test import taddons, tflow\n'), ((5213, 5229), 'mitmproxy.test.tflow.ttcpflow', 'tflow.ttcpflow', ([], {}), '()\n', (5227, 5229), False, 'from mitmproxy.test import taddons, tflow\n'), ((6290, 6343), 'mitmproxy.proxy.server_hooks.ServerConnectionHookData', 'server_hooks.ServerConnectionHookData', (['server', 'client'], {}), '(server, client)\n', (6327, 6343), False, 'from mitmproxy.proxy import layers, server_hooks\n'), ((6540, 6578), 'pytest.raises', 'pytest.raises', (['exceptions.OptionsError'], {}), '(exceptions.OptionsError)\n', (6553, 6578), False, 'import pytest\n'), ((6701, 6739), 'pytest.raises', 'pytest.raises', (['exceptions.OptionsError'], {}), '(exceptions.OptionsError)\n', (6714, 6739), False, 'import pytest\n'), ((518, 545), 'mitmproxy.proxy.layers.modes.HttpProxy', 'layers.modes.HttpProxy', (['ctx'], {}), '(ctx)\n', (540, 545), False, 'from mitmproxy.proxy import layers, server_hooks\n'), ((571, 610), 'mitmproxy.proxy.layers.HttpLayer', 'layers.HttpLayer', (['ctx', 'HTTPMode.regular'], {}), '(ctx, HTTPMode.regular)\n', (587, 610), False, 'from mitmproxy.proxy import layers, server_hooks\n'), ((636, 656), 'mitmproxy.proxy.layers.TCPLayer', 'layers.TCPLayer', (['ctx'], {}), '(ctx)\n', (651, 656), False, 'from mitmproxy.proxy import layers, server_hooks\n'), ((1984, 2020), 'asyncio.open_connection', 'asyncio.open_connection', (['*proxy_addr'], {}), '(*proxy_addr)\n', (2007, 2020), False, 'import asyncio\n'), ((3900, 3936), 'asyncio.open_connection', 'asyncio.open_connection', (['*proxy_addr'], {}), '(*proxy_addr)\n', (3923, 3936), False, 'import asyncio\n'), ((3130, 3146), 'asyncio.sleep', 'asyncio.sleep', (['(0)'], {}), '(0)\n', (3143, 3146), False, 'import asyncio\n')]