Addon for mitmproxy to export request as python code with httpx

mitmproxy_export_httpx.py

import textwrap
import ast
import json

from typing import Any, Union
from urllib.parse import urlparse, urlunparse, parse_qs

from mitmproxy.addonmanager import Loader
from mitmproxy.addons.export import formats
from mitmproxy.flow import Flow
from mitmproxy import ctx


def load(loader: Loader):
    formats["httpx"] = httpx_format
    ctx.log.info("export-httpx loaded")


def httpx_format(flow: Flow) -> str:
    import_httpx = ast.Import(names=[ast.alias(name="httpx", as_=None)])

    client_args = []

    request_args = [
        ast.Constant(value=flow.request.method.upper()),
    ]

    if flow.response and flow.response.http_version == "HTTP/2.0":
        client_args.append(ast.keyword(arg="http2", value=ast.Constant(value=True)))

    pretty_url = flow.request.pretty_url
    parsed_url = urlparse(pretty_url)
    query_params = parse_qs(parsed_url.query)
    result_url = urlunparse(parsed_url._replace(query=""))

    request_args.append(
        ast.Constant(value=result_url),
    )

    if query_params:
        unfolded_single_params = {k: v[0] if len(v) == 1 else v for k, v in query_params.items()}
        request_args.append(ast.keyword(arg="params", value=value_to_ast(unfolded_single_params)))

    if flow.request.headers:
        headers = ast.List(
            elts=[
                ast.Tuple(elts=[ast.Constant(value=k), ast.Constant(value=v)])
                for k, v in flow.request.headers.items()
            ],
        )
        request_args.append(ast.keyword(arg="headers", value=headers))

    if flow.request.content:
        content = flow.request.content
        content_type = flow.request.headers.get("content-type", "")

        try:
            if "application/json" in content_type:
                parsed_json = json.loads(content.decode("utf-8"))
                ast_json = value_to_ast(parsed_json)
                keyword = ast.keyword(arg="json", value=ast_json)
            else:
                keyword = ast.keyword(
                    arg="content", value=ast.Constant(value=content.decode("utf-8"))
                )
        except UnicodeDecodeError:
            keyword = ast.keyword(arg="content", value=ast.Constant(value=content))

        request_args.append(keyword)

    module = ast.Module(
        body=[
            import_httpx,
            ast.Assign(
                targets=[ast.Name(id="client", ctx=ast.Store())],
                value=ast.Call(
                    func=ast.Attribute(
                        value=ast.Name(id="httpx", ctx=ast.Load()),
                        attr="Client",
                        ctx=ast.Load(),
                    ),
                    args=client_args,
                    keywords=[],
                ),
            ),
            ast.Assign(
                targets=[ast.Name(id="request", ctx=ast.Store())],
                value=ast.Call(
                    func=ast.Attribute(
                        value=ast.Name(id="httpx", ctx=ast.Load()),
                        attr="Request",
                        ctx=ast.Load(),
                    ),
                    args=request_args,
                    keywords=[],
                ),
            ),
            ast.Assign(
                targets=[ast.Name(id="response", ctx=ast.Store())],
                value=ast.Call(
                    func=ast.Attribute(
                        value=ast.Name(id="client", ctx=ast.Load()),
                        attr="send",
                        ctx=ast.Load(),
                    ),
                    args=[ast.Name(id="request", ctx=ast.Load())],
                    keywords=[],
                ),
            ),
            ast.Expr(
                value=ast.Call(
                    func=ast.Name(id="print", ctx=ast.Load()),
                    args=[
                        ast.Attribute(
                            value=ast.Name(id="response", ctx=ast.Load()),
                            attr="status_code",
                            ctx=ast.Load(),
                        )
                    ],
                    keywords=[],
                )
            ),
            ast.Expr(
                value=ast.Call(
                    func=ast.Name(id="print", ctx=ast.Load()),
                    args=[
                        ast.Attribute(
                            value=ast.Name(id="response", ctx=ast.Load()),
                            attr="text",
                            ctx=ast.Load(),
                        )
                    ],
                    keywords=[],
                )
            ),
        ],
        type_ignores=[],
    )

    ast.fix_missing_locations(module)
    # ctx.log.info(ast.dump(module, indent=4))
    code = ast.unparse(module)

    try:
        import black

        code = black.format_str(code, mode=black.FileMode())
    except ImportError:
        pass

    return code.strip()


def value_to_ast(value: Any) -> Union[ast.Constant, ast.Dict, ast.List]:
    if isinstance(value, dict):
        return ast.Dict(
            keys=[ast.Constant(value=k) for k in value.keys()],
            values=[value_to_ast(v) for v in value.values()],
        )
    elif isinstance(value, list):
        return ast.List(
            elts=[value_to_ast(item) for item in value],
            ctx=ast.Load(),
        )
    else:
        return ast.Constant(value=value)