Merge pull request #1656 from mhils/improve-export-2

Improve Flow Export

7 files changed, 87 insertions, 162 deletions

diff --git a/mitmproxy/export.py b/mitmproxy/export.py
index d9a88849..0a261509 100644
--- a/mitmproxy/export.py
+++ b/mitmproxy/export.py
@@ -1,9 +1,11 @@
+import io
 import json
+import pprint
 import re
 import textwrap
-import urllib
+from typing import Any
 
-import mitmproxy.net.http
+from mitmproxy import http
 
 
 def _native(s):
@@ -12,14 +14,14 @@ def _native(s):
     return s
 
 
-def dictstr(items, indent):
+def dictstr(items, indent: str) -> str:
     lines = []
     for k, v in items:
         lines.append(indent + "%s: %s,\n" % (repr(_native(k)), repr(_native(v))))
     return "{\n%s}\n" % "".join(lines)
 
 
-def curl_command(flow):
+def curl_command(flow: http.HTTPFlow) -> str:
     data = "curl "
 
     request = flow.request.copy()
@@ -31,8 +33,7 @@ def curl_command(flow):
     if request.method != "GET":
         data += "-X %s " % request.method
 
-    full_url = request.scheme + "://" + request.host + request.path
-    data += "'%s'" % full_url
+    data += "'%s'" % request.url
 
     if request.content:
         data += " --data-binary '%s'" % _native(request.content)
@@ -40,64 +41,54 @@ def curl_command(flow):
     return data
 
 
-def python_code(flow):
-    code = textwrap.dedent("""
-        import requests
-
-        url = '{url}'
-        {headers}{params}{data}
-        response = requests.request(
-            method='{method}',
-            url=url,{args}
-        )
+def python_arg(arg: str, val: Any) -> str:
+    if not val:
+        return ""
+    if arg:
+        arg += "="
+    arg_str = "{}{},\n".format(
+        arg,
+        pprint.pformat(val, 79 - len(arg))
+    )
+    return textwrap.indent(arg_str, " " * 4)
 
-        print(response.text)
-    """).strip()
 
-    components = [urllib.parse.quote(c, safe="") for c in flow.request.path_components]
-    url = flow.request.scheme + "://" + flow.request.host + "/" + "/".join(components)
+def python_code(flow: http.HTTPFlow):
+    code = io.StringIO()
 
-    args = ""
-    headers = ""
-    if flow.request.headers:
-        headers += "\nheaders = %s\n" % dictstr(flow.request.headers.fields, "    ")
-        args += "\n    headers=headers,"
+    def writearg(arg, val):
+        code.write(python_arg(arg, val))
 
-    params = ""
-    if flow.request.query:
-        params = "\nparams = %s\n" % dictstr(flow.request.query.collect(), "    ")
-        args += "\n    params=params,"
+    code.write("import requests\n")
+    code.write("\n")
+    if flow.request.method.lower() in ("get", "post", "put", "head", "delete", "patch"):
+        code.write("response = requests.{}(\n".format(flow.request.method.lower()))
+    else:
+        code.write("response = requests.request(\n")
+        writearg("", flow.request.method)
+    url_without_query = flow.request.url.split("?", 1)[0]
+    writearg("", url_without_query)
 
-    data = ""
-    if flow.request.body:
-        json_obj = is_json(flow.request.headers, flow.request.content)
-        if json_obj:
-            data = "\njson = %s\n" % dictstr(sorted(json_obj.items()), "    ")
-            args += "\n    json=json,"
-        else:
-            data = "\ndata = '''%s'''\n" % _native(flow.request.content)
-            args += "\n    data=data,"
+    writearg("params", list(flow.request.query.fields))
 
-    code = code.format(
-        url=url,
-        headers=headers,
-        params=params,
-        data=data,
-        method=flow.request.method,
-        args=args,
-    )
-    return code
+    headers = flow.request.headers.copy()
+    # requests adds those by default.
+    for x in ("host", "content-length"):
+        headers.pop(x, None)
+    writearg("headers", dict(headers))
+    try:
+        if "json" not in flow.request.headers.get("content-type", ""):
+            raise ValueError()
+        writearg("json", json.loads(flow.request.text))
+    except ValueError:
+        writearg("data", flow.request.content)
 
+    code.seek(code.tell() - 2)  # remove last comma
+    code.write("\n)\n")
+    code.write("\n")
+    code.write("print(response.text)")
 
-def is_json(headers: mitmproxy.net.http.Headers, content: bytes) -> bool:
-    if headers:
-        ct = mitmproxy.net.http.parse_content_type(headers.get("content-type", ""))
-        if ct and "%s/%s" % (ct[0], ct[1]) == "application/json":
-            try:
-                return json.loads(content.decode("utf8", "surrogateescape"))
-            except ValueError:
-                return False
-    return False
+    return code.getvalue()
 
 
 def locust_code(flow):
