diff options
Diffstat (limited to 'mitmproxy/flow_export.py')
| -rw-r--r-- | mitmproxy/flow_export.py | 65 |
1 files changed, 11 insertions, 54 deletions
diff --git a/mitmproxy/flow_export.py b/mitmproxy/flow_export.py index a9fd5cde..e2d46d93 100644 --- a/mitmproxy/flow_export.py +++ b/mitmproxy/flow_export.py @@ -103,10 +103,10 @@ def locust_code(flow): class UserBehavior(TaskSet): def on_start(self): ''' on_start is called when a Locust start before any task is scheduled ''' - self.flow() + self.{name}() @task() - def flow(self): + def {name}(self): url = '{url}' {headers}{params}{data} self.response = self.client.request( @@ -121,10 +121,12 @@ def locust_code(flow): task_set = UserBehavior min_wait = 1000 max_wait = 3000 +""").strip() - """).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) args = "" @@ -146,6 +148,7 @@ def locust_code(flow): args += "\n data=data," code = code.format( + name=name, url=url, headers=headers, params=params, @@ -163,55 +166,9 @@ def locust_code(flow): def locust_task(flow): - code = dedent(""" - @task() - def {name}(self): - url = '{url}' - {headers}{params}{data} - self.response = self.client.request( - method='{method}', - url=url,{args} - ) - """).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) - - 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" % (k, v) for k, v in flow.request.query] - params = "\n params = {\n%s }\n" % "".join(lines) - args += "\n params=params," + 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] - 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 = "\n".join(" " + i for i in code.splitlines()) - - return code + return task_code |