"""
Module representing response payload classes.
"""
import importlib
from http.cookies import SimpleCookie, Morsel
from datetime import (
timedelta,
datetime,
)
from typing import (
Optional,
Dict,
Any,
Union,
)
from duck.etc.statuscodes import responses
from duck.http.headers import Headers
from duck.utils.object_mapping import map_data_to_object
from duck.utils.importer import import_module_once
[docs]
class BaseResponsePayload:
"""
BaseResponsePayload class.
"""
@property
def raw(self):
raise NotImplementedError("Property `raw` should be implemented.")
@property
def cookies(self) -> SimpleCookie:
"""
Property getter for cookies. If not already initialized, it initializes the cookies.
Returns:
SimpleCookie: The cookies for the response.
"""
if not hasattr(self, '_cookies'):
self._cookies = SimpleCookie() # Initialize cookies if not already set
return self._cookies
@cookies.setter
def cookies(self, cookies_obj: SimpleCookie) -> None:
"""
Setter for cookies. Assigns a SimpleCookie object to the internal cookies attribute.
Args:
cookies_obj (SimpleCookie): The SimpleCookie object to set.
"""
if not isinstance(cookies_obj, SimpleCookie):
raise ValueError("Expected a SimpleCookie object.")
self._cookies = cookies_obj
[docs]
def set_multiple_cookies(self, cookies: Dict[str, Dict[str, Any]]) -> None:
"""
Sets multiple cookies at once. Each cookie is specified by a dictionary of attributes.
Args:
cookies (Dict[str, Dict[str, Any]]): A dictionary where the key is the cookie name
and the value is another dictionary of cookie attributes.
"""
for cookie_name, attributes in cookies.items():
self.set_cookie(cookie_name, **attributes)
[docs]
def get_cookie_obj(self, name: str) -> Optional[Morsel]:
"""
Retrieves the cookie object/morsel of a specific cookie by name.
Args:
name (str): The name of the cookie to retrieve.
Returns:
Optional[Morsel]: The cookie object, or None if the cookie does not exist.
"""
return self._cookies.get(name, '') if name in self._cookies else None
[docs]
def get_cookie(self, name: str) -> str:
"""
Retrieves the value of a specific cookie by name.
Args:
name (str): The name of the cookie to retrieve.
Returns:
str: The cookie value, or an empty string if the cookie does not exist.
"""
return self._cookies.get(name, '').value if name in self._cookies else ''
[docs]
def get_cookie_str(self, name: str, include_cookie_name: bool = True) -> str:
"""
Returns the cookie string for the provided name with all fields including max-age, domain, path, etc.
Args:
include_cookie_name (bool): Whether to cookie name e.g. `cookie=something;`. Defaults to True.
"""
cookie_obj = self.get_cookie_obj(name)
if not cookie_obj:
return ""
cookie_str = cookie_obj.output(header="").strip()
if not include_cookie_name:
cookie_str = cookie_str.split(name, 1)[-1].split('=', 1)[-1].strip()
return cookie_str
[docs]
def get_all_cookies(self) -> Dict[str, str]:
"""
Retrieves all cookies as a dictionary.
Returns:
Dict[str, str]: A dictionary of all cookies, where the key is the cookie name and the value is the cookie value (without other cookie properties).
"""
return {key: morsel.value for key, morsel in self._cookies.items()}
[docs]
def set_cookie(
self,
key: str,
value: str = "",
domain: Optional[str] = None,
path: str = "/",
max_age: Optional[Union[int, timedelta]] = None,
expires: Optional[Union[datetime, str]] = None,
secure: bool = False,
httponly: bool = False,
samesite: Optional[str] = "Lax",
) -> None:
"""
Set a custom cookie on the response payload.
Args:
key (str): The name of the cookie (e.g., 'csrftoken').
value (str): The value of the cookie (e.g., 'some_random_token').
domain (Optional[str]): The domain to set for the cookie. Defaults to None.
path (str): The path for the cookie. Defaults to '/' indicating the whole site.
max_age (Optional[Union[int, timedelta]]): The maximum age of the cookie in seconds or as a timedelta object.
expires (Optional[Union[datetime, str]]): The expiration date of the cookie. Defaults to None.
secure (bool): Whether the cookie should only be sent over HTTPS connections. Defaults to False.
httponly (bool): Whether the cookie should be inaccessible to JavaScript. Defaults to False.
samesite (Optional[str]): The SameSite attribute for the cookie. Default is 'Lax'. Other possible values are 'Strict' or 'None'.
Raises:
ValueError: If an invalid value is provided for `samesite`.
"""
# Validate SameSite value
valid_samesite_values = {"Lax", "Strict", "None", None}
if samesite is not None:
samesite = str(samesite).capitalize() # Normalize capitalization
if samesite not in valid_samesite_values:
raise ValueError(f"Invalid 'samesite' value: {samesite}. Must be one of {valid_samesite_values}.")
# Set the cookie
self.cookies[key] = value
cookie = self.cookies[key]
# Optional attributes
if domain:
cookie["domain"] = domain
cookie["path"] = path
if max_age is not None:
# Convert timedelta to seconds if necessary
max_age_seconds = max_age.total_seconds() if isinstance(max_age, timedelta) else max_age
cookie["max-age"] = int(max_age_seconds)
# Calculate expires if not explicitly provided
if expires is None:
expires = datetime.utcnow() + timedelta(seconds=max_age_seconds)
if expires is not None:
if isinstance(expires, datetime):
expires = expires.strftime('%a, %d %b %Y %H:%M:%S GMT')
cookie["expires"] = expires
if secure:
cookie["secure"] = True
if httponly:
cookie["httponly"] = True
if samesite:
cookie["samesite"] = samesite
[docs]
def delete_cookie(self, key: str, path: str = "/", domain: Optional[str] = None) -> None:
"""
Delete a cookie from the response payload by setting its expiration date to the past.
This will prompt the client to remove the cookie.
Args:
key (str): The name of the cookie to delete.
path (str): The path for which the cookie was set. Defaults to "/".
domain (Optional[str]): The domain for which the cookie was set. Defaults to None.
"""
cookie = self.cookies[key] = ""
cookie["path"] = path
cookie["domain"] = domain or ""
cookie["expires"] = datetime.utcfromtimestamp(0).strftime('%a, %d %b %Y %H:%M:%S GMT') # unix epoch
cookie["max-age"] = 0
cookie["secure"] = False
cookie["httponly"] = False
cookie["samesite"] = "Lax"
[docs]
class SimpleHttpResponsePayload(BaseResponsePayload):
"""
A class representing a simplified HTTP response payload, storing basic response metadata such as the
top header (status line) and HTTP headers.
This class is used for lightweight storage and manipulation of HTTP response headers and status-related
information. It does not handle the response body.
Attributes:
cookies: The http.cookies.SimpleCookie object for all cookies.
topheader (str): The status line of the HTTP response (e.g., "200 OK"). Defaults to an empty string.
headers (Headers): A `Headers` object containing the HTTP response headers. Defaults to `None`.
Methods:
- set_topheader(topheader: str):
Updates the top header (status line) of the response.
- set_headers(headers: dict):
Updates the HTTP headers using a dictionary input.
Properties:
- status_code (int):
Extracts and returns the status code from the `topheader`.
- status_message (str):
Extracts and returns the status message from the `topheader`. Defaults to an empty string if the
`topheader` is not set.
- explanation (str):
A placeholder for additional explanation or description of the status. Currently always returns an
empty string.
"""
def __init__(self, topheader: str = "", headers: Optional[Dict] = None):
"""
Initializes a SimpleHttpResponsePayload instance.
Args:
topheader (str): The HTTP response status line (e.g., "HTTP/1.1 200 OK"). Defaults to an empty string.
headers (Optional[Dict]): A dictionary representing the HTTP response headers. Defaults to `None`.
"""
if not topheader or not isinstance(topheader, str):
raise ValueError("The 'topheader' argument must be a non-empty string.")
self.topheader = topheader
self.headers = Headers(headers or {})
@property
def raw(self) -> bytes:
"""
Constructs and returns the raw HTTP response as bytes.
This property generates a raw representation of the HTTP response by combining the `topheader` (status line)
and all headers stored in the `headers` attribute. Each header is formatted as "Header-Name: value", separated
by carriage return and newline characters (`\r\n`). The resulting raw response is encoded as UTF-8.
Returns:
bytes: The raw HTTP response as bytes.
"""
data = f"{self.topheader}"
for header, value in self.headers.titled_headers().items():
data += f"\r\n{header}: {value}"
cookies = self.cookies.output()
data += "\r\n" + cookies
return data.encode("utf-8").strip()
@property
def status_code(self) -> int:
"""
Extracts and returns the HTTP status code from the status line.
Returns:
int: The HTTP status code (e.g., 200).
"""
return int(self.topheader.split(" ", 2)[1])
@property
def status_message(self) -> str:
"""
Extracts and returns the HTTP status message from the status line.
Returns:
str: The HTTP status message (e.g., "OK").
"""
return self.topheader.split(" ", 2)[-1]
@property
def explanation(self) -> str:
"""
Placeholder for additional status explanation.
Returns:
str: Always returns an empty string.
"""
return ""
@property
def http_version(self) -> str:
"""
Extracts and returns the HTTP version from the `topheader`.
This property analyzes the `topheader` (status line) of the HTTP response and retrieves the HTTP version,
such as "HTTP/1.1" or "HTTP/2". If the `topheader` is not set or improperly formatted, it defaults to the
last version in `HttpRequest.SUPPORTED_HTTP_VERSIONS`.
Returns:
str: The HTTP version extracted from the `topheader` or the default version if unavailable.
"""
from duck.http.request import HttpRequest
if self.topheader:
return self.topheader.split(" ", 1)[0]
return HttpRequest.SUPPORTED_HTTP_VERSIONS[-1] # Default to the latest supported HTTP version.
[docs]
def __repr__(self):
return f'<{self.__class__.__name__} "{self.topheader}">'
[docs]
class HttpResponsePayload(BaseResponsePayload):
"""
ResponsePayload class for storing response top header and headers.
Example payload:
```
HTTP/1.1 200 OK\r\n
Connection: close\r\n
Content-Type: text/html\r\n
```
"""
def __init__(self, **kwargs):
"""
Args:
status_code (int): The status code for the response. Default value is 200.
status_message (str): The status message corresponding to the status code.
headers (Headers): The headers to attach to the response payload.
http_version (str): The HTTP version for the server. Default value is 'HTTP/2.1'
explanation (str): A longer explanation of the status message.
"""
from duck.http.request import HttpRequest
self.status_code: int = 200
self.status_message: str = ""
self.headers: Headers = Headers()
self.http_version: str = HttpRequest.SUPPORTED_HTTP_VERSIONS[-1]
self.explanation: str = ""
# map kwargs to self
map_data_to_object(self, kwargs)
@property
def raw(self) -> bytes:
"""
Returns the raw bytes representing the payload (header section).
"""
topheader = "%s %d %s" % (
self.http_version,
self.status_code,
self.status_message,
)
data = f"{topheader}"
for header, value in self.headers.titled_headers().items():
data += f"\r\n{header}: {value}"
cookies = self.cookies.output()
data += "\r\n" + cookies
return data.encode("utf-8").strip()
[docs]
def parse_status(
self,
code: int,
msg: str = None,
explanation: str = None):
"""
Parses topheader status for the payload.
"""
if code in responses:
# status code exists
short_msg, long_msg = responses[code]
if not msg:
msg = short_msg
if not explanation:
explanation = long_msg
else:
if not msg:
msg = "???"
if not explanation:
if msg != "???":
explanation = msg
else:
explanation = "???"
self.status_message = msg
self.status_code = code
self.explanation = explanation
[docs]
def __repr__(self):
return f"<{self.__class__.__name__} {self.status_code}>"
