diff options
| -rw-r--r-- | docs/dev/models.rst | 2 | ||||
| -rw-r--r-- | examples/redirect_requests.py | 7 | ||||
| -rw-r--r-- | mitmproxy/console/flowview.py | 5 | ||||
| -rw-r--r-- | netlib/http/response.py | 54 | ||||
| -rw-r--r-- | 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 7260f1f7..a333fb06 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 36594bcd..8cde1bfd 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 6d74be65..15be379b 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 85f54940..7866c142 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 b3c2f736..c7b1b646 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") |