Merge branch 'master' of https://github.com/mitmproxy/mitmproxy

10 files changed, 110 insertions, 30 deletions

diff --git a/mitmproxy/console/common.py b/mitmproxy/console/common.py
index c29ffddc..141735ef 100644
--- a/mitmproxy/console/common.py
+++ b/mitmproxy/console/common.py
@@ -4,7 +4,6 @@ import urwid
 import urwid.util
 import os
 
-from netlib.http import CONTENT_MISSING
 import netlib.utils
 
 from .. import utils
@@ -256,7 +255,7 @@ def copy_flow_format_data(part, scope, flow):
     else:
         data = ""
         if scope in ("q", "a"):
-            if flow.request.content is None or flow.request.content == CONTENT_MISSING:
+            if flow.request.content is None:
                 return None, "Request content is missing"
             with decoded(flow.request):
                 if part == "h":
@@ -269,7 +268,7 @@ def copy_flow_format_data(part, scope, flow):
             # Add padding between request and response
             data += "\r\n" * 2
         if scope in ("s", "a") and flow.response:
-            if flow.response.content is None or flow.response.content == CONTENT_MISSING:
+            if flow.response.content is None:
                 return None, "Response content is missing"
             with decoded(flow.response):
                 if part == "h":
@@ -286,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))
@@ -418,7 +419,7 @@ def format_flow(f, focus, extended=False, hostheader=False, marked=False):
     if f.response:
         if f.response.content:
             contentdesc = netlib.utils.pretty_size(len(f.response.content))
-        elif f.response.content == CONTENT_MISSING:
+        elif f.response.content is None:
             contentdesc = "[content missing]"
         else:
             contentdesc = "[no content]"
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 f74ab140..b761a924 100644
--- a/mitmproxy/console/flowview.py
+++ b/mitmproxy/console/flowview.py
@@ -7,7 +7,7 @@ import math
 import urwid
 
 from netlib import odict
-from netlib.http import CONTENT_MISSING, Headers
+from netlib.http import Headers
 from . import common, grideditor, signals, searchable, tabs
 from . import flowdetailview
 from .. import utils, controller, contentviews
@@ -169,7 +169,7 @@ class FlowView(tabs.Tabs):
             self.show()
 
     def content_view(self, viewmode, message):
-        if message.content == CONTENT_MISSING:
+        if message.content is None:
             msg, body = "", [urwid.Text([("error", "[content missing]")])]
             return msg, body
         else:
@@ -510,7 +510,7 @@ class FlowView(tabs.Tabs):
 
     def delete_body(self, t):
         if t == "m":
-            val = CONTENT_MISSING
+            val = None
         else:
             val = None
         if self.tab_offset == TAB_REQ:
@@ -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/dump.py b/mitmproxy/dump.py
index 7b4609b4..631e4d2e 100644
--- a/mitmproxy/dump.py
+++ b/mitmproxy/dump.py
@@ -5,7 +5,6 @@ import click
 import itertools
 
 from netlib import tcp
-from netlib.http import CONTENT_MISSING
 import netlib.utils
 from . import flow, filt, contentviews
 from .exceptions import ContentViewException
@@ -180,7 +179,7 @@ class DumpMaster(flow.FlowMaster):
             )
             self.echo(headers, indent=4)
         if self.o.flow_detail >= 3:
-            if message.content == CONTENT_MISSING:
+            if message.content is None:
                 self.echo("(content missing)", indent=4)
             elif message.content:
                 self.echo("")
@@ -283,7 +282,7 @@ class DumpMaster(flow.FlowMaster):
         code = click.style(str(code), fg=code_color, bold=True, blink=(code == 418))
         reason = click.style(flow.response.reason, fg=code_color, bold=True)
 
-        if flow.response.content == CONTENT_MISSING:
+        if flow.response.content is None:
             size = "(content missing)"
         else:
             size = netlib.utils.pretty_size(len(flow.response.content))
diff --git a/mitmproxy/flow.py b/mitmproxy/flow.py
index d656bc4d..fbf102b5 100644
--- a/mitmproxy/flow.py
+++ b/mitmproxy/flow.py
@@ -16,7 +16,7 @@ from six.moves import urllib
 
 from netlib import wsgi
 from netlib.exceptions import HttpException
-from netlib.http import CONTENT_MISSING, Headers, http1
+from netlib.http import Headers, http1
 from . import controller, tnetstring, filt, script, version, flow_format_compat
 from .onboarding import app
 from .proxy.config import HostMatcher