@@ -111,7 +102,7 @@ def locust_code(flow):
 
             @task()
             def {name}(self):
-                url = '{url}'
+                url = self.locust.host + '{path}'
                 {headers}{params}{data}
                 self.response = self.client.request(
                     method='{method}',
@@ -127,13 +118,12 @@ def locust_code(flow):
             max_wait = 3000
 """).strip()
 
-    components = [urllib.parse.quote(c, safe="") for c in flow.request.path_components]
-    name = re.sub('\W|^(?=\d)', '_', "_".join(components))
-    if name == "" or name is None:
+    name = re.sub('\W|^(?=\d)', '_', flow.request.path.strip("/").split("?", 1)[0])
+    if not name:
         new_name = "_".join([str(flow.request.host), str(flow.request.timestamp_start)])
         name = re.sub('\W|^(?=\d)', '_', new_name)
 
-    url = flow.request.scheme + "://" + flow.request.host + "/" + "/".join(components)
+    path_without_query = flow.request.path.split("?")[0]
 
     args = ""
     headers = ""
@@ -148,7 +138,11 @@ def locust_code(flow):
 
     params = ""
     if flow.request.query:
-        lines = ["            %s: %s,\n" % (repr(k), repr(v)) for k, v in flow.request.query.collect()]
+        lines = [
+            "            %s: %s,\n" % (repr(k), repr(v))
+            for k, v in
+            flow.request.query.collect()
+        ]
         params = "\n        params = {\n%s        }\n" % "".join(lines)
         args += "\n            params=params,"
 
@@ -159,7 +153,7 @@ def locust_code(flow):
 
     code = code.format(
         name=name,
-        url=url,
+        path=path_without_query,
         headers=headers,
         params=params,
         data=data,
@@ -167,12 +161,6 @@ def locust_code(flow):
         args=args,
     )
 
-    host = flow.request.scheme + "://" + flow.request.host
-    code = code.replace(host, "' + self.locust.host + '")
-    code = code.replace(urllib.parse.quote_plus(host), "' + quote_plus(self.locust.host) + '")
-    code = code.replace(urllib.parse.quote(host), "' + quote(self.locust.host) + '")
-    code = code.replace("'' + ", "")
-
     return code
 
 
diff --git a/mitmproxy/http.py b/mitmproxy/http.py
index 99e126fe..1905791d 100644
--- a/mitmproxy/http.py
+++ b/mitmproxy/http.py
@@ -1,9 +1,11 @@
 import cgi
 
 from mitmproxy import flow
+
 from mitmproxy.net import http
 from mitmproxy import version
 from mitmproxy.net import tcp
+from mitmproxy import connections  # noqa
 
 
 class HTTPRequest(http.Request):
@@ -155,22 +157,22 @@ class HTTPFlow(flow.Flow):
     def __init__(self, client_conn, server_conn, live=None):
         super().__init__("http", client_conn, server_conn, live)
 
-        self.request = None
+        self.request = None  # type: HTTPRequest
         """ :py:class:`HTTPRequest` object """
-        self.response = None
+        self.response = None  # type: HTTPResponse
         """ :py:class:`HTTPResponse` object """
-        self.error = None
+        self.error = None  # type: flow.Error
         """ :py:class:`Error` object
 
         Note that it's possible for a Flow to have both a response and an error
         object. This might happen, for instance, when a response was received
         from the server, but there was an error sending it back to the client.
         """
-        self.server_conn = server_conn
+        self.server_conn = server_conn  # type: connections.ServerConnection
         """ :py:class:`ServerConnection` object """
-        self.client_conn = client_conn
+        self.client_conn = client_conn  # type: connections.ClientConnection
         """:py:class:`ClientConnection` object """
-        self.intercepted = False
+        self.intercepted = False  # type: bool
         """ Is this flow currently being intercepted? """
 
     _stateobject_attributes = flow.Flow._stateobject_attributes.copy()
diff --git a/test/mitmproxy/data/test_flow_export/python_get.py b/test/mitmproxy/data/test_flow_export/python_get.py
index af8f7c81..e9ed072a 100644
--- a/test/mitmproxy/data/test_flow_export/python_get.py
+++ b/test/mitmproxy/data/test_flow_export/python_get.py
@@ -1,22 +1,9 @@
 import requests
 
-url = 'http://address/path'
-
-headers = {
-    'header': 'qvalue',
-    'content-length': '7',
-}
-
-params = {
-    'a': ['foo', 'bar'],
-    'b': 'baz',
-}
-
-response = requests.request(
-    method='GET',
-    url=url,
-    headers=headers,
-    params=params,
+response = requests.get(
+    'http://address:22/path',
+    params=[('a', 'foo'), ('a', 'bar'), ('b', 'baz')],
+    headers={'header': 'qvalue'}
 )
 
-print(response.text)
+print(response.text)
\ No newline at end of file
diff --git a/test/mitmproxy/data/test_flow_export/python_patch.py b/test/mitmproxy/data/test_flow_export/python_patch.py
index 159e802f..d83a57b9 100644
--- a/test/mitmproxy/data/test_flow_export/python_patch.py
+++ b/test/mitmproxy/data/test_flow_export/python_patch.py
@@ -1,24 +1,10 @@
 import requests
 
-url = 'http://address/path'
-
-headers = {
-    'header': 'qvalue',
-    'content-length': '7',
-}
-
-params = {
-    'query': 'param',
-}
-
-data = '''content'''
-
-response = requests.request(
-    method='PATCH',
-    url=url,
-    headers=headers,
-    params=params,
-    data=data,
+response = requests.patch(
+    'http://address:22/path',
+    params=[('query', 'param')],
+    headers={'header': 'qvalue'},
+    data=b'content'
 )
 
-print(response.text)
+print(response.text)
\ No newline at end of file
diff --git a/test/mitmproxy/data/test_flow_export/python_post.py b/test/mitmproxy/data/test_flow_export/python_post.py
index b13f6441..6254adfb 100644
--- a/test/mitmproxy/data/test_flow_export/python_post.py
+++ b/test/mitmproxy/data/test_flow_export/python_post.py
@@ -1,13 +1,8 @@
 import requests
 
-url = 'http://address/path'
-
-data = '''content'''
-
-response = requests.request(
-    method='POST',
-    url=url,
-    data=data,
+response = requests.post(
+    'http://address:22/path',
+    data=b'content'
 )
 
 print(response.text)
diff --git a/test/mitmproxy/data/test_flow_export/python_post_json.py b/test/mitmproxy/data/test_flow_export/python_post_json.py
index 5ef110f3..d6ae6357 100644
--- a/test/mitmproxy/data/test_flow_export/python_post_json.py
+++ b/test/mitmproxy/data/test_flow_export/python_post_json.py
@@ -1,23 +1,9 @@
 import requests
 
-url = 'http://address/path'
-
-headers = {
-    'content-type': 'application/json',
-}
-
-
-json = {
-    'email': 'example@example.com',
-    'name': 'example',
-}
-
-
-response = requests.request(
-    method='POST',
-    url=url,
-    headers=headers,
-    json=json,
+response = requests.post(
+    'http://address:22/path',
+    headers={'content-type': 'application/json'},
+    json={'email': 'example@example.com', 'name': 'example'}
 )
 
-print(response.text)
+print(response.text)
\ No newline at end of file
diff --git a/test/mitmproxy/test_flow_export.py b/test/mitmproxy/test_flow_export.py
index df0ccb77..86234616 100644
--- a/test/mitmproxy/test_flow_export.py
+++ b/test/mitmproxy/test_flow_export.py
@@ -34,17 +34,17 @@ def req_patch():
 class TestExportCurlCommand:
     def test_get(self):
         flow = tutils.tflow(req=req_get())
-        result = """curl -H 'header:qvalue' -H 'content-length:7' 'http://address/path?a=foo&a=bar&b=baz'"""
+        result = """curl -H 'header:qvalue' -H 'content-length:7' 'http://address:22/path?a=foo&a=bar&b=baz'"""
         assert export.curl_command(flow) == result
 
     def test_post(self):
         flow = tutils.tflow(req=req_post())
-        result = """curl -X POST 'http://address/path' --data-binary 'content'"""
+        result = """curl -X POST 'http://address:22/path' --data-binary 'content'"""
         assert export.curl_command(flow) == result
 
     def test_patch(self):
         flow = tutils.tflow(req=req_patch())
-        result = """curl -H 'header:qvalue' -H 'content-length:7' -X PATCH 'http://address/path?query=param' --data-binary 'content'"""
+        result = """curl -H 'header:qvalue' -H 'content-length:7' -X PATCH 'http://address:22/path?query=param' --data-binary 'content'"""
         assert export.curl_command(flow) == result
 
 
@@ -100,25 +100,6 @@ class TestExportLocustTask:
         python_equals("data/test_flow_export/locust_task_patch.py", export.locust_task(flow))
 
 
-class TestIsJson:
-    def test_empty(self):
-        assert export.is_json(None, None) is False
-
-    def test_json_type(self):
-        headers = Headers(content_type="application/json")
-        assert export.is_json(headers, b"foobar") is False
-
-    def test_valid(self):
-        headers = Headers(content_type="application/foobar")
-        j = export.is_json(headers, b'{"name": "example", "email": "example@example.com"}')
-        assert j is False
-
-    def test_valid2(self):
-        headers = Headers(content_type="application/json")
-        j = export.is_json(headers, b'{"name": "example", "email": "example@example.com"}')
-        assert isinstance(j, dict)
-
-
 class TestURL:
     def test_url(self):
         flow = tutils.tflow()