codekingpro/portable-devtools

1from __future__ import annotations2 3import asyncio4import logging5import time6from collections.abc import Sequence7from types import TracebackType8from typing import cast9from typing import Literal10 11import mitmproxy.types12from mitmproxy import command13from mitmproxy import ctx14from mitmproxy import exceptions15from mitmproxy import flow16from mitmproxy import http17from mitmproxy import io18from mitmproxy.connection import ConnectionState19from mitmproxy.connection import Server20from mitmproxy.hooks import UpdateHook21from mitmproxy.log import ALERT22from mitmproxy.options import Options23from mitmproxy.proxy import commands24from mitmproxy.proxy import events25from mitmproxy.proxy import layers26from mitmproxy.proxy import server27from mitmproxy.proxy.context import Context28from mitmproxy.proxy.layer import CommandGenerator29from mitmproxy.proxy.layers.http import HTTPMode30from mitmproxy.proxy.mode_specs import UpstreamMode31from mitmproxy.utils import asyncio_utils32 33logger = logging.getLogger(__name__)34 35 36class MockServer(layers.http.HttpConnection):37    """38    A mock HTTP "server" that just pretends it received a full HTTP request,39    which is then processed by the proxy core.40    """41 42    flow: http.HTTPFlow43 44    def __init__(self, flow: http.HTTPFlow, context: Context):45        super().__init__(context, context.client)46        self.flow = flow47 48    def _handle_event(self, event: events.Event) -> CommandGenerator[None]:49        if isinstance(event, events.Start):50            content = self.flow.request.raw_content51            self.flow.request.timestamp_start = self.flow.request.timestamp_end = (52                time.time()53            )54            yield layers.http.ReceiveHttp(55                layers.http.RequestHeaders(56                    1,57                    self.flow.request,58                    end_stream=not (content or self.flow.request.trailers),59                    replay_flow=self.flow,60                )61            )62            if content:63                yield layers.http.ReceiveHttp(layers.http.RequestData(1, content))64            if self.flow.request.trailers:  # pragma: no cover65                # TODO: Cover this once we support HTTP/1 trailers.66                yield layers.http.ReceiveHttp(67                    layers.http.RequestTrailers(1, self.flow.request.trailers)68                )69            yield layers.http.ReceiveHttp(layers.http.RequestEndOfMessage(1))70        elif isinstance(71            event,72            (73                layers.http.ResponseHeaders,74                layers.http.ResponseData,75                layers.http.ResponseTrailers,76                layers.http.ResponseEndOfMessage,77                layers.http.ResponseProtocolError,78            ),79        ):80            pass81        else:  # pragma: no cover82            logger.warning(f"Unexpected event during replay: {event}")83 84 85class ReplayHandler(server.ConnectionHandler):86    layer: layers.HttpLayer87 88    def __init__(self, flow: http.HTTPFlow, options: Options) -> None:89        client = flow.client_conn.copy()90        client.state = ConnectionState.OPEN91 92        context = Context(client, options)93        context.server = Server(address=(flow.request.host, flow.request.port))94        if flow.request.scheme == "https":95            context.server.tls = True96            context.server.sni = flow.request.pretty_host97        if options.mode and options.mode[0].startswith("upstream:"):98            mode = UpstreamMode.parse(options.mode[0])99            assert isinstance(mode, UpstreamMode)  # remove once mypy supports Self.100            context.server.via = flow.server_conn.via = (mode.scheme, mode.address)101 102        super().__init__(context)103 104        if options.mode and options.mode[0].startswith("upstream:"):105            self.layer = layers.HttpLayer(context, HTTPMode.upstream)106        else:107            self.layer = layers.HttpLayer(context, HTTPMode.transparent)108        self.layer.connections[client] = MockServer(flow, context.fork())109        self.flow = flow110        self.done = asyncio.Event()111 112    async def replay(self) -> None:113        await self.server_event(events.Start())114        await self.done.wait()115 116    def log(117        self,118        message: str,119        level: int = logging.INFO,120        exc_info: Literal[True]121        | tuple[type[BaseException] | None, BaseException | None, TracebackType | None]122        | None = None,123    ) -> None:124        assert isinstance(level, int)125        logger.log(level=level, msg=f"[replay] {message}")126 127    async def handle_hook(self, hook: commands.StartHook) -> None:128        (data,) = hook.args()129        await ctx.master.addons.handle_lifecycle(hook)130        if isinstance(data, flow.Flow):131            await data.wait_for_resume()132        if isinstance(hook, (layers.http.HttpResponseHook, layers.http.HttpErrorHook)):133            if self.transports:134                # close server connections135                for x in self.transports.values():136                    if x.handler:137                        x.handler.cancel()138                await asyncio.wait(139                    [x.handler for x in self.transports.values() if x.handler]140                )141            # signal completion142            self.done.set()143 144 145class ClientPlayback:146    playback_task: asyncio.Task | None = None147    inflight: http.HTTPFlow | None148    queue: asyncio.Queue149    options: Options150    replay_tasks: set[asyncio.Task]151 152    def __init__(self):153        self.queue = asyncio.Queue()154        self.inflight = None155        self.task = None156        self.replay_tasks = set()157 158    def running(self):159        self.options = ctx.options160        self.playback_task = asyncio_utils.create_task(161            self.playback(),162            name="client playback",163            keep_ref=False,164        )165 166    async def done(self):167        if self.playback_task:168            self.playback_task.cancel()169            try:170                await self.playback_task171            except asyncio.CancelledError:172                pass173 174    async def playback(self):175        while True:176            self.inflight = await self.queue.get()177            try:178                assert self.inflight179                h = ReplayHandler(self.inflight, self.options)180                if ctx.options.client_replay_concurrency == -1:181                    t = asyncio_utils.create_task(182                        h.replay(),183                        name="client playback awaiting response",184                        keep_ref=False,185                    )186                    # keep a reference so this is not garbage collected187                    self.replay_tasks.add(t)188                    t.add_done_callback(self.replay_tasks.remove)189                else:190                    await h.replay()191            except Exception:192                logger.exception(f"Client replay has crashed!")193            self.queue.task_done()194            self.inflight = None195 196    def check(self, f: flow.Flow) -> str | None:197        if f.live or f == self.inflight:198            return "Can't replay live flow."199        if f.intercepted:200            return "Can't replay intercepted flow."201        if isinstance(f, http.HTTPFlow):202            if not f.request:203                return "Can't replay flow with missing request."204            if f.request.raw_content is None:205                return "Can't replay flow with missing content."206            if f.websocket is not None:207                return "Can't replay WebSocket flows."208        else:209            return "Can only replay HTTP flows."210        return None211 212    def load(self, loader):213        loader.add_option(214            "client_replay",215            Sequence[str],216            [],217            "Replay client requests from a saved file.",218        )219        loader.add_option(220            "client_replay_concurrency",221            int,222            1,223            "Concurrency limit on in-flight client replay requests. Currently the only valid values are 1 and -1 (no limit).",224        )225 226    def configure(self, updated):227        if "client_replay" in updated and ctx.options.client_replay:228            try:229                flows = io.read_flows_from_paths(ctx.options.client_replay)230            except exceptions.FlowReadException as e:231                raise exceptions.OptionsError(str(e))232            self.start_replay(flows)233 234        if "client_replay_concurrency" in updated:235            if ctx.options.client_replay_concurrency not in [-1, 1]:236                raise exceptions.OptionsError(237                    "Currently the only valid client_replay_concurrency values are -1 and 1."238                )239 240    @command.command("replay.client.count")241    def count(self) -> int:242        """243        Approximate number of flows queued for replay.244        """245        return self.queue.qsize() + int(bool(self.inflight))246 247    @command.command("replay.client.stop")248    def stop_replay(self) -> None:249        """250        Clear the replay queue.251        """252        updated = []253        while True:254            try:255                f = self.queue.get_nowait()256            except asyncio.QueueEmpty:257                break258            else:259                self.queue.task_done()260                f.revert()261                updated.append(f)262 263        ctx.master.addons.trigger(UpdateHook(updated))264        logger.log(ALERT, "Client replay queue cleared.")265 266    @command.command("replay.client")267    def start_replay(self, flows: Sequence[flow.Flow]) -> None:268        """269        Add flows to the replay queue, skipping flows that can't be replayed.270        """271        updated: list[http.HTTPFlow] = []272        for f in flows:273            err = self.check(f)274            if err:275                logger.warning(err)276                continue277 278            http_flow = cast(http.HTTPFlow, f)279 280            # Prepare the flow for replay281            http_flow.backup()282            http_flow.is_replay = "request"283            http_flow.response = None284            http_flow.error = None285            self.queue.put_nowait(http_flow)286            updated.append(http_flow)287        ctx.master.addons.trigger(UpdateHook(updated))288 289    @command.command("replay.client.file")290    def load_file(self, path: mitmproxy.types.Path) -> None:291        """292        Load flows from file, and add them to the replay queue.293        """294        try:295            flows = io.read_flows_from_paths([path])296        except exceptions.FlowReadException as e:297            raise exceptions.CommandError(str(e))298        self.start_replay(flows)299