예제 #1
0
        async def send_wrapper(message: Message) -> None:
            if message["type"] == "http.response.start":
                session: Session = scope["session"]
                if session.is_modified and not session.is_empty:
                    # We have session data to persist (data was changed, cleared, etc).
                    nonlocal session_id
                    session_id = await scope["session"].persist()

                    headers = MutableHeaders(scope=message)
                    header_value = "%s=%s; path=/; Max-Age=%d; %s" % (
                        self.session_cookie,
                        session_id,
                        self.max_age,
                        self.security_flags,
                    )
                    headers.append("Set-Cookie", header_value)
                elif session.is_loaded and session.is_empty:
                    # no interactions to session were done
                    headers = MutableHeaders(scope=message)
                    header_value = "%s=%s; %s" % (
                        self.session_cookie,
                        "null; path=/; expires=Thu, 01 Jan 1970 00:00:00 GMT;",
                        self.security_flags,
                    )
                    headers.append("Set-Cookie", header_value)
            await send(message)
예제 #2
0
 async def send_wrapper(message: Message) -> None:
     if message["type"] == "http.response.start":
         if scope["session"]:
             # We have session data to persist.
             data = b64encode(
                 json.dumps(scope["session"]).encode("utf-8"))
             data = self.signer.sign(data)
             headers = MutableHeaders(scope=message)
             header_value = "{session_cookie}={data}; path={path}; {max_age}{security_flags}".format(  # noqa E501
                 session_cookie=self.session_cookie,
                 data=data.decode("utf-8"),
                 path=self.path,
                 max_age=f"Max-Age={self.max_age}; "
                 if self.max_age else "",
                 security_flags=self.security_flags,
             )
             headers.append("Set-Cookie", header_value)
         elif not initial_session_was_empty:
             # The session has been cleared.
             headers = MutableHeaders(scope=message)
             header_value = "{session_cookie}={data}; path={path}; {expires}{security_flags}".format(  # noqa E501
                 session_cookie=self.session_cookie,
                 data="null",
                 path=self.path,
                 expires="expires=Thu, 01 Jan 1970 00:00:00 GMT; ",
                 security_flags=self.security_flags,
             )
             headers.append("Set-Cookie", header_value)
     await send(message)
예제 #3
0
        async def send_wrapper(message: Message) -> None:
            if message["type"] == "http.response.start":
                if scope["session"]:
                    if "exp" not in scope["session"]:
                        scope["session"]["exp"] = int(
                            time.time()) + self.max_age
                    data = jwt.encode(self.jwt_header, scope["session"],
                                      str(self.jwt_secret.encode))

                    headers = MutableHeaders(scope=message)
                    header_value = "%s=%s; path=/; Max-Age=%d; %s" % (
                        self.session_cookie,
                        data.decode("utf-8"),
                        self.max_age,
                        self.security_flags,
                    )
                    if self.domain:  # pragma: no cover
                        header_value += f"; domain={self.domain}"
                    headers.append("Set-Cookie", header_value)
                elif not initial_session_was_empty:
                    # The session has been cleared.
                    headers = MutableHeaders(scope=message)
                    header_value = "%s=%s; %s" % (
                        self.session_cookie,
                        "null; path=/; expires=Thu, 01 Jan 1970 00:00:00 GMT;",
                        self.security_flags,
                    )
                    if self.domain:  # pragma: no cover
                        header_value += f"; domain={self.domain}"
                    headers.append("Set-Cookie", header_value)
            await send(message)
예제 #4
0
 async def send_wrapper(message: Message) -> None:
     if message["type"] == "http.response.start":
         path = scope.get("root_path", "") or "/"
         if scope["session"]:
             # We have session data to persist.
             data = b64encode(
                 json.dumps(scope["session"]).encode("utf-8"))
             data = self.signer.sign(data)
             headers = MutableHeaders(scope=message)
             header_value = "%s=%s; path=%s; Max-Age=%d; %s" % (
                 self.session_cookie,
                 data.decode("utf-8"),
                 path,
                 self.max_age,
                 self.security_flags,
             )
             headers.append("Set-Cookie", header_value)
         elif not initial_session_was_empty:
             # The session has been cleared.
             headers = MutableHeaders(scope=message)
             header_value = "{}={}; {}".format(
                 self.session_cookie,
                 f"null; path={path}; expires=Thu, 01 Jan 1970 00:00:00 GMT;",
                 self.security_flags,
             )
             headers.append("Set-Cookie", header_value)
     await send(message)
