<commit_before>import asyncio
import traceback
import typing

import mitmproxy.types
from mitmproxy import command
from mitmproxy import ctx
from mitmproxy import exceptions
from mitmproxy import flow
from mitmproxy import http
from mitmproxy import io
from mitmproxy.addons.proxyserver import AsyncReply
from mitmproxy.net import server_spec
from mitmproxy.options import Options
from mitmproxy.proxy.protocol.http import HTTPMode
from mitmproxy.proxy2 import commands, events, layers, server
from mitmproxy.proxy2.context import Context, Server
from mitmproxy.proxy2.layer import CommandGenerator


class MockServer(layers.http.HttpConnection):
    """
    A mock HTTP "server" that just pretends it received a full HTTP request,
    which is then processed by the proxy core.
    """
    flow: http.HTTPFlow

    def __init__(self, flow: http.HTTPFlow, context: Context):
        super().__init__(context, context.client)
        self.flow = flow

    def _handle_event(self, event: events.Event) -> CommandGenerator[None]:
        if isinstance(event, events.Start):
            yield layers.http.ReceiveHttp(layers.http.RequestHeaders(1, self.flow.request))
            if self.flow.request.raw_content:
                yield layers.http.ReceiveHttp(layers.http.RequestData(1, self.flow.request.raw_content))
            yield layers.http.ReceiveHttp(layers.http.RequestEndOfMessage(1))
        elif isinstance(event, (
                layers.http.ResponseHeaders,
                layers.http.ResponseData,
                layers.http.ResponseEndOfMessage,
                layers.http.ResponseProtocolError,
        )):
            pass
        else:
            ctx.log(f"Unexpected event during replay: {events}")


class ReplayHandler(server.ConnectionHandler):
    def __init__(self, flow: http.HTTPFlow, options: Options) -> None:
        client = flow.client_conn.copy()

        context = Context(client, options)
        context.server = Server(
            (flow.request.host, flow.request.port)
        )
        context.server.tls = flow.request.scheme == "https"
        if options.mode.startswith("upstream:"):
            context.server.via = server_spec.parse_with_mode(options.mode)[1]

        super().__init__(context)

        self.layer = layers.HttpLayer(context, HTTPMode.transparent)
        self.layer.connections[client] = MockServer(flow, context.fork())
        self.flow = flow
        self.done = asyncio.Event()

    async def replay(self) -> None:
        self.server_event(events.Start())
        await self.done.wait()

    def log(self, message: str, level: str = "info") -> None:
        ctx.log(f"[replay] {message}", level)

    async def handle_hook(self, hook: commands.Hook) -> None:
        data, = hook.as_tuple()
        data.reply = AsyncReply(data)
        await ctx.master.addons.handle_lifecycle(hook.name, data)
        await data.reply.done.wait()
        if isinstance(hook, (layers.http.HttpResponseHook, layers.http.HttpErrorHook)):
            if self.transports:
                # close server connections
                for x in self.transports.values():
                    x.handler.cancel()
                await asyncio.wait([x.handler for x in self.transports.values()])
            # signal completion
            self.done.set()


class ClientPlayback:
    playback_task: typing.Optional[asyncio.Task]
    inflight: typing.Optional[http.HTTPFlow]
    queue: asyncio.Queue
    options: Options

    def __init__(self):
        self.queue = asyncio.Queue()
        self.inflight = None
        self.task = None

    def running(self):
        self.playback_task = asyncio.create_task(self.playback())
        self.options = ctx.options

    def done(self):
        self.playback_task.cancel()

    async def playback(self):
        try:
            while True:
                self.inflight = await self.queue.get()
                try:
                    h = ReplayHandler(self.inflight, self.options)
                    await h.replay()
                except Exception:
                    ctx.log(f"Client replay has crashed!\n{traceback.format_exc()}", "error")
                self.inflight = None
        except asyncio.CancelledError:
            return

    def check(self, f: flow.Flow) -> typing.Optional[str]:
        if f.live:
            return "Can't replay live flow."
        if f.intercepted:
            return "Can't replay intercepted flow."
        if isinstance(f, http.HTTPFlow):
            if not f.request:
                return "Can't replay flow with missing request."
            if f.request.raw_content is None:
                return "Can't replay flow with missing content."
        else:
            return "Can only replay HTTP flows."

    def load(self, loader):
        loader.add_option(
            "client_replay", typing.Sequence[str], [],
            "Replay client requests from a saved file."
        )

    def configure(self, updated):
        if "client_replay" in updated and ctx.options.client_replay:
            try:
                flows = io.read_flows_from_paths(ctx.options.client_replay)<commit_msg>[sans-io] clientplayback: don't catch cancellederror<commit_after>import asyncio
import traceback
import typing

