Merge pull request #916 from dufferzafar/export-code

Export any flow request as Python code

libmproxy/console/common.py

@@ -8,6 +8,7 @@ from netlib.http import CONTENT_MISSING
 import netlib.utils
 
 from .. import utils
+from .. import flow_export
 from ..models import decoded
 from . import signals
 
@@ -281,56 +282,13 @@ def copy_flow_format_data(part, scope, flow):
 
 
 def export_prompt(k, flow):
-    if k == "c":
-        copy_as_curl_command(flow)
-    elif k == "p":
-        copy_as_python_code(flow)
-    elif k == "r":
-        copy_as_raw_request(flow)
-
-
-def copy_as_curl_command(flow):
-    data = "curl "
-
-    for k, v in flow.request.headers.fields:
-        data += "-H '%s:%s' " % (k, v)
-
-    if flow.request.method != "GET":
-        data += "-X %s " % flow.request.method
-
-    full_url = flow.request.scheme + "://" + flow.request.host + flow.request.path
-    data += "'%s'" % full_url
-
-    if flow.request.content:
-        data += " --data-binary '%s'" % flow.request.content
-
-    copy_to_clipboard_or_prompt(data)
-
-
-def copy_as_python_code(flow):
-    if flow.request.method != "GET":
-        signals.status_message.send(message="Currently, only GET methods are supported")
-        return
-
-    data = ("import requests\n"
-            "headers = {%s}\n"
-            "url = '%s'\n"
-            "resp = requests.get(url, headers=headers)")
-
-    headers = "\n"
-    for k, v in flow.request.headers.fields:
-        headers += "    '%s': '%s',\n" % (k, v)
-
-    full_url = flow.request.scheme + "://" + flow.request.host + flow.request.path
-
-    data = data % (headers, full_url)
-
-    copy_to_clipboard_or_prompt(data)
-
-
-def copy_as_raw_request(flow):
-    data = netlib.http.http1.assemble_request(flow.request)
-    copy_to_clipboard_or_prompt(data)
+    exporters = {
+        "c": flow_export.curl_command,
+        "p": flow_export.python_code,
+        "r": flow_export.raw_request,
+    }
+    if k in exporters:
+        copy_to_clipboard_or_prompt(exporters[k](flow))
 
 
 def copy_to_clipboard_or_prompt(data):

libmproxy/flow_export.py

@@ -0,0 +1,73 @@
+import urllib
+import netlib.http
+from textwrap import dedent
+
+
+def curl_command(flow):
+    data = "curl "
+
+    for k, v in flow.request.headers.fields:
+        data += "-H '%s:%s' " % (k, v)
+
+    if flow.request.method != "GET":
+        data += "-X %s " % flow.request.method
+
+    full_url = flow.request.scheme + "://" + flow.request.host + flow.request.path
+    data += "'%s'" % full_url
+
+    if flow.request.content:
+        data += " --data-binary '%s'" % flow.request.content
+
+    return data
+
+
+def python_code(flow):
+    code = dedent("""
+        import requests
+
+        url = '{url}'
+        {headers}{params}{data}
+        response = requests.request(
+            method='{method}',
+            url=url,{args}
+        )
+
+        print(response.text)
+    """).strip()
+
+    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:
+        lines = ["    '%s': '%s',\n" % (k, v) for k, v in flow.request.headers.fields]
+        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=url,
+        headers=headers,
+        params=params,
+        data=data,
+        method=flow.request.method,
+        args=args,
+    )
+
+    return code
+
+
+def raw_request(flow):
+    data = netlib.http.http1.assemble_request(flow.request)
+    return data

test/test_flow_export.py

@@ -0,0 +1,148 @@
+from textwrap import dedent
+
+import netlib.tutils
+from libmproxy import flow_export
+from . import tutils
+
+req_get = netlib.tutils.treq(
+    method='GET',
+    content='',
+)
+
+req_post = netlib.tutils.treq(
+    method='POST',
+    headers=None,
+)
+
+req_patch = netlib.tutils.treq(
+    method='PATCH',
+    path=b"/path?query=param",
+)
+
+
+class TestExportCurlCommand():
+
+    def test_get(self):
+        flow = tutils.tflow(req=req_get)
+        result = """curl -H 'header:qvalue' -H 'content-length:7' 'http://address/path'"""
+        assert flow_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'"""
+        assert flow_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'"""
+        assert flow_export.curl_command(flow) == result
+
+
+class TestExportPythonCode():
+
+    def test_get(self):
+        flow = tutils.tflow(req=req_get)
+        result = dedent("""
+            import requests
+
+            url = 'http://address/path'
+
+            headers = {
+                'header': 'qvalue',
+                'content-length': '7',
+            }
+
+            response = requests.request(
+                method='GET',
+                url=url,
+                headers=headers,
+            )
+
+            print(response.text)
+        """).strip()
+        assert flow_export.python_code(flow) == result
+
+    def test_post(self):
+        flow = tutils.tflow(req=req_post)
+        result = dedent("""
+            import requests
+
+            url = 'http://address/path'
+
+            data = '''content'''
+
+            response = requests.request(
+                method='POST',
+                url=url,
+                data=data,
+            )
+
+            print(response.text)
+        """).strip()
+        assert flow_export.python_code(flow) == result
+
+    def test_patch(self):
+        flow = tutils.tflow(req=req_patch)
+        result = dedent("""
+            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,
+            )
+
+            print(response.text)
+        """).strip()
+        assert flow_export.python_code(flow) == result
+
+
+def TestRawRequest():
+
+    def test_get(self):
+        flow = tutils.tflow(req=req_get)
+        result = dedent("""
+            GET /path HTTP/1.1\r
+            header: qvalue\r
+            content-length: 7\r
+            host: address:22\r
+            \r
+        """).strip(" ").lstrip()
+        assert flow_export.raw_request(flow) == result
+
+    def test_post(self):
+        flow = tutils.tflow(req=req_post)
+        result = dedent("""
+            POST /path HTTP/1.1\r
+            host: address:22\r
+            \r
+            content
+        """).strip()
+        assert flow_export.raw_request(flow) == result
+
+    def test_patch(self):
+        flow = tutils.tflow(req=req_patch)
+        result = dedent("""
+            PATCH /path?query=param HTTP/1.1\r
+            header: qvalue\r
+            content-length: 7\r
+            host: address:22\r
+            \r
+            content
+        """).strip()
+        assert flow_export.raw_request(flow) == result