aboutsummaryrefslogtreecommitdiffstats
diff options
context:
space:
mode:
Diffstat
-rw-r--r--mitmproxy/console/common.py2
-rw-r--r--mitmproxy/console/flowlist.py2
-rw-r--r--mitmproxy/console/flowview.py2
-rw-r--r--mitmproxy/flow_export.py84
-rw-r--r--test/mitmproxy/test_flow_export.py193
5 files changed, 281 insertions, 2 deletions
diff --git a/mitmproxy/console/common.py b/mitmproxy/console/common.py
index 746a67f1..141735ef 100644
--- a/mitmproxy/console/common.py
+++ b/mitmproxy/console/common.py
@@ -285,6 +285,8 @@ def export_prompt(k, flow):
"c": flow_export.curl_command,
"p": flow_export.python_code,
"r": flow_export.raw_request,
+ "l": flow_export.locust_code,
+ "t": flow_export.locust_task,
}
if k in exporters:
copy_to_clipboard_or_prompt(exporters[k](flow))
diff --git a/mitmproxy/console/flowlist.py b/mitmproxy/console/flowlist.py
index c2201055..78b30231 100644
--- a/mitmproxy/console/flowlist.py
+++ b/mitmproxy/console/flowlist.py
@@ -265,6 +265,8 @@ class ConnectionItem(urwid.WidgetWrap):
("as curl command", "c"),
("as python code", "p"),
("as raw request", "r"),
+ ("as locust code", "l"),
+ ("as locust task", "t"),
),
callback = common.export_prompt,
args = (self.flow,)
diff --git a/mitmproxy/console/flowview.py b/mitmproxy/console/flowview.py
index 3b86c920..b761a924 100644
--- a/mitmproxy/console/flowview.py
+++ b/mitmproxy/console/flowview.py
@@ -585,6 +585,8 @@ class FlowView(tabs.Tabs):
("as curl command", "c"),
("as python code", "p"),
("as raw request", "r"),
+ ("as locust code", "l"),
+ ("as locust task", "t"),
),
callback = common.export_prompt,
args = (self.flow,)
diff --git a/mitmproxy/flow_export.py b/mitmproxy/flow_export.py
index 6333de57..e2ba7161 100644
--- a/mitmproxy/flow_export.py
+++ b/mitmproxy/flow_export.py
@@ -1,10 +1,11 @@
import json
-import urllib
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 "
@@ -38,7 +39,7 @@ def python_code(flow):
print(response.text)
""").strip()
- components = map(lambda x: urllib.quote(x, safe=""), flow.request.path_components)
+ components = map(lambda x: quote(x, safe=""), flow.request.path_components)
url = flow.request.scheme + "://" + flow.request.host + "/" + "/".join(components)
args = ""
@@ -93,3 +94,82 @@ def is_json(headers, 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)
+
+ 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,"
+
+ 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
diff --git a/test/mitmproxy/test_flow_export.py b/test/mitmproxy/test_flow_export.py
index d98759bb..75a8090f 100644
--- a/test/mitmproxy/test_flow_export.py
+++ b/test/mitmproxy/test_flow_export.py
@@ -179,6 +179,199 @@ class TestRawRequest():
""").strip()
assert flow_export.raw_request(flow) == result
+class TestExportLocustCode():
+
+ def test_get(self):
+ flow = tutils.tflow(req=req_get)
+ result = """
+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.path()
+
+ @task()
+ def path(self):
+ url = self.locust.host + '/path'
+
+ headers = {
+ 'header': 'qvalue',
+ 'content-length': '7',
+ }
+
+ self.response = self.client.request(
+ method='GET',
+ url=url,
+ headers=headers,
+ )
+
+ ### Additional tasks can go here ###
+
+
+class WebsiteUser(HttpLocust):
+ task_set = UserBehavior
+ min_wait = 1000
+ max_wait = 3000
+ """.strip()
+
+ assert flow_export.locust_code(flow) == result
+
+ def test_post(self):
+ req_post.content = '''content'''
+ req_post.headers = ''
+ flow = tutils.tflow(req=req_post)
+ result = """
+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.path()
+
+ @task()
+ def path(self):
+ url = self.locust.host + '/path'
+
+ data = '''content'''
+
+ self.response = self.client.request(
+ method='POST',
+ url=url,
+ data=data,
+ )
+
+ ### Additional tasks can go here ###
+
+
+class WebsiteUser(HttpLocust):
+ task_set = UserBehavior
+ min_wait = 1000
+ max_wait = 3000
+
+ """.strip()
+
+ assert flow_export.locust_code(flow) == result
+
+
+ def test_patch(self):
+ flow = tutils.tflow(req=req_patch)
+ result = """
+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.path()
+
+ @task()
+ def path(self):
+ url = self.locust.host + '/path'
+
+ headers = {
+ 'header': 'qvalue',
+ 'content-length': '7',
+ }
+
+ params = {
+ 'query': 'param',
+ }
+
+ data = '''content'''
+
+ self.response = self.client.request(
+ method='PATCH',
+ url=url,
+ headers=headers,
+ params=params,
+ data=data,
+ )
+
+ ### Additional tasks can go here ###
+
+
+class WebsiteUser(HttpLocust):
+ task_set = UserBehavior
+ min_wait = 1000
+ max_wait = 3000
+
+ """.strip()
+
+ assert flow_export.locust_code(flow) == result
+
+
+class TestExportLocustTask():
+
+ def test_get(self):
+ flow = tutils.tflow(req=req_get)
+ result = ' ' + """
+ @task()
+ def path(self):
+ url = self.locust.host + '/path'
+
+ headers = {
+ 'header': 'qvalue',
+ 'content-length': '7',
+ }
+
+ self.response = self.client.request(
+ method='GET',
+ url=url,
+ headers=headers,
+ )
+ """.strip() + '\n'
+
+ assert flow_export.locust_task(flow) == result
+
+ def test_post(self):
+ flow = tutils.tflow(req=req_post)
+ result = ' ' + """
+ @task()
+ def path(self):
+ url = self.locust.host + '/path'
+
+ data = '''content'''
+
+ self.response = self.client.request(
+ method='POST',
+ url=url,
+ data=data,
+ )
+ """.strip() + '\n'
+
+ assert flow_export.locust_task(flow) == result
+
+
+ def test_patch(self):
+ flow = tutils.tflow(req=req_patch)
+ result = ' ' + """
+ @task()
+ def path(self):
+ url = self.locust.host + '/path'
+
+ headers = {
+ 'header': 'qvalue',
+ 'content-length': '7',
+ }
+
+ params = {
+ 'query': 'param',
+ }
+
+ data = '''content'''
+
+ self.response = self.client.request(
+ method='PATCH',
+ url=url,
+ headers=headers,
+ params=params,
+ data=data,
+ )
+ """.strip() + '\n'
+
+ assert flow_export.locust_task(flow) == result
+
+
class TestIsJson():
def test_empty(self):
generated by cgit v1.2.3 (git 2.25.1) at 2025-04-27 06:04:04 +0000