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 .. autoclass:: Response
.. automethod:: make
.. rubric:: Data .. rubric:: Data
.. autoattribute:: http_version .. autoattribute:: http_version
.. autoattribute:: status_code .. autoattribute:: status_code

View File

@ -2,7 +2,6 @@
This example shows two ways to redirect flows to other destinations. This example shows two ways to redirect flows to other destinations.
""" """
from mitmproxy.models import HTTPResponse from mitmproxy.models import HTTPResponse
from netlib.http import Headers
def request(flow): def request(flow):
@ -12,11 +11,7 @@ def request(flow):
# Method 1: Answer with a locally generated response # Method 1: Answer with a locally generated response
if flow.request.pretty_host.endswith("example.com"): if flow.request.pretty_host.endswith("example.com"):
resp = HTTPResponse( resp = HTTPResponse.make(200, b"Hello World", {"Content-Type": "text/html"})
b"HTTP/1.1", 200, b"OK",
Headers(Content_Type="text/html"),
b"helloworld"
)
flow.reply.send(resp) flow.reply.send(resp)
# Method 2: Redirect the request to a different server # Method 2: Redirect the request to a different server

View File

@ -374,10 +374,7 @@ class FlowView(tabs.Tabs):
message = self.flow.request message = self.flow.request
else: else:
if not self.flow.response: if not self.flow.response:
self.flow.response = models.HTTPResponse( self.flow.response = models.HTTPResponse.make(200, b"")
self.flow.request.http_version,
200, b"OK", Headers(), b""
)
message = self.flow.response message = self.flow.response
self.flow.backup() self.flow.backup()

View File

@ -1,14 +1,19 @@
from __future__ import absolute_import, print_function, division from __future__ import absolute_import, print_function, division
from email.utils import parsedate_tz, formatdate, mktime_tz
import time
import six 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 cookies
from netlib.http import headers as nheaders from netlib.http import headers as nheaders
from netlib.http import message from netlib.http import message
from netlib import multidict from netlib.http import status_codes
from netlib import human 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): class ResponseData(message.MessageData):
@ -54,6 +59,45 @@ class Response(message.Message):
details=details 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 @property
def status_code(self): def status_code(self):
""" """

View File

@ -5,6 +5,7 @@ import email
import time import time
from netlib.http import Headers from netlib.http import Headers
from netlib.http import Response
from netlib.http.cookies import CookieAttrs from netlib.http.cookies import CookieAttrs
from netlib.tutils import raises, tresp from netlib.tutils import raises, tresp
from .test_message import _test_passthrough_attr, _test_decoded_attr from .test_message import _test_passthrough_attr, _test_decoded_attr
@ -28,6 +29,25 @@ class TestResponseCore(object):
response.content = None response.content = None
assert repr(response) == "Response(200 OK, no content)" 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): def test_status_code(self):
_test_passthrough_attr(tresp(), "status_code") _test_passthrough_attr(tresp(), "status_code")