Merge pull request #969 from dufferzafar/pretty-export

Indent JSON body while exporting it as code

2 files changed, 56 insertions, 3 deletions

diff --git a/mitmproxy/flow_export.py b/mitmproxy/flow_export.py
index 52145516..6333de57 100644
--- a/mitmproxy/flow_export.py
+++ b/mitmproxy/flow_export.py
@@ -1,7 +1,10 @@
+import json
 import urllib
-import netlib.http
 from textwrap import dedent
 
+import netlib.http
+from netlib.utils import parse_content_type
+
 
 def curl_command(flow):
     data = "curl "
@@ -53,8 +56,16 @@ def python_code(flow):
 
     data = ""
     if flow.request.body:
-        data = "\ndata = '''%s'''\n" % flow.request.body
-        args += "\n    data=data,"
+        json_obj = is_json(flow.request.headers, flow.request.body)
+        if json_obj:
+            # Without the separators field json.dumps() produces
+            # trailing white spaces: https://bugs.python.org/issue16333
+            data = json.dumps(json_obj, indent=4, separators=(',', ': '))
+            data = "\njson = %s\n" % data
+            args += "\n    json=json,"
+        else:
+            data = "\ndata = '''%s'''\n" % flow.request.body
+            args += "\n    data=data,"
 
     code = code.format(
         url=url,
@@ -71,3 +82,14 @@ def python_code(flow):
 def raw_request(flow):
     data = netlib.http.http1.assemble_request(flow.request)
     return data
+
+
+def is_json(headers, content):
+    if headers:
+        ct = parse_content_type(headers.get("content-type", ""))
+        if ct and "%s/%s" % (ct[0], ct[1]) == "application/json":
+            try:
+                return json.loads(content)
+            except ValueError:
+                return False
+    return False
diff --git a/test/mitmproxy/test_flow_export.py b/test/mitmproxy/test_flow_export.py
index 2dce3fd6..3dc07427 100644
--- a/test/mitmproxy/test_flow_export.py
+++ b/test/mitmproxy/test_flow_export.py
@@ -1,6 +1,8 @@
+import json
 from textwrap import dedent
 
 import netlib.tutils
+from netlib.http import Headers
 from mitmproxy import flow_export
 from . import tutils
 
@@ -81,6 +83,35 @@ class TestExportPythonCode():
         """).strip()
         assert flow_export.python_code(flow) == result
 
+    def test_post_json(self):
+        req_post.content = '{"name": "example", "email": "example@example.com"}'
+        req_post.headers = Headers(content_type="application/json")
+        flow = tutils.tflow(req=req_post)
+        result = dedent("""
+            import requests
+
+            url = 'http://address/path'
+
+            headers = {
+                'content-type': 'application/json',
+            }
+
+            json = {
+                "name": "example",
+                "email": "example@example.com"
+            }
+
+            response = requests.request(
+                method='POST',
+                url=url,
+                headers=headers,
+                json=json,
+            )
+
+            print(response.text)
+        """).strip()
+        assert flow_export.python_code(flow) == result
+
     def test_patch(self):
         flow = tutils.tflow(req=req_patch)
         result = dedent("""