Export any flow as Python code, not just GET

The generic code idea was taken from Postman's code export:
https://www.getpostman.com/

1 files changed, 19 insertions, 5 deletions

diff --git a/libmproxy/console/common.py b/libmproxy/console/common.py
index 93c4df72..f96d7dd6 100644
--- a/libmproxy/console/common.py
+++ b/libmproxy/console/common.py
@@ -3,6 +3,7 @@ from __future__ import absolute_import
 import urwid
 import urwid.util
 import os
+import urllib
 
 from netlib.http import CONTENT_MISSING
 import netlib.utils
@@ -305,14 +306,11 @@ def copy_as_curl_command(flow):
 
 
 def copy_as_python_code(flow):
-    if flow.request.method != "GET":
-        signals.status_message.send(message="Currently, only GET methods are supported")
-        return
 
     code = """import requests
 
 url = '{url}'
-{headers}
+{headers}{params}{data}
 response = requests.request(
     method='{method}',
     url=url,{args}
@@ -320,6 +318,9 @@ response = requests.request(
 
 print(response.text)"""
 
+    components = map(lambda x: urllib.quote(x, safe=""), flow.request.path_components)
+    url = flow.request.scheme + "://" + flow.request.host + "/" + "/".join(components)
+
     args = ""
     headers = ""
     if flow.request.headers:
@@ -327,9 +328,22 @@ print(response.text)"""
         headers += "\nheaders = {\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="\nparams = {\n%s}\n" % "".join(lines)
+        args += "\n    params=params,"
+
+    data = ""
+    if flow.request.body:
+        data="\ndata = '''%s'''\n" % flow.request.body
+        args += "\n    data=data,"
+
     code = code.format(
-        url=flow.request.pretty_url,
+        url=url,
         headers=headers,
+        params=params,
+        data=data,
         method=flow.request.method,
         args=args,
     )