예제 #5
0
 async def sender(message: Message) -> None:
     if message["type"] == "http.response.start":
         if scope["session"]:
             # We have session data to persist.
             data = b64encode(
                 json.dumps(scope["session"]).encode("utf-8"))
             data = self.signer.sign(data)
             headers = MutableHeaders(scope=message)
             header_value = "%s=%s; path=/; Max-Age=%d; %s" % (
                 self.session_cookie,
                 data.decode("utf-8"),
                 self.max_age,
                 self.security_flags,
             )
             headers.append("Set-Cookie", header_value)
         elif not was_empty_session:
             # The session has been cleared.
             headers = MutableHeaders(scope=message)
             header_value = "%s=%s; %s" % (
                 self.session_cookie,
                 "null; path=/; expires=Thu, 01 Jan 1970 00:00:00 GMT;",
                 self.security_flags,
             )
             headers.append("Set-Cookie", header_value)
     await send(message)
예제 #6
0
        def send_wrapper(message: Message):
            if message["type"] == "http.response.start":
                counter = traced_counter.get()
                if counter and isinstance(counter, Counter):
                    headers = MutableHeaders(scope=message)
                    headers.append("x-dagster-call-counts", json.dumps(counter.counts()))

            return send(message)
예제 #7
0
 async def send_wrapper(message: Message) -> None:
     if message['type'] == 'http.response.start':
         if scope['session']:
             _data = base64.b64encode(
                 json.dumps(scope['session']).encode('utf-8'))
             _data = self.signer.sign(_data)
             headers = MutableHeaders(scope=message)
             headers.append('session', _data.decode('utf-8'))
     await send(message)
예제 #8
0
 async def send(self, message, send, request):
     if message["type"] != "http.response.start":
         await send(message)
         return
     headers = MutableHeaders(scope=message)
     req_headers = {k.lower(): v for k, v in dict(request.headers).items()}
     headers.append("x-b3-traceid", req_headers.get("x-b3-traceid", ""))
     headers.append("x-b3-sampled", req_headers.get("x-b3-sampled", ""))
     await send(message)
예제 #9
0
    async def enrich_response(self, arg) -> None:
        value = str(context.get(self.key))

        # for ContextMiddleware
        if isinstance(arg, Response):
            arg.headers[self.key] = value
        # for ContextPureMiddleware
        else:
            if arg["type"] == "http.response.start":
                headers = MutableHeaders(scope=arg)
                headers.append(self.key, value)
예제 #10
0
 def set_cookie(
     self,
     message: Message,
     value: str,
     max_age: int = None,
 ):
     headers = MutableHeaders(scope=message)
     headers.append("Cache-Control", "no-cache")
     headers.append(
         "Set-Cookie",
         f"{self.session_cookie}={value};"
         f" path={self.path};"
         f" Max-Age={max_age or self.max_age};"
         f" {self.security_flags}",
     )
예제 #11
0
 async def send_wrapper(message: Message) -> None:
     if message['type'] == 'http.response.start':
         if scope['session']:
             # We have session data to persist.
             headers = MutableHeaders(scope=message)
             header_value = '%s=%s; path=/; Max-Age=%d; %s' % (
                 self.session_cookie,
                 self.backend.encode(scope['session']),
                 self.max_age,
                 self.security_flags,
             )
             headers.append('Set-Cookie', header_value)
         elif not initial_session_was_empty:
             # The session has been cleared.
             headers = MutableHeaders(scope=message)
             header_value = '%s=%s; %s' % (
                 self.session_cookie,
                 'null; path=/; expires=Thu, 01 Jan 1970 00:00:00 GMT;',
                 self.security_flags,
             )
             headers.append('Set-Cookie', header_value)
     await send(message)
예제 #12
0
 async def send_wrapper(self, msg: Message):
     if msg["type"] == "http.response.start":
         result = self.rl_res
         headers = MutableHeaders(scope=msg)
         headers.append("X-Rate-Limit-Limit", str(result.consumed_points))
         headers.append("X-Rate-Limit-Remaining",
                        str(result.remaining_points))
         headers.append("X-Rate-Limit-Reset", str(result.ms_before_next))
     await self.send(msg)
예제 #13
0
        async def send_wrapper(message: Message, **kwargs) -> None:
            if message["type"] == "http.response.start":

                session_key = scope.pop("__session_key", str(uuid4()))

                if scope["session"]:

                    if (self.backend_type == BackendType.cookie
                            or not self.session_backend):
                        cookie_data = scope["session"]
                    else:
                        await self.session_backend.set(session_key,
                                                       scope["session"],
                                                       self.max_age)
                        cookie_data = {
                            self._cookie_session_id_field: session_key
                        }

                    data = b64encode(json.dumps(cookie_data).encode("utf-8"))
                    data = self.signer.sign(data)

                    headers = MutableHeaders(scope=message)
                    header_value = self._construct_cookie(clear=False,
                                                          data=data)
                    headers.append("Set-Cookie", header_value)

                elif not initial_session_was_empty:

                    if self.session_backend and self.backend_type != BackendType.cookie:
                        await self.session_backend.delete(session_key)

                    headers = MutableHeaders(scope=message)
                    header_value = self._construct_cookie(clear=True)
                    headers.append("Set-Cookie", header_value)

            await send(message)
