Merge pull request #2161 from mhils/make-request

Add Request.make

5 files changed, 84 insertions, 24 deletions

diff --git a/mitmproxy/master.py b/mitmproxy/master.py
index 79747a97..57d47394 100644
--- a/mitmproxy/master.py
+++ b/mitmproxy/master.py
@@ -15,7 +15,6 @@ from mitmproxy import log
 from mitmproxy import io
 from mitmproxy.proxy.protocol import http_replay
 from mitmproxy.types import basethread
-import mitmproxy.net.http
 
 from . import ctx as mitmproxy_ctx
 
@@ -122,27 +121,18 @@ class Master:
         self.should_exit.set()
         self.addons.done()
 
-    def create_request(self, method, scheme, host, port, path):
+    def create_request(self, method, url):
         """
-            this method creates a new artificial and minimalist request also adds it to flowlist
+        Create a new artificial and minimalist request also adds it to flowlist.
+
+        Raises:
+            ValueError, if the url is malformed.
         """
+        req = http.HTTPRequest.make(method, url)
         c = connections.ClientConnection.make_dummy(("", 0))
-        s = connections.ServerConnection.make_dummy((host, port))
+        s = connections.ServerConnection.make_dummy((req.host, req.port))
 
         f = http.HTTPFlow(c, s)
-        headers = mitmproxy.net.http.Headers()
-
-        req = http.HTTPRequest(
-            "absolute",
-            method,
-            scheme,
-            host,
-            port,
-            path,
-            b"HTTP/1.1",
-            headers,
-            b""
-        )
         f.request = req
         self.load_flow(f)
         return f
diff --git a/mitmproxy/net/http/request.py b/mitmproxy/net/http/request.py
index b961e1e4..90a1f924 100644
--- a/mitmproxy/net/http/request.py
+++ b/mitmproxy/net/http/request.py
@@ -1,6 +1,6 @@
 import re
 import urllib
-from typing import Optional
+from typing import Optional, AnyStr, Dict, Iterable, Tuple, Union
 
 from mitmproxy.types import multidict
 from mitmproxy.utils import strutils
@@ -77,6 +77,53 @@ class Request(message.Message):
             self.method, hostport, path
         )
 
+    @classmethod
+    def make(
+            cls,
+            method: str,
+            url: str,
+            content: AnyStr = b"",
+            headers: Union[Dict[AnyStr, AnyStr], Iterable[Tuple[bytes, bytes]]] = ()
+    ):
+        """
+        Simplified API for creating request objects.
+        """
+        req = cls(
+            "absolute",
+            method,
+            "",
+            "",
+            "",
+            "",
+            "HTTP/1.1",
+            (),
+            b""
+        )
+
+        req.url = url
+
+        # Headers can be list or dict, we differentiate here.
+        if isinstance(headers, dict):
+            req.headers = nheaders.Headers(**headers)
+        elif isinstance(headers, Iterable):
+            req.headers = nheaders.Headers(headers)
+        else:
+            raise TypeError("Expected headers to be an iterable or dict, but is {}.".format(
+                type(headers).__name__
+            ))
+
+        # Assign this manually to update the content-length header.
+        if isinstance(content, bytes):
+            req.content = content
+        elif isinstance(content, str):
+            req.text = content
+        else:
+            raise TypeError("Expected content to be str or bytes, but is {}.".format(
+                type(content).__name__
+            ))
+
+        return req
+
     def replace(self, pattern, repl, flags=0, count=0):
         """
             Replaces a regular expression pattern with repl in the headers, the
diff --git a/mitmproxy/tools/console/flowlist.py b/mitmproxy/tools/console/flowlist.py
index d2e28d35..5fe86975 100644
--- a/mitmproxy/tools/console/flowlist.py
+++ b/mitmproxy/tools/console/flowlist.py
@@ -1,6 +1,5 @@
 import urwid
 
-import mitmproxy.net.http.url
 from mitmproxy import exceptions
 from mitmproxy.tools.console import common
 from mitmproxy.tools.console import signals
@@ -339,12 +338,10 @@ class FlowListBox(urwid.ListBox):
 
     def new_request(self, url, method):
         try:
-            parts = mitmproxy.net.http.url.parse(str(url))
+            f = self.master.create_request(method, url)
         except ValueError as e:
             signals.status_message.send(message = "Invalid URL: " + str(e))
             return
-        scheme, host, port, path = parts
-        f = self.master.create_request(method, scheme, host, port, path)
         self.master.view.focus.flow = f
 
     def keypress(self, size, key):
diff --git a/test/mitmproxy/net/http/test_request.py b/test/mitmproxy/net/http/test_request.py
index 90ec31fe..ce49002c 100644
--- a/test/mitmproxy/net/http/test_request.py
+++ b/test/mitmproxy/net/http/test_request.py
@@ -1,7 +1,7 @@
 from unittest import mock
 import pytest
 
-from mitmproxy.net.http import Headers
+from mitmproxy.net.http import Headers, Request
 from mitmproxy.test.tutils import treq
 from .test_message import _test_decoded_attr, _test_passthrough_attr
 
@@ -35,6 +35,32 @@ class TestRequestCore:
         request.host = None
         assert repr(request) == "Request(GET /path)"
 
+    def test_make(self):
+        r = Request.make("GET", "https://example.com/")
+        assert r.method == "GET"
+        assert r.scheme == "https"
+        assert r.host == "example.com"
+        assert r.port == 443
+        assert r.path == "/"
+
+        r = Request.make("GET", "https://example.com/", "content", {"Foo": "bar"})
+        assert r.content == b"content"
+        assert r.headers["content-length"] == "7"
+        assert r.headers["Foo"] == "bar"
+
+        Request.make("GET", "https://example.com/", content=b"content")
+        with pytest.raises(TypeError):
+            Request.make("GET", "https://example.com/", content=42)
+
+        r = Request.make("GET", "https://example.com/", headers=[(b"foo", b"bar")])
+        assert r.headers["foo"] == "bar"
+
+        r = Request.make("GET", "https://example.com/", headers=({"foo": "baz"}))
+        assert r.headers["foo"] == "baz"
+
+        with pytest.raises(TypeError):
+            Request.make("GET", "https://example.com/", headers=42)
+
     def test_replace(self):
         r = treq()
         r.path = b"foobarfoo"
diff --git a/test/mitmproxy/test_flow.py b/test/mitmproxy/test_flow.py
index f4d32cbb..7a37fbd1 100644
--- a/test/mitmproxy/test_flow.py
+++ b/test/mitmproxy/test_flow.py
@@ -140,7 +140,7 @@ class TestFlowMaster:
 
     def test_create_flow(self):
         fm = master.Master(None, DummyServer())
-        assert fm.create_request("GET", "http", "example.com", 80, "/")
+        assert fm.create_request("GET", "http://example.com/")
 
     def test_all(self):
         s = tservers.TestState()