mitmproxy/mitmproxy/flow_export.py

import json
from textwrap import dedent

import netlib.http
from netlib.utils import parse_content_type

import re
from six.moves.urllib.parse import urlparse, quote, quote_plus

def curl_command(flow):
    data = "curl "

    for k, v in flow.request.headers.fields:
        data += "-H '%s:%s' " % (k, v)

    if flow.request.method != "GET":
        data += "-X %s " % flow.request.method

    full_url = flow.request.scheme + "://" + flow.request.host + flow.request.path
    data += "'%s'" % full_url

    if flow.request.content:
        data += " --data-binary '%s'" % flow.request.content

    return data


def python_code(flow):
    code = dedent("""
        import requests

        url = '{url}'
        {headers}{params}{data}
        response = requests.request(
            method='{method}',
            url=url,{args}
        )

        print(response.text)
    """).strip()

    components = map(lambda x: quote(x, safe=""), flow.request.path_components)
    url = flow.request.scheme + "://" + flow.request.host + "/" + "/".join(components)

    args = ""
    headers = ""
    if flow.request.headers:
        lines = ["    '%s': '%s',\n" % (k, v) for k, v in flow.request.headers.fields]
        headers += "\nheaders = {\n%s}\n" % "".join(lines)
        args += "\n    headers=headers,"

    params = ""
    if flow.request.query:
        lines = ["    %s: %s,\n" % (repr(k), repr(v)) for k, v in flow.request.query.to_dict().items()]
        params = "\nparams = {\n%s}\n" % "".join(lines)
        args += "\n    params=params,"

    data = ""
    if flow.request.body:
        json_obj = is_json(flow.request.headers, flow.request.body)
        if json_obj:
            # Without the separators field json.dumps() produces
            # trailing white spaces: https://bugs.python.org/issue16333
            data = json.dumps(json_obj, indent=4, separators=(',', ': '))
            data = "\njson = %s\n" % data
            args += "\n    json=json,"
        else:
            data = "\ndata = '''%s'''\n" % flow.request.body
            args += "\n    data=data,"

    code = code.format(
        url=url,
        headers=headers,
        params=params,
        data=data,
        method=flow.request.method,
        args=args,
    )

    return code


def raw_request(flow):
    data = netlib.http.http1.assemble_request(flow.request)
    return data


def is_json(headers, content):
    if headers:
        ct = parse_content_type(headers.get("content-type", ""))
        if ct and "%s/%s" % (ct[0], ct[1]) == "application/json":
            try:
                return json.loads(content)
            except ValueError:
                return False
    return False


def locust_code(flow):
    code = dedent("""
        from locust import HttpLocust, TaskSet, task

        class UserBehavior(TaskSet):
            def on_start(self):
                ''' on_start is called when a Locust start before any task is scheduled '''
                self.{name}()

            @task()
            def {name}(self):
                url = '{url}'
                {headers}{params}{data}
                self.response = self.client.request(
                    method='{method}',
                    url=url,{args}
                )

            ### Additional tasks can go here ###


        class WebsiteUser(HttpLocust):
            task_set = UserBehavior
            min_wait = 1000
            max_wait = 3000
""").strip()


    components = map(lambda x: quote(x, safe=""), flow.request.path_components)
    file_name = "_".join(components)
    name = re.sub('\W|^(?=\d)', '_', file_name)
    url = flow.request.scheme + "://" + flow.request.host + "/" + "/".join(components)
    if name == "" or name is None:
        new_name = "_".join([str(flow.request.host) , str(flow.request.timestamp_start)])
        name = re.sub('\W|^(?=\d)', '_', new_name)
    args = ""
    headers = ""
    if flow.request.headers:
        lines = ["            '%s': '%s',\n" % (k, v) for k, v in flow.request.headers.fields if k.lower() not in ["host", "cookie"]]
        headers += "\n        headers = {\n%s        }\n" % "".join(lines)
        args += "\n            headers=headers,"

    params = ""
    if flow.request.query:
        lines = ["            %s: %s,\n" % (repr(k), repr(v)) for k, v in flow.request.query.to_dict().items()]
        params = "\n        params = {\n%s        }\n" % "".join(lines)
        args += "\n            params=params,"

    data = ""
    if flow.request.body:
        data = "\n        data = '''%s'''\n" % flow.request.body
        args += "\n            data=data,"

    code = code.format(
        name=name,
        url=url,
        headers=headers,
        params=params,
        data=data,
        method=flow.request.method,
        args=args,
    )

    host = flow.request.scheme + "://" + flow.request.host
    code = code.replace(host, "' + self.locust.host + '")
    code = code.replace(quote_plus(host), "' + quote_plus(self.locust.host) + '")
    code = code.replace(quote(host), "' + quote(self.locust.host) + '")
    code = code.replace("'' + ", "")

    return code


def locust_task(flow):
    code = locust_code(flow)
    start_task = len(code.split('@task')[0]) - 4
    end_task = -19 - len(code.split('### Additional')[1])
    task_code = code[start_task:end_task]

    return task_code