codekingpro/portable-devtools

1import csv2import io3import logging4import os.path5from collections.abc import Sequence6from typing import Any7 8import pyperclip9 10import mitmproxy.types11from mitmproxy import certs12from mitmproxy import command13from mitmproxy import exceptions14from mitmproxy import flow15from mitmproxy import http16from mitmproxy.log import ALERT17 18logger = logging.getLogger(__name__)19 20 21def headername(spec: str):22    if not (spec.startswith("header[") and spec.endswith("]")):23        raise exceptions.CommandError("Invalid header spec: %s" % spec)24    return spec[len("header[") : -1].strip()25 26 27def is_addr(v):28    return isinstance(v, tuple) and len(v) > 129 30 31def extract(cut: str, f: flow.Flow) -> str | bytes:32    # Hack for https://github.com/mitmproxy/mitmproxy/issues/6721:33    # Make "save body" keybind work for WebSocket flows.34    # Ideally the keybind would be smarter and this here can get removed.35    if (36        isinstance(f, http.HTTPFlow)37        and f.websocket38        and cut in ("request.content", "response.content")39    ):40        return f.websocket._get_formatted_messages()41 42    path = cut.split(".")43    current: Any = f44    for i, spec in enumerate(path):45        if spec.startswith("_"):46            raise exceptions.CommandError("Can't access internal attribute %s" % spec)47 48        part = getattr(current, spec, None)49        if i == len(path) - 1:50            if spec == "port" and is_addr(current):51                return str(current[1])52            if spec == "host" and is_addr(current):53                return str(current[0])54            elif spec.startswith("header["):55                if not current:56                    return ""57                return current.headers.get(headername(spec), "")58            elif isinstance(part, bytes):59                return part60            elif isinstance(part, bool):61                return "true" if part else "false"62            elif isinstance(part, certs.Cert):  # pragma: no cover63                return part.to_pem().decode("ascii")64            elif (65                isinstance(part, list)66                and len(part) > 067                and isinstance(part[0], certs.Cert)68            ):69                # TODO: currently this extracts only the very first cert as PEM-encoded string.70                return part[0].to_pem().decode("ascii")71        current = part72    return str(current or "")73 74 75def extract_str(cut: str, f: flow.Flow) -> str:76    ret = extract(cut, f)77    if isinstance(ret, bytes):78        return repr(ret)79    else:80        return ret81 82 83class Cut:84    @command.command("cut")85    def cut(86        self,87        flows: Sequence[flow.Flow],88        cuts: mitmproxy.types.CutSpec,89    ) -> mitmproxy.types.Data:90        """91        Cut data from a set of flows. Cut specifications are attribute paths92        from the base of the flow object, with a few conveniences - "port"93        and "host" retrieve parts of an address tuple, ".header[key]"94        retrieves a header value. Return values converted to strings or95        bytes: SSL certificates are converted to PEM format, bools are "true"96        or "false", "bytes" are preserved, and all other values are97        converted to strings.98        """99        ret: list[list[str | bytes]] = []100        for f in flows:101            ret.append([extract(c, f) for c in cuts])102        return ret  # type: ignore103 104    @command.command("cut.save")105    def save(106        self,107        flows: Sequence[flow.Flow],108        cuts: mitmproxy.types.CutSpec,109        path: mitmproxy.types.Path,110    ) -> None:111        """112        Save cuts to file. If there are multiple flows or cuts, the format113        is UTF-8 encoded CSV. If there is exactly one row and one column,114        the data is written to file as-is, with raw bytes preserved. If the115        path is prefixed with a "+", values are appended if there is an116        existing file.117        """118        append = False119        if path.startswith("+"):120            append = True121            epath = os.path.expanduser(path[1:])122            path = mitmproxy.types.Path(epath)123        try:124            if len(cuts) == 1 and len(flows) == 1:125                with open(path, "ab" if append else "wb") as fp:126                    if fp.tell() > 0:127                        # We're appending to a file that already exists and has content128                        fp.write(b"\n")129                    v = extract(cuts[0], flows[0])130                    if isinstance(v, bytes):131                        fp.write(v)132                    else:133                        fp.write(v.encode("utf8"))134                logger.log(ALERT, "Saved single cut.")135            else:136                with open(137                    path, "a" if append else "w", newline="", encoding="utf8"138                ) as tfp:139                    writer = csv.writer(tfp)140                    for f in flows:141                        vals = [extract_str(c, f) for c in cuts]142                        writer.writerow(vals)143                logger.log(144                    ALERT,145                    "Saved %s cuts over %d flows as CSV." % (len(cuts), len(flows)),146                )147        except OSError as e:148            logger.error(str(e))149 150    @command.command("cut.clip")151    def clip(152        self,153        flows: Sequence[flow.Flow],154        cuts: mitmproxy.types.CutSpec,155    ) -> None:156        """157        Send cuts to the clipboard. If there are multiple flows or cuts, the158        format is UTF-8 encoded CSV. If there is exactly one row and one159        column, the data is written to file as-is, with raw bytes preserved.160        """161        v: str | bytes162        fp = io.StringIO(newline="")163        if len(cuts) == 1 and len(flows) == 1:164            v = extract_str(cuts[0], flows[0])165            fp.write(v)166            logger.log(ALERT, "Clipped single cut.")167        else:168            writer = csv.writer(fp)169            for f in flows:170                vals = [extract_str(c, f) for c in cuts]171                writer.writerow(vals)172            logger.log(ALERT, "Clipped %s cuts as CSV." % len(cuts))173        try:174            pyperclip.copy(fp.getvalue())175        except pyperclip.PyperclipException as e:176            logger.error(str(e))177