codekingpro/portable-devtools

1import logging2import re3from collections.abc import Sequence4 5from mitmproxy import ctx6from mitmproxy import exceptions7from mitmproxy.addons.modifyheaders import ModifySpec8from mitmproxy.addons.modifyheaders import parse_modify_spec9from mitmproxy.log import ALERT10 11logger = logging.getLogger(__name__)12 13 14class ModifyBody:15    def __init__(self) -> None:16        self.replacements: list[ModifySpec] = []17 18    def load(self, loader):19        loader.add_option(20            "modify_body",21            Sequence[str],22            [],23            """24            Replacement pattern of the form "[/flow-filter]/regex/[@]replacement", where25            the separator can be any character. The @ allows to provide a file path that26            is used to read the replacement string.27            """,28        )29 30    def configure(self, updated):31        if "modify_body" in updated:32            self.replacements = []33            for option in ctx.options.modify_body:34                try:35                    spec = parse_modify_spec(option, True)36                except ValueError as e:37                    raise exceptions.OptionsError(38                        f"Cannot parse modify_body option {option}: {e}"39                    ) from e40 41                self.replacements.append(spec)42 43        stream_and_modify_conflict = (44            ctx.options.modify_body45            and ctx.options.stream_large_bodies46            and ("modify_body" in updated or "stream_large_bodies" in updated)47        )48        if stream_and_modify_conflict:49            logger.log(50                ALERT,51                "Both modify_body and stream_large_bodies are active. "52                "Streamed bodies will not be modified.",53            )54 55    def request(self, flow):56        if flow.response or flow.error or not flow.live:57            return58        self.run(flow)59 60    def response(self, flow):61        if flow.error or not flow.live:62            return63        self.run(flow)64 65    def run(self, flow):66        for spec in self.replacements:67            if spec.matches(flow):68                try:69                    replacement = spec.read_replacement()70                except OSError as e:71                    logging.warning(f"Could not read replacement file: {e}")72                    continue73                if flow.response:74                    flow.response.content = re.sub(75                        spec.subject,76                        lambda _: replacement,77                        flow.response.content,78                        flags=re.DOTALL,79                    )80                else:81                    flow.request.content = re.sub(82                        spec.subject,83                        lambda _: replacement,84                        flow.request.content,85                        flags=re.DOTALL,86                    )87