codekingpro/portable-devtools

1import asyncio2import logging3import os.path4import sys5from typing import BinaryIO6from typing import Optional7 8from mitmproxy import command9from mitmproxy import ctx10from mitmproxy import exceptions11from mitmproxy import flowfilter12from mitmproxy import io13from mitmproxy.utils import asyncio_utils14 15logger = logging.getLogger(__name__)16 17 18class ReadFile:19    """20    An addon that handles reading from file on startup.21    """22 23    def __init__(self):24        self.filter = None25        self._read_task: asyncio.Task | None = None26 27    def load(self, loader):28        loader.add_option("rfile", Optional[str], None, "Read flows from file.")29        loader.add_option(30            "readfile_filter", Optional[str], None, "Read only matching flows."31        )32 33    def configure(self, updated):34        if "readfile_filter" in updated:35            if ctx.options.readfile_filter:36                try:37                    self.filter = flowfilter.parse(ctx.options.readfile_filter)38                except ValueError as e:39                    raise exceptions.OptionsError(str(e)) from e40            else:41                self.filter = None42 43    async def load_flows(self, fo: BinaryIO) -> int:44        cnt = 045        freader = io.FlowReader(fo)46        try:47            for flow in freader.stream():48                if self.filter and not self.filter(flow):49                    continue50                await ctx.master.load_flow(flow)51                cnt += 152        except (OSError, exceptions.FlowReadException) as e:53            if cnt:54                logging.warning("Flow file corrupted - loaded %i flows." % cnt)55            else:56                logging.error("Flow file corrupted.")57            raise exceptions.FlowReadException(str(e)) from e58        else:59            return cnt60 61    async def load_flows_from_path(self, path: str) -> int:62        path = os.path.expanduser(path)63        try:64            with open(path, "rb") as f:65                return await self.load_flows(f)66        except OSError as e:67            logging.error(f"Cannot load flows: {e}")68            raise exceptions.FlowReadException(str(e)) from e69 70    async def doread(self, rfile: str) -> None:71        try:72            await self.load_flows_from_path(rfile)73        except exceptions.FlowReadException as e:74            logger.exception(f"Failed to read {ctx.options.rfile}: {e}")75 76    def running(self):77        if ctx.options.rfile:78            self._read_task = asyncio_utils.create_task(79                self.doread(ctx.options.rfile),80                name="readfile",81                keep_ref=False,82            )83 84    @command.command("readfile.reading")85    def reading(self) -> bool:86        return bool(self._read_task and not self._read_task.done())87 88 89class ReadFileStdin(ReadFile):90    """Support the special case of "-" for reading from stdin"""91 92    async def load_flows_from_path(self, path: str) -> int:93        if path == "-":  # pragma: no cover94            # Need to think about how to test this. This function is scheduled95            # onto the event loop, where a sys.stdin mock has no effect.96            return await self.load_flows(sys.stdin.buffer)97        else:98            return await super().load_flows_from_path(path)99