@@ -942,7 +942,7 @@ class FlowMaster(controller.Master):
             return "Can't replay live request."
         if f.intercepted:
             return "Can't replay while intercepting..."
-        if f.request.content == CONTENT_MISSING:
+        if f.request.content is None:
             return "Can't replay request with missing content..."
         if f.request:
             f.backup()
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/mitmproxy/models/http.py b/mitmproxy/models/http.py
index a80e11b0..40460182 100644
--- a/mitmproxy/models/http.py
+++ b/mitmproxy/models/http.py
@@ -84,9 +84,8 @@ class HTTPRequest(MessageMixin, Request):
 
         headers: Headers object
 
-        content: Content of the request, None, or CONTENT_MISSING if there
-        is content associated, but not present. CONTENT_MISSING evaluates
-        to False to make checking for the presence of content natural.
+        content: Content of the request, the value is None if there is content
+        associated, but not present.
 
         form_in: The request form which mitmproxy has received. The following
         values are possible:
@@ -226,9 +225,8 @@ class HTTPResponse(MessageMixin, Response):
 
         headers: Headers object
 
-        content: Content of the request, None, or CONTENT_MISSING if there
-        is content associated, but not present. CONTENT_MISSING evaluates
-        to False to make checking for the presence of content natural.
+        content: Content of the response, the value is None if there is content
+        associated, but not present.
 
         timestamp_start: Timestamp indicating when request transmission started
 
diff --git a/mitmproxy/protocol/http.py b/mitmproxy/protocol/http.py
index a2745eac..56c5f9ea 100644
--- a/mitmproxy/protocol/http.py
+++ b/mitmproxy/protocol/http.py
@@ -6,7 +6,7 @@ import six
 
 from netlib import tcp
 from netlib.exceptions import HttpException, HttpReadDisconnect, NetlibException
-from netlib.http import Headers, CONTENT_MISSING
+from netlib.http import Headers
 
 from h2.exceptions import H2Error
 
@@ -50,8 +50,8 @@ class _HttpTransmissionLayer(Layer):
         yield "this is a generator"  # pragma: no cover
 
     def send_response(self, response):
-        if response.content == CONTENT_MISSING:
-            raise HttpException("Cannot assemble flow with CONTENT_MISSING")
+        if response.content is None:
+            raise HttpException("Cannot assemble flow with missing content")
         self.send_response_headers(response)
         self.send_response_body(response, [response.content])
 
@@ -231,7 +231,8 @@ class HttpLayer(Layer):
                     six.reraise(ProtocolException, ProtocolException(
                         "Error in HTTP connection: %s" % repr(e)), sys.exc_info()[2])
             finally:
-                flow.live = False
+                if flow:
+                    flow.live = False
 
     def get_request_from_client(self):
         request = self.read_request()
@@ -318,7 +319,7 @@ class HttpLayer(Layer):
             raise Kill()
 
         if flow.response.stream:
-            flow.response.data.content = CONTENT_MISSING
+            flow.response.data.content = None
         else:
             flow.response.data.content = b"".join(self.read_response_body(
                 flow.request,
diff --git a/mitmproxy/script/script.py b/mitmproxy/script/script.py
index f142daca..5a8334c4 100644
--- a/mitmproxy/script/script.py
+++ b/mitmproxy/script/script.py
@@ -47,7 +47,7 @@ class Script(object):
         if os.name == "nt":  # pragma: no cover
             backslashes = shlex.split(command, posix=False)[0].count("\\")
             command = command.replace("\\", "\\\\", backslashes)
-        args = shlex.split(command)  # pragma: nocover
+        args = shlex.split(command)  # pragma: no cover
         args[0] = os.path.expanduser(args[0])
         if not os.path.exists(args[0]):
             raise ScriptException(
diff --git a/mitmproxy/web/app.py b/mitmproxy/web/app.py
index 2cac2ab9..cf3c8bdd 100644
--- a/mitmproxy/web/app.py
+++ b/mitmproxy/web/app.py
@@ -8,7 +8,6 @@ import logging
 import json
 import base64
 
-from netlib.http import CONTENT_MISSING
 from .. import version, filt
 
 
@@ -26,10 +25,8 @@ def _strip_content(flow_state):
                 continue
             if message["content"]:
                 message["contentLength"] = len(message["content"])
-            elif message["content"] == CONTENT_MISSING:
-                message["contentLength"] = None
             else:
-                message["contentLength"] = 0
+                message["contentLength"] = None
             del message["content"]
 
     if "backup" in flow_state: