codekingpro/portable-devtools

1from collections.abc import Sequence2from typing import NamedTuple3 4from mitmproxy import ctx5from mitmproxy import exceptions6from mitmproxy import flowfilter7from mitmproxy import http8from mitmproxy import version9from mitmproxy.net.http.status_codes import NO_RESPONSE10 11 12class BlockSpec(NamedTuple):13    matches: flowfilter.TFilter14    status_code: int15 16 17def parse_spec(option: str) -> BlockSpec:18    """19    Parses strings in the following format, enforces number of segments:20 21        /flow-filter/status22 23    """24    sep, rem = option[0], option[1:]25 26    parts = rem.split(sep, 2)27    if len(parts) != 2:28        raise ValueError("Invalid number of parameters (2 are expected)")29    flow_patt, status = parts30    try:31        status_code = int(status)32    except ValueError:33        raise ValueError(f"Invalid HTTP status code: {status}")34    flow_filter = flowfilter.parse(flow_patt)35 36    return BlockSpec(matches=flow_filter, status_code=status_code)37 38 39class BlockList:40    def __init__(self) -> None:41        self.items: list[BlockSpec] = []42 43    def load(self, loader):44        loader.add_option(45            "block_list",46            Sequence[str],47            [],48            """49            Block matching requests and return an empty response with the specified HTTP status.50            Option syntax is "/flow-filter/status-code", where flow-filter describes51            which requests this rule should be applied to and status-code is the HTTP status code to return for52            blocked requests. The separator ("/" in the example) can be any character.53            Setting a non-standard status code of 444 will close the connection without sending a response.54            """,55        )56 57    def configure(self, updated):58        if "block_list" in updated:59            self.items = []60            for option in ctx.options.block_list:61                try:62                    spec = parse_spec(option)63                except ValueError as e:64                    raise exceptions.OptionsError(65                        f"Cannot parse block_list option {option}: {e}"66                    ) from e67                self.items.append(spec)68 69    def request(self, flow: http.HTTPFlow) -> None:70        if flow.response or flow.error or not flow.live:71            return72 73        for spec in self.items:74            if spec.matches(flow):75                flow.metadata["blocklisted"] = True76                if spec.status_code == NO_RESPONSE:77                    flow.kill()78                else:79                    flow.response = http.Response.make(80                        spec.status_code, headers={"Server": version.MITMPROXY}81                    )82