예제 #14
0
class Response:
    """
    Outgoing HTTP response class. `Response` instance is ASGI3 application.

    **Attributes**

    * headers (`MutableHeaders`): The response headers, case-insensitive dictionary.
        To set values having same key, use `headers.append()` .
    * cookies (`SimpleCookie`): Dict-like http cookies. `Set-Cookie` header refers this.
        You can set cookie-attributes.
    * status (`int`): The response's status code.
    * streaming (`Optional[AsyncGenerator]`): Async generator for streaming. If set,
        other response body attrs like `media` are ignored.
    * mimetype (`str`): The mediatype of the response body.
    * reraise (`bool`): In ErrorHandler, if set true, reraise the exception after
        sending data.

    """

    __slots__ = (
        "_redirect_to",
        "_starlette_resp",
        "_body",
        "_text",
        "_content",
        "_json",
        "headers",
        "cookies",
        "status_code",
        "streaming",
        "reraise",
    )

    _redirect_to: Optional[tuple[str, Optional[str]]]
    _starlette_resp: type[StarletteResponse]
    _body: Any
    _text: Optional[str]
    _content: Optional[bytes]
    _json: Union[addict.Dict, list, None]

    headers: MutableHeaders
    cookies: SimpleCookie
    status_code: int
    streaming: Optional[AsyncGenerator]
    reraise: bool

    def __init__(self) -> None:
        """Do not use manually."""
        self._redirect_to = None
        self._starlette_resp = StarletteResponse
        self._body = None
        self._text = None
        self._content = None
        self._json = None

        self.headers = MutableHeaders()
        self.cookies = SimpleCookie()
        self.status_code = HTTPStatus.OK
        self.streaming: Optional[AsyncGenerator] = None
        self.reraise = False

    async def __call__(self, scope: Scope, receive: Receive, send: Send):
        self._set_cookies_to_headers()
        if self.streaming:
            self._starlette_resp = StreamingResponse
            self._body = self.streaming
        elif self._json is not None:
            self._starlette_resp = JSONResponse
            self._body = self.json
            self.headers["content-type"] = "application/json"
        elif self._redirect_to is not None:
            url, qs = self._redirect_to
            if qs is None:
                qs = scope.get("query_string", b"").decode("ascii")

            if qs:
                qs = f"?{qs}"

            self._body = url + qs

        app = self._starlette_resp(self._body,
                                   status_code=self.status_code,
                                   headers=self.headers)

        return await app(scope, receive, send)

    def set_status(self, status: int) -> "Response":
        """
        Set HTTP status code.

        **Args**

        * status (`int`): HTTP status code.

        **Returns**

        * `spangle.models.http.Response` : Return self.

        """
        self.status_code = status
        return self

    def set_cookie(
        self,
        key: str,
        value: str = "",
        max_age: Optional[int] = None,
        expires: Optional[int] = None,
        path: Optional[str] = "/",
        comment: Optional[str] = None,
        domain: Optional[str] = None,
        secure: bool = False,
        httponly: bool = True,
        version: Optional[int] = None,
        samesite: Optional[str] = "Lax",
    ) -> "Response":
        """
        Set cookie value to given key with params.

        **Args**

        * key (`str`)
        * value (`str`)

        Cookie options:

        * max_age (`Optional[int]`)
        * expires (`Optional[int]`)
        * path (`Optional[str]`)
        * comment (`Optional[str]`)
        * domain (`Optional[str]`)
        * secure (`bool`)
        * httponly (`bool`)
        * version (`Optional[int]`)
        * samesite (`Optional[str]`)

        **Returns**

        * `spangle.models.http.Response`: Return self.

        """
        self.cookies[key] = value
        if max_age is not None:
            self.cookies[key]["max-age"] = max_age
        if expires is not None:
            self.cookies[key]["expires"] = expires
        if path is not None:
            self.cookies[key]["path"] = path
        if comment is not None:
            self.cookies[key]["comment"] = comment
        if domain is not None:
            self.cookies[key]["domain"] = domain
        if version is not None:
            self.cookies[key]["version"] = version
        if samesite is not None:
            self.cookies[key]["samesite"] = samesite

        self.cookies[key]["secure"] = secure
        self.cookies[key]["httponly"] = httponly
        return self

    def delete_cookie(self,
                      key: str,
                      path: str = "/",
                      domain: str = None) -> "Response":
        """
        Remove cookie value from client.

        **Args**

        * key (`str`)

        Cookie options:

        * path (`str`)
        * domain (`str`)

        **Returns**

        * `spangle.models.http.Response`: Return self.

        """
        self.set_cookie(key, expires=0, max_age=0, path=path, domain=domain)
        return self

    @property
    def json(self) -> Any:
        """
        (`Any`): A dict sent to the client. Default-type: `"application/json"` .
            You can set values like `resp.json.keyName.you = "want"` .

        """
        if self._json is None:
            self._json = addict.Dict()
        return self._json

    @json.setter
    def json(self, v: Any):
        self._json = v

    @property
    def text(self) -> Optional[str]:
        """
        (`str`): A unicode string of the response body. Default-type: `"text/plain"` .
        """
        return self._text

    @text.setter
    def text(self, t: str):
        self.set_text(t)

    @property
    def content(self) -> Optional[bytes]:
        """
        (`bytes`): Bytes of the response body. Default-type:
            `"application/octet-stream"` .
        """
        return self._content

    @content.setter
    def content(self, c: bytes):
        self.set_content(c)

    def add_header(self, key: str, value: str) -> "Response":
        """
        Append new header. To overwrite, use `spangle.models.http.Response.set_header` .

        **Args**

        * key (`str`): Header's key.
        * value (`str`): Header's value.

        **Returns**

        * `spangle.models.http.Response`: Return self.

        """
        self.headers.append(key, value)
        return self

    def set_header(self, key: str, value: str) -> "Response":
        """
        Set HTTP header value to given key. It overwrites value if exists.

        **Args**

        * key (`str`): Header's key.
        * value (`str`): Header's value.

        **Returns**

        * `spangle.models.http.Response`: Return self.

        """
        self.headers[key] = value
        return self

    def set_text(self, text: str, content_type="text/plain") -> "Response":
        """
        Set given text to response body with content type.

        **Args**

        * text (`str`): Response body as UTF-8 string.
        * content_type (`str`): Response content type.

        **Returns**

        * `spangle.models.http.Response`: Return self.

        """
        self._text = text
        self._body = self.text
        mark_utf8 = "; charset=utf-8"
        if mark_utf8 not in content_type.lower():
            content_type = f"{content_type}{mark_utf8}"
        self.headers["content-type"] = content_type
        self._starlette_resp = StarletteResponse
        return self

    def set_content(self,
                    content: bytes,
                    content_type="application/octet-stream") -> "Response":
        """
        Set bytes to response body with content type.

        **Args**

        * content (`bytes`): Response body as bytes.
        * content_type (`str`): Response content type.

        **Returns**

        * `spangle.models.http.Response`: Return self.

        """
        self._content = content
        self._body = self._content
        self.headers["content-type"] = content_type
        self._starlette_resp = StarletteResponse
        return self

    async def load_template(self,
                            template_name: str,
                            content_type="text/html",
                            **params) -> "Response":
        """
        Load `jinja2` template, render, set headers & text.

        **Args**

        * template_name (`str`): The template `"path/name"` .
        * content_type (`str`): `"text/html"` .
        * **params: Variables used in the template. `api` is reserved by
            `spangle.api.Api` instance by default.

        **Returns**

        * `spangle.models.http.Response`: Return self.

        **Raises**

        * `ValueError`: Missing `jinja2` env in `spangle.api.Api` instance.
        * `NotFoundError`: Missing requested template.

        """
        jinja_env = use_api()._jinja_env
        if jinja_env is None:
            raise ValueError("Set jinja env.")

        try:
            template = jinja_env.get_template(template_name)
        except jinja2.exceptions.TemplateNotFound:
            raise NotFoundError

        rendered = await template.render_async(**params)
        # Set result.
        self.set_text(rendered, content_type)

        return self

    def redirect(
        self,
        *,
        view: type = None,
        params: dict = None,
        url: str = None,
        status=HTTPStatus.TEMPORARY_REDIRECT,
        query_string: Optional[str] = None,
    ) -> "Response":
        """
        Set redirect view/location. Positional args are not allowed.

        If both `view` and `url` are set, `url` is ignored.

        **Args**

        * view (`Type`): View class that the client redirect to.
        * params (`dict`): Dynamic URL params passed to the view.
        * url (`str`): The location out of the app.
        * status (`int`): HTTP status code. Must be `300<=status<400` .

        **Returns**

        * `spangle.models.http.Response`: Return self.

        """

        if bool(view) is bool(url):
            raise TypeError("Set only one location; view-class or url.")
        if not (300 <= status < 400):
            raise ValueError("Set correct status.")

        self.status_code = status
        self._starlette_resp = RedirectResponse
        if view:
            url_for = use_api().url_for
            redirect_to = url_for(view, params)
        elif url:
            redirect_to = url
        else:
            raise TypeError("Set only one location; view-class or url.")

        self._redirect_to = redirect_to, query_string
        return self

    def _set_cookies_to_headers(self) -> None:
        if not self.cookies:
            return
        cookies = self.cookies.output(header="").split("\r\n")
        for c in cookies:
            self.headers.append("Set-Cookie", c.lstrip())