From f27028f58e50492c8c8e15e47278bf444d692246 Mon Sep 17 00:00:00 2001
From: Maximilian Hils <git@maximilianhils.com>
Date: Tue, 23 Aug 2016 00:17:06 -0700
Subject: [PATCH] introduce Response.make for simple response creation

---
 docs/dev/models.rst               |  2 ++
 examples/redirect_requests.py     |  7 +---
 mitmproxy/console/flowview.py     |  5 +--
 netlib/http/response.py           | 54 ++++++++++++++++++++++++++++---
 test/netlib/http/test_response.py | 20 ++++++++++++
 5 files changed, 73 insertions(+), 15 deletions(-)

diff --git a/docs/dev/models.rst b/docs/dev/models.rst
index 7260f1f75..a333fb06a 100644
--- a/docs/dev/models.rst
+++ b/docs/dev/models.rst
@@ -39,6 +39,8 @@ Datastructures
 
     .. autoclass:: Response
 
+        .. automethod:: make
+
         .. rubric:: Data
         .. autoattribute:: http_version
         .. autoattribute:: status_code
diff --git a/examples/redirect_requests.py b/examples/redirect_requests.py
index 36594bcd8..8cde1bfd2 100644
--- a/examples/redirect_requests.py
+++ b/examples/redirect_requests.py
@@ -2,7 +2,6 @@
 This example shows two ways to redirect flows to other destinations.
 """
 from mitmproxy.models import HTTPResponse
-from netlib.http import Headers
 
 
 def request(flow):
@@ -12,11 +11,7 @@ def request(flow):
 
     # Method 1: Answer with a locally generated response
     if flow.request.pretty_host.endswith("example.com"):
-        resp = HTTPResponse(
-            b"HTTP/1.1", 200, b"OK",
-            Headers(Content_Type="text/html"),
-            b"helloworld"
-        )
+        resp = HTTPResponse.make(200, b"Hello World", {"Content-Type": "text/html"})
         flow.reply.send(resp)
 
     # Method 2: Redirect the request to a different server
diff --git a/mitmproxy/console/flowview.py b/mitmproxy/console/flowview.py
index 6d74be657..15be379b4 100644
--- a/mitmproxy/console/flowview.py
+++ b/mitmproxy/console/flowview.py
@@ -374,10 +374,7 @@ class FlowView(tabs.Tabs):
             message = self.flow.request
         else:
             if not self.flow.response:
-                self.flow.response = models.HTTPResponse(
-                    self.flow.request.http_version,
-                    200, b"OK", Headers(), b""
-                )
+                self.flow.response = models.HTTPResponse.make(200, b"")
             message = self.flow.response
 
         self.flow.backup()
diff --git a/netlib/http/response.py b/netlib/http/response.py
index 85f549407..7866c142a 100644
--- a/netlib/http/response.py
+++ b/netlib/http/response.py
@@ -1,14 +1,19 @@
 from __future__ import absolute_import, print_function, division
 
-from email.utils import parsedate_tz, formatdate, mktime_tz
-import time
 import six
-
+import time
+from email.utils import parsedate_tz, formatdate, mktime_tz
+from netlib import human
+from netlib import multidict
 from netlib.http import cookies
 from netlib.http import headers as nheaders
 from netlib.http import message
-from netlib import multidict
-from netlib import human
+from netlib.http import status_codes
+from typing import AnyStr  # noqa
+from typing import Dict  # noqa
+from typing import Iterable  # noqa
+from typing import Tuple  # noqa
+from typing import Union  # noqa
 
 
 class ResponseData(message.MessageData):
@@ -54,6 +59,45 @@ class Response(message.Message):
             details=details
         )
 
+    @classmethod
+    def make(
+            cls,
+            status_code=200,  # type: int
+            content=b"",  # type: AnyStr
+            headers=()  # type: Union[Dict[AnyStr, AnyStr], Iterable[Tuple[bytes, bytes]]]
+    ):
+        """
+        Simplified API for creating response objects.
+        """
+        resp = cls(
+            b"HTTP/1.1",
+            status_code,
+            status_codes.RESPONSES.get(status_code, "").encode(),
+            (),
+            None
+        )
+        # Assign this manually to update the content-length header.
+        if isinstance(content, bytes):
+            resp.content = content
+        elif isinstance(content, str):
+            resp.text = content
+        else:
+            raise TypeError("Expected content to be str or bytes, but is {}.".format(
+                type(content).__name__
+            ))
+
+        # Headers can be list or dict, we differentiate here.
+        if isinstance(headers, dict):
+            resp.headers = nheaders.Headers(**headers)
+        elif isinstance(headers, Iterable):
+            resp.headers = nheaders.Headers(headers)
+        else:
+            raise TypeError("Expected headers to be an iterable or dict, but is {}.".format(
+                type(headers).__name__
+            ))
+
+        return resp
+
     @property
     def status_code(self):
         """
diff --git a/test/netlib/http/test_response.py b/test/netlib/http/test_response.py
index b3c2f7364..c7b1b6469 100644
--- a/test/netlib/http/test_response.py
+++ b/test/netlib/http/test_response.py
@@ -5,6 +5,7 @@ import email
 import time
 
 from netlib.http import Headers
+from netlib.http import Response
 from netlib.http.cookies import CookieAttrs
 from netlib.tutils import raises, tresp
 from .test_message import _test_passthrough_attr, _test_decoded_attr
@@ -28,6 +29,25 @@ class TestResponseCore(object):
         response.content = None
         assert repr(response) == "Response(200 OK, no content)"
 
+    def test_make(self):
+        r = Response.make()
+        assert r.status_code == 200
+        assert r.content == b""
+
+        Response.make(content=b"foo")
+        Response.make(content="foo")
+        with raises(TypeError):
+            Response.make(content=42)
+
+        r = Response.make(headers=[(b"foo", b"bar")])
+        assert r.headers["foo"] == "bar"
+
+        r = Response.make(headers=({"foo": "baz"}))
+        assert r.headers["foo"] == "baz"
+
+        with raises(TypeError):
+            Response.make(headers=42)
+
     def test_status_code(self):
         _test_passthrough_attr(tresp(), "status_code")