codekingpro/portable-devtools

1import logging2import shlex3from collections.abc import Callable4from collections.abc import Sequence5 6import pyperclip7 8import mitmproxy.types9from mitmproxy import command10from mitmproxy import ctx11from mitmproxy import exceptions12from mitmproxy import flow13from mitmproxy import http14from mitmproxy.net.http.http1 import assemble15from mitmproxy.utils import strutils16 17 18def cleanup_request(f: flow.Flow) -> http.Request:19    if not getattr(f, "request", None):20        raise exceptions.CommandError("Can't export flow with no request.")21    assert isinstance(f, http.HTTPFlow)22    request = f.request.copy()23    request.decode(strict=False)24    return request25 26 27def pop_headers(request: http.Request) -> None:28    """Remove some headers that are redundant for curl/httpie export."""29    request.headers.pop("content-length", None)30 31    if request.headers.get("host", "") == request.host:32        request.headers.pop("host")33    if request.headers.get(":authority", "") == request.host:34        request.headers.pop(":authority")35 36 37def cleanup_response(f: flow.Flow) -> http.Response:38    if not getattr(f, "response", None):39        raise exceptions.CommandError("Can't export flow with no response.")40    assert isinstance(f, http.HTTPFlow)41    response = f.response.copy()  # type: ignore42    response.decode(strict=False)43    return response44 45 46def request_content_for_console(request: http.Request) -> str:47    try:48        text = request.get_text(strict=True)49        assert text50    except ValueError:51        # shlex.quote doesn't support a bytes object52        # see https://github.com/python/cpython/pull/1087153        raise exceptions.CommandError("Request content must be valid unicode")54    escape_control_chars = {chr(i): f"\\x{i:02x}" for i in range(32)}55    escaped_text = "".join(escape_control_chars.get(x, x) for x in text)56    if any(char in escape_control_chars for char in text):57        # Escaped chars need to be unescaped by the shell to be properly inperpreted by curl and httpie58        return f'"$(printf {shlex.quote(escaped_text)})"'59 60    return shlex.quote(escaped_text)61 62 63def curl_command(f: flow.Flow) -> str:64    request = cleanup_request(f)65    pop_headers(request)66 67    args = ["curl"]68 69    server_addr = f.server_conn.peername[0] if f.server_conn.peername else None70 71    if (72        ctx.options.export_preserve_original_ip73        and server_addr74        and request.pretty_host != server_addr75    ):76        resolve = f"{request.pretty_host}:{request.port}:[{server_addr}]"77        args.append("--resolve")78        args.append(resolve)79 80    for k, v in request.headers.items(multi=True):81        if k.lower() == "accept-encoding":82            args.append("--compressed")83        else:84            args += ["-H", f"{k}: {v}"]85 86    if request.method != "GET":87        if not request.content:88            # curl will not calculate content-length if there is no content89            # some server/verb combinations require content-length headers90            # (ex. nginx and POST)91            args += ["-H", "content-length: 0"]92 93        args += ["-X", request.method]94 95    args.append(request.pretty_url)96 97    command = " ".join(shlex.quote(arg) for arg in args)98    if request.content:99        command += f" -d {request_content_for_console(request)}"100    return command101 102 103def httpie_command(f: flow.Flow) -> str:104    request = cleanup_request(f)105    pop_headers(request)106 107    # TODO: Once https://github.com/httpie/httpie/issues/414 is implemented, we108    # should ensure we always connect to the IP address specified in the flow,109    # similar to how it's done in curl_command.110    url = request.pretty_url111 112    args = ["http", request.method, url]113    for k, v in request.headers.items(multi=True):114        args.append(f"{k}: {v}")115    cmd = " ".join(shlex.quote(arg) for arg in args)116    if request.content:117        cmd += " <<< " + request_content_for_console(request)118    return cmd119 120 121def raw_request(f: flow.Flow) -> bytes:122    request = cleanup_request(f)123    if request.raw_content is None:124        raise exceptions.CommandError("Request content missing.")125    return assemble.assemble_request(request)126 127 128def raw_response(f: flow.Flow) -> bytes:129    response = cleanup_response(f)130    if response.raw_content is None:131        raise exceptions.CommandError("Response content missing.")132    return assemble.assemble_response(response)133 134 135def raw(f: flow.Flow, separator=b"\r\n\r\n") -> bytes:136    """Return either the request or response if only one exists, otherwise return both"""137    request_present = (138        isinstance(f, http.HTTPFlow) and f.request and f.request.raw_content is not None139    )140    response_present = (141        isinstance(f, http.HTTPFlow)142        and f.response143        and f.response.raw_content is not None144    )145 146    if request_present and response_present:147        parts = [raw_request(f), raw_response(f)]148        if isinstance(f, http.HTTPFlow) and f.websocket:149            parts.append(f.websocket._get_formatted_messages())150        return separator.join(parts)151    elif request_present:152        return raw_request(f)153    elif response_present:154        return raw_response(f)155    else:156        raise exceptions.CommandError("Can't export flow with no request or response.")157 158 159formats: dict[str, Callable[[flow.Flow], str | bytes]] = dict(160    curl=curl_command,161    httpie=httpie_command,162    raw=raw,163    raw_request=raw_request,164    raw_response=raw_response,165)166 167 168class Export:169    def load(self, loader):170        loader.add_option(171            "export_preserve_original_ip",172            bool,173            False,174            """175            When exporting a request as an external command, make an effort to176            connect to the same IP as in the original request. This helps with177            reproducibility in cases where the behaviour depends on the178            particular host we are connecting to. Currently this only affects179            curl exports.180            """,181        )182 183    @command.command("export.formats")184    def formats(self) -> Sequence[str]:185        """186        Return a list of the supported export formats.187        """188        return list(sorted(formats.keys()))189 190    @command.command("export.file")191    def file(self, format: str, flow: flow.Flow, path: mitmproxy.types.Path) -> None:192        """193        Export a flow to path.194        """195        if format not in formats:196            raise exceptions.CommandError("No such export format: %s" % format)197        v = formats[format](flow)198        try:199            with open(path, "wb") as fp:200                if isinstance(v, bytes):201                    fp.write(v)202                else:203                    fp.write(v.encode("utf-8", "surrogateescape"))204        except OSError as e:205            logging.error(str(e))206 207    @command.command("export.clip")208    def clip(self, format: str, f: flow.Flow) -> None:209        """210        Export a flow to the system clipboard.211        """212        content = self.export_str(format, f)213        try:214            pyperclip.copy(content)215        except pyperclip.PyperclipException as e:216            logging.error(str(e))217 218    @command.command("export")219    def export_str(self, format: str, f: flow.Flow) -> str:220        """221        Export a flow and return the result.222        """223        if format not in formats:224            raise exceptions.CommandError("No such export format: %s" % format)225 226        content = formats[format](f)227        # The individual formatters may return surrogate-escaped UTF-8, but that may blow up in later steps.228        # For example, pyperclip on macOS does not like surrogates.229        # To fix this, We first surrogate-encode and then backslash-decode.230        content = strutils.always_bytes(content, "utf8", "surrogateescape")231        content = strutils.always_str(content, "utf8", "backslashreplace")232        return content233