codekingpro/portable-devtools

1import hashlib2import logging3import urllib4from collections.abc import Hashable5from collections.abc import Sequence6from typing import Any7 8import mitmproxy.types9from mitmproxy import command10from mitmproxy import ctx11from mitmproxy import exceptions12from mitmproxy import flow13from mitmproxy import hooks14from mitmproxy import http15from mitmproxy import io16 17logger = logging.getLogger(__name__)18 19HASH_OPTIONS = [20    "server_replay_ignore_content",21    "server_replay_ignore_host",22    "server_replay_ignore_params",23    "server_replay_ignore_payload_params",24    "server_replay_ignore_port",25    "server_replay_use_headers",26]27 28 29class ServerPlayback:30    flowmap: dict[Hashable, list[http.HTTPFlow]]31    configured: bool32 33    def __init__(self):34        self.flowmap = {}35        self.configured = False36 37    def load(self, loader):38        loader.add_option(39            "server_replay_kill_extra",40            bool,41            False,42            "Kill extra requests during replay (for which no replayable response was found)."43            "[Deprecated, prefer to use server_replay_extra='kill']",44        )45        loader.add_option(46            "server_replay_extra",47            str,48            "forward",49            "Behaviour for extra requests during replay for which no replayable response was found. "50            "Setting a numeric string value will return an empty HTTP response with the respective status code.",51            choices=["forward", "kill", "204", "400", "404", "500"],52        )53        loader.add_option(54            "server_replay_reuse",55            bool,56            False,57            """58            Don't remove flows from server replay state after use. This makes it59            possible to replay same response multiple times.60            """,61        )62        loader.add_option(63            "server_replay_nopop",64            bool,65            False,66            """67            Deprecated alias for `server_replay_reuse`.68            """,69        )70        loader.add_option(71            "server_replay_refresh",72            bool,73            True,74            """75            Refresh server replay responses by adjusting date, expires and76            last-modified headers, as well as adjusting cookie expiration.77            """,78        )79        loader.add_option(80            "server_replay_use_headers",81            Sequence[str],82            [],83            """84            Request headers that need to match while searching for a saved flow85            to replay.86            """,87        )88        loader.add_option(89            "server_replay",90            Sequence[str],91            [],92            "Replay server responses from a saved file.",93        )94        loader.add_option(95            "server_replay_ignore_content",96            bool,97            False,98            "Ignore request content while searching for a saved flow to replay.",99        )100        loader.add_option(101            "server_replay_ignore_params",102            Sequence[str],103            [],104            """105            Request parameters to be ignored while searching for a saved flow106            to replay.107            """,108        )109        loader.add_option(110            "server_replay_ignore_payload_params",111            Sequence[str],112            [],113            """114            Request payload parameters (application/x-www-form-urlencoded or115            multipart/form-data) to be ignored while searching for a saved flow116            to replay.117            """,118        )119        loader.add_option(120            "server_replay_ignore_host",121            bool,122            False,123            """124            Ignore request destination host while searching for a saved flow125            to replay.126            """,127        )128        loader.add_option(129            "server_replay_ignore_port",130            bool,131            False,132            """133            Ignore request destination port while searching for a saved flow134            to replay.135            """,136        )137 138    @command.command("replay.server")139    def load_flows(self, flows: Sequence[flow.Flow]) -> None:140        """141        Replay server responses from flows.142        """143        self.flowmap = {}144        self.add_flows(flows)145 146    @command.command("replay.server.add")147    def add_flows(self, flows: Sequence[flow.Flow]) -> None:148        """149        Add responses from flows to server replay list.150        """151        for f in flows:152            if isinstance(f, http.HTTPFlow):153                lst = self.flowmap.setdefault(self._hash(f), [])154                lst.append(f)155        ctx.master.addons.trigger(hooks.UpdateHook([]))156 157    @command.command("replay.server.file")158    def load_file(self, path: mitmproxy.types.Path) -> None:159        try:160            flows = io.read_flows_from_paths([path])161        except exceptions.FlowReadException as e:162            raise exceptions.CommandError(str(e))163        self.load_flows(flows)164 165    @command.command("replay.server.stop")166    def clear(self) -> None:167        """168        Stop server replay.169        """170        self.flowmap = {}171        ctx.master.addons.trigger(hooks.UpdateHook([]))172 173    @command.command("replay.server.count")174    def count(self) -> int:175        return sum(len(i) for i in self.flowmap.values())176 177    def _hash(self, flow: http.HTTPFlow) -> Hashable:178        """179        Calculates a loose hash of the flow request.180        """181        r = flow.request182        _, _, path, _, query, _ = urllib.parse.urlparse(r.url)183        queriesArray = urllib.parse.parse_qsl(query, keep_blank_values=True)184 185        key: list[Any] = [str(r.scheme), str(r.method), str(path)]186        if not ctx.options.server_replay_ignore_content:187            if ctx.options.server_replay_ignore_payload_params and r.multipart_form:188                key.extend(189                    (k, v)190                    for k, v in r.multipart_form.items(multi=True)191                    if k.decode(errors="replace")192                    not in ctx.options.server_replay_ignore_payload_params193                )194            elif ctx.options.server_replay_ignore_payload_params and r.urlencoded_form:195                key.extend(196                    (k, v)197                    for k, v in r.urlencoded_form.items(multi=True)198                    if k not in ctx.options.server_replay_ignore_payload_params199                )200            else:201                key.append(str(r.raw_content))202 203        if not ctx.options.server_replay_ignore_host:204            key.append(r.pretty_host)205        if not ctx.options.server_replay_ignore_port:206            key.append(r.port)207 208        filtered = []209        ignore_params = ctx.options.server_replay_ignore_params or []210        for p in queriesArray:211            if p[0] not in ignore_params:212                filtered.append(p)213        for p in filtered:214            key.append(p[0])215            key.append(p[1])216 217        if ctx.options.server_replay_use_headers:218            headers = []219            for i in ctx.options.server_replay_use_headers:220                v = r.headers.get(i)221                headers.append((i, v))222            key.append(headers)223        return hashlib.sha256(repr(key).encode("utf8", "surrogateescape")).digest()224 225    def next_flow(self, flow: http.HTTPFlow) -> http.HTTPFlow | None:226        """227        Returns the next flow object, or None if no matching flow was228        found.229        """230        hash = self._hash(flow)231        if hash in self.flowmap:232            if ctx.options.server_replay_reuse or ctx.options.server_replay_nopop:233                return next(234                    (flow for flow in self.flowmap[hash] if flow.response), None235                )236            else:237                ret = self.flowmap[hash].pop(0)238                while not ret.response:239                    if self.flowmap[hash]:240                        ret = self.flowmap[hash].pop(0)241                    else:242                        del self.flowmap[hash]243                        return None244                if not self.flowmap[hash]:245                    del self.flowmap[hash]246                return ret247        else:248            return None249 250    def configure(self, updated):251        if ctx.options.server_replay_kill_extra:252            logger.warning(253                "server_replay_kill_extra has been deprecated, "254                "please update your config to use server_replay_extra='kill'."255            )256        if ctx.options.server_replay_nopop:  # pragma: no cover257            logger.error(258                "server_replay_nopop has been renamed to server_replay_reuse, please update your config."259            )260        if not self.configured and ctx.options.server_replay:261            self.configured = True262            try:263                flows = io.read_flows_from_paths(ctx.options.server_replay)264            except exceptions.FlowReadException as e:265                raise exceptions.OptionsError(str(e))266            self.load_flows(flows)267        if any(option in updated for option in HASH_OPTIONS):268            self.recompute_hashes()269 270    def recompute_hashes(self) -> None:271        """272        Rebuild flowmap if the hashing method has changed during execution,273        see https://github.com/mitmproxy/mitmproxy/issues/4506274        """275        flows = [flow for lst in self.flowmap.values() for flow in lst]276        self.load_flows(flows)277 278    def request(self, f: http.HTTPFlow) -> None:279        if self.flowmap:280            rflow = self.next_flow(f)281            if rflow:282                assert rflow.response283                response = rflow.response.copy()284                if ctx.options.server_replay_refresh:285                    response.refresh()286                f.response = response287                f.is_replay = "response"288            elif (289                ctx.options.server_replay_kill_extra290                or ctx.options.server_replay_extra == "kill"291            ):292                logging.warning(293                    "server_playback: killed non-replay request {}".format(294                        f.request.url295                    )296                )297                f.kill()298            elif ctx.options.server_replay_extra != "forward":299                logging.warning(300                    "server_playback: returned {} non-replay request {}".format(301                        ctx.options.server_replay_extra, f.request.url302                    )303                )304                f.response = http.Response.make(int(ctx.options.server_replay_extra))305                f.is_replay = "response"306