codekingpro/portable-devtools

1from __future__ import annotations2 3import asyncio4 5from mitmproxy import ctx6from mitmproxy.utils import asyncio_utils7 8 9class KeepServing:10    def load(self, loader):11        loader.add_option(12            "keepserving",13            bool,14            False,15            """16            Continue serving after client playback, server playback or file17            read. This option is ignored by interactive tools, which always keep18            serving.19            """,20        )21 22    def keepgoing(self) -> bool:23        # Checking for proxyserver.active_connections is important for server replay,24        # the addon may report that replay is finished but not the entire response has been sent yet.25        # (https://github.com/mitmproxy/mitmproxy/issues/7569)26        checks = [27            "readfile.reading",28            "replay.client.count",29            "replay.server.count",30            "proxyserver.active_connections",31        ]32        return any([ctx.master.commands.call(c) for c in checks])33 34    def shutdown(self):  # pragma: no cover35        ctx.master.shutdown()36 37    async def watch(self):38        while True:39            await asyncio.sleep(0.1)40            if not self.keepgoing():41                self.shutdown()42 43    def running(self):44        opts = [45            ctx.options.client_replay,46            ctx.options.server_replay,47            ctx.options.rfile,48        ]49        if any(opts) and not ctx.options.keepserving:50            asyncio_utils.create_task(51                self.watch(),52                name="keepserving",53                keep_ref=True,54            )55