introduce Response.make for simple response creation

This commit is contained in:
Maximilian Hils 2016-08-23 00:17:06 -07:00
parent 798ce96bd0
commit f27028f58e
5 changed files with 73 additions and 15 deletions

View File

@ -39,6 +39,8 @@ Datastructures
.. autoclass:: Response
.. automethod:: make
.. rubric:: Data
.. autoattribute:: http_version
.. autoattribute:: status_code

View File

@ -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

View File

@ -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()

View File

@ -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):
"""

View File

@ -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")