78 lines
2.2 KiB
Python
78 lines
2.2 KiB
Python
"""Helpers shared by the security tests."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from collections.abc import Iterable
|
|
from dataclasses import dataclass, field
|
|
from typing import Any
|
|
|
|
from starlette.types import ASGIApp, Message, Scope
|
|
|
|
|
|
@dataclass
|
|
class AsgiResult:
|
|
"""What one raw ASGI call produced and how much of the request body it consumed."""
|
|
|
|
status: int = 0
|
|
body: bytes = b""
|
|
headers: dict[str, str] = field(default_factory=dict)
|
|
bytes_read: int = 0
|
|
receive_calls: int = 0
|
|
|
|
def json(self) -> Any:
|
|
return json.loads(self.body)
|
|
|
|
|
|
async def asgi_request(
|
|
app: ASGIApp,
|
|
method: str,
|
|
path: str,
|
|
*,
|
|
headers: dict[str, str] | None = None,
|
|
chunks: Iterable[bytes] = (),
|
|
host: str = "mcp.example.com",
|
|
scheme: str = "https",
|
|
client: tuple[str, int] = ("198.51.100.7", 4000),
|
|
) -> AsgiResult:
|
|
"""Drive ``app`` with a counting ``receive`` so tests can assert how much body was read."""
|
|
result = AsgiResult()
|
|
pending = list(chunks)
|
|
raw_headers = [(b"host", host.encode())]
|
|
for name, value in (headers or {}).items():
|
|
raw_headers.append((name.lower().encode(), value.encode()))
|
|
scope: Scope = {
|
|
"type": "http",
|
|
"asgi": {"version": "3.0"},
|
|
"http_version": "1.1",
|
|
"method": method,
|
|
"scheme": scheme,
|
|
"path": path,
|
|
"raw_path": path.encode(),
|
|
"query_string": b"",
|
|
"root_path": "",
|
|
"headers": raw_headers,
|
|
"client": client,
|
|
"server": (host, 443),
|
|
}
|
|
|
|
async def receive() -> Message:
|
|
result.receive_calls += 1
|
|
if not pending:
|
|
return {"type": "http.disconnect"}
|
|
chunk = pending.pop(0)
|
|
result.bytes_read += len(chunk)
|
|
return {"type": "http.request", "body": chunk, "more_body": bool(pending)}
|
|
|
|
async def send(message: Message) -> None:
|
|
if message["type"] == "http.response.start":
|
|
result.status = message["status"]
|
|
result.headers = {
|
|
key.decode().lower(): value.decode() for key, value in message["headers"]
|
|
}
|
|
elif message["type"] == "http.response.body":
|
|
result.body += message.get("body", b"")
|
|
|
|
await app(scope, receive, send)
|
|
return result
|