96 lines
2.9 KiB
Python
96 lines
2.9 KiB
Python
from __future__ import annotations
|
|
|
|
import base64
|
|
from typing import Any
|
|
from urllib.parse import unquote
|
|
|
|
from mangum.types import Headers, LambdaConfig
|
|
|
|
|
|
def maybe_encode_body(body: str | bytes, *, is_base64: bool) -> bytes:
|
|
body = body or b""
|
|
if is_base64:
|
|
body = base64.b64decode(body)
|
|
elif not isinstance(body, bytes):
|
|
body = body.encode()
|
|
|
|
return body
|
|
|
|
|
|
def get_server_and_port(headers: dict[str, Any]) -> tuple[str, int]:
|
|
server_name = headers.get("host", "mangum")
|
|
if ":" not in server_name:
|
|
server_port = headers.get("x-forwarded-port", 80)
|
|
else:
|
|
server_name, server_port = server_name.split(":") # pragma: no cover
|
|
server = (server_name, int(server_port))
|
|
|
|
return server
|
|
|
|
|
|
def strip_api_gateway_path(path: str, *, api_gateway_base_path: str) -> str:
|
|
if not path:
|
|
return "/"
|
|
|
|
if api_gateway_base_path and api_gateway_base_path != "/":
|
|
if not api_gateway_base_path.startswith("/"):
|
|
api_gateway_base_path = f"/{api_gateway_base_path}"
|
|
if path.startswith(api_gateway_base_path):
|
|
path = path[len(api_gateway_base_path) :]
|
|
|
|
return unquote(path)
|
|
|
|
|
|
def handle_multi_value_headers(
|
|
response_headers: Headers,
|
|
) -> tuple[dict[str, str], dict[str, list[str]]]:
|
|
headers: dict[str, str] = {}
|
|
multi_value_headers: dict[str, list[str]] = {}
|
|
for key, value in response_headers:
|
|
lower_key = key.decode().lower()
|
|
if lower_key in multi_value_headers:
|
|
multi_value_headers[lower_key].append(value.decode())
|
|
elif lower_key in headers:
|
|
# Move existing to multi_value_headers and append current
|
|
multi_value_headers[lower_key] = [
|
|
headers[lower_key],
|
|
value.decode(),
|
|
]
|
|
del headers[lower_key]
|
|
else:
|
|
headers[lower_key] = value.decode()
|
|
return headers, multi_value_headers
|
|
|
|
|
|
def handle_base64_response_body(
|
|
body: bytes,
|
|
headers: dict[str, str],
|
|
text_mime_types: list[str],
|
|
) -> tuple[str, bool]:
|
|
is_base64_encoded = False
|
|
output_body = ""
|
|
if body != b"":
|
|
for text_mime_type in text_mime_types:
|
|
if text_mime_type in headers.get("content-type", ""):
|
|
try:
|
|
output_body = body.decode()
|
|
except UnicodeDecodeError:
|
|
output_body = base64.b64encode(body).decode()
|
|
is_base64_encoded = True
|
|
break
|
|
else:
|
|
output_body = base64.b64encode(body).decode()
|
|
is_base64_encoded = True
|
|
|
|
return output_body, is_base64_encoded
|
|
|
|
|
|
def handle_exclude_headers(headers: dict[str, Any], config: LambdaConfig) -> dict[str, Any]:
|
|
finalized_headers = {}
|
|
for header_key, header_value in headers.items():
|
|
if header_key in config["exclude_headers"]:
|
|
continue
|
|
finalized_headers[header_key] = header_value
|
|
|
|
return finalized_headers
|