codekingpro/portable-devtools

1import logging2import re3from collections.abc import Sequence4from pathlib import Path5from typing import NamedTuple6 7from mitmproxy import ctx8from mitmproxy import exceptions9from mitmproxy import flowfilter10from mitmproxy import http11from mitmproxy.http import Headers12from mitmproxy.utils import strutils13from mitmproxy.utils.spec import parse_spec14 15 16class ModifySpec(NamedTuple):17    matches: flowfilter.TFilter18    subject: bytes19    replacement_str: str20 21    def read_replacement(self) -> bytes:22        """23        Process the replacement str. This usually just involves converting it to bytes.24        However, if it starts with `@`, we interpret the rest as a file path to read from.25 26        Raises:27            - IOError if the file cannot be read.28        """29        if self.replacement_str.startswith("@"):30            return Path(self.replacement_str[1:]).expanduser().read_bytes()31        else:32            # We could cache this at some point, but unlikely to be a problem.33            return strutils.escaped_str_to_bytes(self.replacement_str)34 35 36def parse_modify_spec(option: str, subject_is_regex: bool) -> ModifySpec:37    flow_filter, subject_str, replacement = parse_spec(option)38 39    subject = strutils.escaped_str_to_bytes(subject_str)40    if subject_is_regex:41        try:42            re.compile(subject)43        except re.error as e:44            raise ValueError(f"Invalid regular expression {subject!r} ({e})")45 46    spec = ModifySpec(flow_filter, subject, replacement)47 48    try:49        spec.read_replacement()50    except OSError as e:51        raise ValueError(f"Invalid file path: {replacement[1:]} ({e})")52 53    return spec54 55 56class ModifyHeaders:57    def __init__(self) -> None:58        self.replacements: list[ModifySpec] = []59 60    def load(self, loader):61        loader.add_option(62            "modify_headers",63            Sequence[str],64            [],65            """66            Header modify pattern of the form "[/flow-filter]/header-name/[@]header-value", where the67            separator can be any character. The @ allows to provide a file path that is used to read68            the header value string. An empty header-value removes existing header-name headers.69            """,70        )71 72    def configure(self, updated):73        if "modify_headers" in updated:74            self.replacements = []75            for option in ctx.options.modify_headers:76                try:77                    spec = parse_modify_spec(option, False)78                except ValueError as e:79                    raise exceptions.OptionsError(80                        f"Cannot parse modify_headers option {option}: {e}"81                    ) from e82                self.replacements.append(spec)83 84    def requestheaders(self, flow):85        if flow.response or flow.error or not flow.live:86            return87        self.run(flow, flow.request.headers)88 89    def responseheaders(self, flow):90        if flow.error or not flow.live:91            return92        self.run(flow, flow.response.headers)93 94    def run(self, flow: http.HTTPFlow, hdrs: Headers) -> None:95        matches = []96 97        # first check all the filters against the original, unmodified flow98        for spec in self.replacements:99            matches.append(spec.matches(flow))100 101        # unset all specified headers102        for i, spec in enumerate(self.replacements):103            if matches[i]:104                hdrs.pop(spec.subject, None)105 106        # set all specified headers if the replacement string is not empty107 108        for i, spec in enumerate(self.replacements):109            if matches[i]:110                try:111                    replacement = spec.read_replacement()112                except OSError as e:113                    logging.warning(f"Could not read replacement file: {e}")114                    continue115                else:116                    if replacement:117                        hdrs.add(spec.subject, replacement)118