import mitmproxy.types
from mitmproxy import command
from mitmproxy import ctx
from mitmproxy import exceptions
from mitmproxy import flow
from mitmproxy import http
from mitmproxy import io
from mitmproxy.addons.proxyserver import AsyncReply
from mitmproxy.net import server_spec
from mitmproxy.options import Options
from mitmproxy.proxy.protocol.http import HTTPMode
from mitmproxy.proxy2 import commands, events, layers, server
from mitmproxy.proxy2.context import Context, Server
from mitmproxy.proxy2.layer import CommandGenerator


class MockServer(layers.http.HttpConnection):
    """
    A mock HTTP "server" that just pretends it received a full HTTP request,
    which is then processed by the proxy core.
    """
    flow: http.HTTPFlow

    def __init__(self, flow: http.HTTPFlow, context: Context):
        super().__init__(context, context.client)
        self.flow = flow

    def _handle_event(self, event: events.Event) -> CommandGenerator[None]:
        if isinstance(event, events.Start):
            yield layers.http.ReceiveHttp(layers.http.RequestHeaders(1, self.flow.request))
            if self.flow.request.raw_content:
                yield layers.http.ReceiveHttp(layers.http.RequestData(1, self.flow.request.raw_content))
            yield layers.http.ReceiveHttp(layers.http.RequestEndOfMessage(1))
        elif isinstance(event, (
                layers.http.ResponseHeaders,
                layers.http.ResponseData,
                layers.http.ResponseEndOfMessage,
                layers.http.ResponseProtocolError,
        )):
            pass
        else:
            ctx.log(f"Unexpected event during replay: {events}")


class ReplayHandler(server.ConnectionHandler):
    def __init__(self, flow: http.HTTPFlow, options: Options) -> None:
        client = flow.client_conn.copy()

        context = Context(client, options)
        context.server = Server(
            (flow.request.host, flow.request.port)
        )
        context.server.tls = flow.request.scheme == "https"
        if options.mode.startswith("upstream:"):
            context.server.via = server_spec.parse_with_mode(options.mode)[1]

        super().__init__(context)

        self.layer = layers.HttpLayer(context, HTTPMode.transparent)
        self.layer.connections[client] = MockServer(flow, context.fork())
        self.flow = flow
        self.done = asyncio.Event()

    async def replay(self) -> None:
        self.server_event(events.Start())
        await self.done.wait()

    def log(self, message: str, level: str = "info") -> None:
        ctx.log(f"[replay] {message}", level)

    async def handle_hook(self, hook: commands.Hook) -> None:
        data, = hook.as_tuple()
        data.reply = AsyncReply(data)
        await ctx.master.addons.handle_lifecycle(hook.name, data)
        await data.reply.done.wait()
        if isinstance(hook, (layers.http.HttpResponseHook, layers.http.HttpErrorHook)):
            if self.transports:
                # close server connections
                for x in self.transports.values():
                    x.handler.cancel()
                await asyncio.wait([x.handler for x in self.transports.values()])
            # signal completion
            self.done.set()


class ClientPlayback:
    playback_task: typing.Optional[asyncio.Task]
    inflight: typing.Optional[http.HTTPFlow]
    queue: asyncio.Queue
    options: Options

    def __init__(self):
        self.queue = asyncio.Queue()
        self.inflight = None
        self.task = None

    def running(self):
        self.playback_task = asyncio.create_task(self.playback())
        self.options = ctx.options

    def done(self):
        self.playback_task.cancel()

    async def playback(self):
        while True:
            self.inflight = await self.queue.get()
            try:
                h = ReplayHandler(self.inflight, self.options)
                await h.replay()
            except Exception:
                ctx.log(f"Client replay has crashed!\n{traceback.format_exc()}", "error")
            self.inflight = None

    def check(self, f: flow.Flow) -> typing.Optional[str]:
        if f.live:
            return "Can't replay live flow."
        if f.intercepted:
            return "Can't replay intercepted flow."
        if isinstance(f, http.HTTPFlow):
            if not f.request:
                return "Can't replay flow with missing request."
            if f.request.raw_content is None:
                return "Can't replay flow with missing content."
        else:
            return "Can only replay HTTP flows."

    def load(self, loader):
        loader.add_option(
            "client_replay", typing.Sequence[str], [],
            "Replay client requests from a saved file."
        )

    def configure(self, updated):
        if "client_replay" in updated and ctx.options.client_replay:
            try:
                flows = io.read_flows_from_paths(ctx.options.client_replay)
{'fraction_non_alphanumeric': 0.0971540726202159, 'fraction_numerical': 0.0013738959764474975, 'mean_word_length': 2.3666997026759167, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 15, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3093960', 'n_tokens_mistral': 2854, 'n_tokens_neox': 2764, 'n_words': 701}