Merge pull request #1050 from zlorb/master

Flow export to locust.io load test tool [http://locust.io]

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):