|
| 1 | +"""Traffic capture, ``urn:agentenv:traffic/v1``: an env that records the packets its containers send and receive serves |
| 2 | +them as capture files through its card. ``list`` (``TRAFFIC_LIST_OP``) answers a ``TrafficCaptureListing``, each file a |
| 3 | +``TrafficCapture`` named for the container it holds, oldest first; ``get`` (``TRAFFIC_GET_OP``) streams the file a |
| 4 | +``list`` entry names, given as its ``name`` param, as ``PCAP_MEDIA_TYPE``. Both are framework-only: they take the env's |
| 5 | +control token (``client.control_headers``). ``serves_traffic`` reads whether a card advertises it, |
| 6 | +``list_traffic_captures`` and ``fetch_traffic_capture`` call it as a card advertises it.""" |
| 7 | +from __future__ import annotations |
| 8 | + |
| 9 | +from datetime import datetime |
| 10 | + |
| 11 | +import httpx |
| 12 | +from pydantic import BaseModel, ConfigDict |
| 13 | + |
| 14 | +from .client import ExtensionCallError, find_extension_method, invoke_extension, request_headers |
| 15 | + |
| 16 | +TRAFFIC_EXTENSION_URI = "urn:agentenv:traffic/v1" |
| 17 | +TRAFFIC_LIST_OP = "list" |
| 18 | +TRAFFIC_GET_OP = "get" |
| 19 | +TRAFFIC_FORMAT = "pcap" |
| 20 | +PCAP_MEDIA_TYPE = "application/vnd.tcpdump.pcap" |
| 21 | +LIST_TIMEOUT_S = 30.0 |
| 22 | +FETCH_TIMEOUT = httpx.Timeout(600.0, connect=10.0) |
| 23 | + |
| 24 | + |
| 25 | +class TrafficCapture(BaseModel): |
| 26 | + """One capture file a ``list`` names: its ``name``, which ``get`` takes, the ``container`` whose packets it holds, |
| 27 | + its size in ``bytes`` and when it was last written. Fields it doesn't name are ignored, so a reader keeps reading |
| 28 | + an env that says more.""" |
| 29 | + |
| 30 | + model_config = ConfigDict(extra="ignore") |
| 31 | + |
| 32 | + name: str |
| 33 | + container: str |
| 34 | + bytes: int |
| 35 | + modified_at: datetime |
| 36 | + |
| 37 | + |
| 38 | +class TrafficCaptureListing(BaseModel): |
| 39 | + """``list``'s answer: the files' ``format`` and each ``TrafficCapture``, oldest first within a container, so its |
| 40 | + files read in the order their packets were captured.""" |
| 41 | + |
| 42 | + model_config = ConfigDict(extra="ignore") |
| 43 | + |
| 44 | + format: str = TRAFFIC_FORMAT |
| 45 | + captures: list[TrafficCapture] = [] |
| 46 | + |
| 47 | + |
| 48 | +def serves_traffic(card: dict) -> bool: |
| 49 | + """Whether ``card`` advertises traffic/v1's ``list``; no network.""" |
| 50 | + return isinstance(card, dict) and find_extension_method(card, TRAFFIC_EXTENSION_URI, TRAFFIC_LIST_OP) is not None |
| 51 | + |
| 52 | + |
| 53 | +async def list_traffic_captures(base_url: str, card: dict, *, timeout: float = LIST_TIMEOUT_S, verify: bool = True, |
| 54 | + headers: dict[str, str] | None = None) -> TrafficCaptureListing: |
| 55 | + """The captures the env whose ``card`` is read from ``base_url`` serves, through its ``list``. Raises |
| 56 | + ``ValueError`` when the card doesn't advertise it or the answer isn't a ``TrafficCaptureListing``, httpx's error |
| 57 | + when it can't be asked, and ``ExtensionCallError`` for a non-2xx answer.""" |
| 58 | + answer = await invoke_extension(base_url, card, TRAFFIC_EXTENSION_URI, timeout=timeout, verify=verify, |
| 59 | + method=TRAFFIC_LIST_OP, headers=headers) |
| 60 | + return TrafficCaptureListing.model_validate(answer) |
| 61 | + |
| 62 | + |
| 63 | +async def fetch_traffic_capture(base_url: str, card: dict, name: str, path: str, *, |
| 64 | + timeout: httpx.Timeout | float = FETCH_TIMEOUT, verify: bool = True, |
| 65 | + headers: dict[str, str] | None = None) -> int: |
| 66 | + """Stream the capture ``name``, as a ``list`` entry names it, from the env whose ``card`` is read from |
| 67 | + ``base_url`` through its ``get`` to the file ``path``, and return its size. Raises ``ValueError`` when the card |
| 68 | + doesn't advertise ``get``, httpx's error when it can't be asked, and ``ExtensionCallError`` for a non-2xx answer, |
| 69 | + with nothing written.""" |
| 70 | + op = find_extension_method(card, TRAFFIC_EXTENSION_URI, TRAFFIC_GET_OP) |
| 71 | + if op is None or not op["endpoint"]: |
| 72 | + raise ValueError(f"card does not advertise {TRAFFIC_EXTENSION_URI} {TRAFFIC_GET_OP}") |
| 73 | + url = f"{base_url}{op['endpoint']}" |
| 74 | + params = {"name": name} |
| 75 | + request = {"params": params} if op["method"] == "GET" else {"json": params} |
| 76 | + size = 0 |
| 77 | + async with httpx.AsyncClient(verify=verify, timeout=timeout) as client: |
| 78 | + async with client.stream(op["method"], url, headers=request_headers(url, headers), **request) as response: |
| 79 | + if response.status_code >= 400: |
| 80 | + await response.aread() |
| 81 | + raise ExtensionCallError(response) |
| 82 | + with open(path, "wb") as out: |
| 83 | + async for chunk in response.aiter_bytes(): |
| 84 | + out.write(chunk) |
| 85 | + size += len(chunk) |
| 86 | + return size |
0 commit comments