Flow export to locust.io load test tool.

2 files changed, 308 insertions, 4 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
diff --git a/test/mitmproxy/test_flow_export.py b/test/mitmproxy/test_flow_export.py
index 62161d5d..a4264e10 100644
--- a/test/mitmproxy/test_flow_export.py
+++ b/test/mitmproxy/test_flow_export.py
@@ -142,7 +142,7 @@ class TestExportPythonCode():
         assert flow_export.python_code(flow) == result
 
 
-def TestRawRequest():
+class TestRawRequest():
 
     def test_get(self):
         flow = tutils.tflow(req=req_get)
@@ -159,9 +159,10 @@ def TestRawRequest():
         flow = tutils.tflow(req=req_post)
         result = dedent("""
             POST /path HTTP/1.1\r
+            content-type: application/json\r
             host: address:22\r
             \r
-            content
+            {"name": "example", "email": "example@example.com"}
         """).strip()
         assert flow_export.raw_request(flow) == result
 
@@ -176,3 +177,187 @@ def TestRawRequest():
             content
         """).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.flow()
+
+    @task()
+    def flow(self):
+        url = '' + self.locust.host +'/path'
+        
+        headers = {
+            'header': 'qvalue',
+            'content-length': '7',
+        }
+
+        self.response = self.client.request(
+            method='GET',
+            url=url,
+            headers=headers,
+        )
+
+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.flow()
+
+    @task()
+    def flow(self):
+        url = '' + self.locust.host +'/path'
+        
+        data = '''content'''
+
+        self.response = self.client.request(
+            method='POST',
+            url=url,
+            data=data,
+        )
+
+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.flow()
+
+    @task()
+    def flow(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,
+        )
+
+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()
+
+        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()
+
+        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()
+
+        assert flow_export.locust_task(flow) == result