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")

[((52, 9, 52, 22), 'mitmproxy.addons.proxyserver.Proxyserver', 'Proxyserver', ({}, {}), '()', False, 'from mitmproxy.addons.proxyserver import Proxyserver\n'), ((98, 9, 98, 22), 'mitmproxy.addons.proxyserver.Proxyserver', 'Proxyserver', ({}, {}), '()', False, 'from mitmproxy.addons.proxyserver import Proxyserver\n'), ((123, 9, 123, 22), 'mitmproxy.addons.proxyserver.Proxyserver', 'Proxyserver', ({}, {}), '()', False, 'from mitmproxy.addons.proxyserver import Proxyserver\n'), ((158, 9, 158, 22), 'mitmproxy.addons.proxyserver.Proxyserver', 'Proxyserver', ({}, {}), '()', False, 'from mitmproxy.addons.proxyserver import Proxyserver\n'), ((168, 13, 168, 27), 'mitmproxy.test.tflow.tserver_conn', 'tserver_conn', ({}, {}), '()', False, 'from mitmproxy.test.tflow import tclient_conn, tserver_conn\n'), ((169, 13, 169, 27), 'mitmproxy.test.tflow.tclient_conn', 'tclient_conn', ({}, {}), '()', False, 'from mitmproxy.test.tflow import tclient_conn, tserver_conn\n'), ((171, 9, 171, 22), 'mitmproxy.addons.proxyserver.Proxyserver', 'Proxyserver', ({}, {}), '()', False, 'from mitmproxy.addons.proxyserver import Proxyserver\n'), ((182, 9, 182, 22), 'mitmproxy.addons.proxyserver.Proxyserver', 'Proxyserver', ({}, {}), '()', False, 'from mitmproxy.addons.proxyserver import Proxyserver\n'), ((36, 19, 36, 68), 'asyncio.start_server', 'asyncio.start_server', ({(36, 40, 36, 51): 'handle_conn', (36, 53, 36, 64): '"""127.0.0.1"""', (36, 66, 36, 67): '(0)'}, {}), "(handle_conn, '127.0.0.1', 0)", False, 'import asyncio\n'), ((53, 9, 53, 28), 'mitmproxy.test.taddons.context', 'taddons.context', ({(53, 25, 53, 27): 'ps'}, {}), '(ps)', False, 'from mitmproxy.test import taddons, tflow\n'), ((99, 9, 99, 28), 'mitmproxy.test.taddons.context', 'taddons.context', ({(99, 25, 99, 27): 'ps'}, {}), '(ps)', False, 'from mitmproxy.test import taddons, tflow\n'), ((124, 9, 124, 28), 'mitmproxy.test.taddons.context', 'taddons.context', ({(124, 25, 124, 27): 'ps'}, {}), '(ps)', False, 'from mitmproxy.test import taddons, tflow\n'), ((159, 9, 159, 28), 'mitmproxy.test.taddons.context', 'taddons.context', ({(159, 25, 159, 27): 'ps'}, {}), '(ps)', False, 'from mitmproxy.test import taddons, tflow\n'), ((172, 9, 172, 28), 'mitmproxy.test.taddons.context', 'taddons.context', ({(172, 25, 172, 27): 'ps'}, {}), '(ps)', False, 'from mitmproxy.test import taddons, tflow\n'), ((183, 9, 183, 28), 'mitmproxy.test.taddons.context', 'taddons.context', ({(183, 25, 183, 27): 'ps'}, {}), '(ps)', False, 'from mitmproxy.test import taddons, tflow\n'), ((126, 12, 126, 25), 'mitmproxy.test.tflow.tflow', 'tflow.tflow', ({}, {}), '()', False, 'from mitmproxy.test import taddons, tflow\n'), ((132, 12, 132, 25), 'mitmproxy.test.tflow.tflow', 'tflow.tflow', ({}, {}), '()', False, 'from mitmproxy.test import taddons, tflow\n'), ((139, 12, 139, 34), 'mitmproxy.test.tflow.twebsocketflow', 'tflow.twebsocketflow', ({}, {}), '()', False, 'from mitmproxy.test import taddons, tflow\n'), ((145, 12, 145, 28), 'mitmproxy.test.tflow.ttcpflow', 'tflow.ttcpflow', ({}, {}), '()', False, 'from mitmproxy.test import taddons, tflow\n'), ((176, 12, 176, 65), 'mitmproxy.proxy.server_hooks.ServerConnectionHookData', 'server_hooks.ServerConnectionHookData', ({(176, 50, 176, 56): 'server', (176, 58, 176, 64): 'client'}, {}), '(server, client)', False, 'from mitmproxy.proxy import layers, server_hooks\n'), ((184, 13, 184, 51), 'pytest.raises', 'pytest.raises', ({(184, 27, 184, 50): 'exceptions.OptionsError'}, {}), '(exceptions.OptionsError)', False, 'import pytest\n'), ((188, 13, 188, 51), 'pytest.raises', 'pytest.raises', ({(188, 27, 188, 50): 'exceptions.OptionsError'}, {}), '(exceptions.OptionsError)', False, 'import pytest\n'), ((19, 24, 19, 51), 'mitmproxy.proxy.layers.modes.HttpProxy', 'layers.modes.HttpProxy', ({(19, 47, 19, 50): 'ctx'}, {}), '(ctx)', False, 'from mitmproxy.proxy import layers, server_hooks\n'), ((20, 24, 20, 63), 'mitmproxy.proxy.layers.HttpLayer', 'layers.HttpLayer', ({(20, 41, 20, 44): 'ctx', (20, 46, 20, 62): 'HTTPMode.regular'}, {}), '(ctx, HTTPMode.regular)', False, 'from mitmproxy.proxy import layers, server_hooks\n'), ((21, 24, 21, 44), 'mitmproxy.proxy.layers.TCPLayer', 'layers.TCPLayer', ({(21, 40, 21, 43): 'ctx'}, {}), '(ctx)', False, 'from mitmproxy.proxy import layers, server_hooks\n'), ((64, 35, 64, 71), 'asyncio.open_connection', 'asyncio.open_connection', ({(64, 59, 64, 70): '*proxy_addr'}, {}), '(*proxy_addr)', False, 'import asyncio\n'), ((107, 35, 107, 71), 'asyncio.open_connection', 'asyncio.open_connection', ({(107, 59, 107, 70): '*proxy_addr'}, {}), '(*proxy_addr)', False, 'import asyncio\n'), ((88, 22, 88, 38), 'asyncio.sleep', 'asyncio.sleep', ({(88, 36, 88, 37): '(0)'}, {}), '(0)', False, 'import asyncio\n')]