1 files changed, 121 insertions, 2 deletions
diff --git a/mitmproxy/flow_export.py b/mitmproxy/flow_export.py
index 6333de57..a1e07953 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,121 @@ 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.flow()
+
+            @task()
+            def flow(self):
+                url = '{url}'
+                {headers}{params}{data}
+                self.response = self.client.request(
+                    method='{method}',
+                    url=url,{args}
+                )
+
+        class WebsiteUser(HttpLocust):
+            task_set = UserBehavior
+            min_wait=1000
+            max_wait=3000
+
+    """).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 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(
+        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) +'")
+
+    return code
+
+
+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,"
+
+    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