codekingpro/portable-devtools

1import io2import typing3from collections.abc import Iterable4from pathlib import Path5from typing import Any6 7from ruamel.yaml import YAML8 9from .. import ctx10from .. import http11from ..dns import DNSMessage12from ..flow import Flow13from ..tcp import TCPMessage14from ..udp import UDPMessage15from ..utils import strutils16from ..websocket import WebSocketMessage17from ._api import Metadata18 19type ContentviewMessage = (20    http.Message | TCPMessage | UDPMessage | WebSocketMessage | DNSMessage21)22 23 24def make_metadata(25    message: ContentviewMessage,26    flow: Flow,27) -> Metadata:28    metadata = Metadata(29        flow=flow,30        protobuf_definitions=Path(ctx.options.protobuf_definitions).expanduser()31        if ctx.options.protobuf_definitions32        else None,33    )34 35    match message:36        case http.Message():37            metadata.http_message = message38            if ctype := message.headers.get("content-type"):39                if ct := http.parse_content_type(ctype):40                    metadata.content_type = f"{ct[0]}/{ct[1]}"41        case TCPMessage():42            metadata.tcp_message = message43        case UDPMessage():44            metadata.udp_message = message45        case WebSocketMessage():46            metadata.websocket_message = message47        case DNSMessage():48            metadata.dns_message = message49        case other:  # pragma: no cover50            typing.assert_never(other)51 52    return metadata53 54 55def get_data(56    message: ContentviewMessage,57) -> tuple[bytes | None, str]:58    content: bytes | None59    try:60        content = message.content61    except ValueError:62        assert isinstance(message, http.Message)63        content = message.raw_content64        enc = "[cannot decode]"65    else:66        if isinstance(message, http.Message) and content != message.raw_content:67            enc = "[decoded {}]".format(message.headers.get("content-encoding"))68        else:69            enc = ""70 71    return content, enc72 73 74def yaml_dumps(d: Any) -> str:75    if not d:76        return ""77    out = io.StringIO()78    YAML(typ="rt", pure=True).dump(d, out)79    return out.getvalue()80 81 82def yaml_loads(yaml: str) -> Any:83    return YAML(typ="safe", pure=True).load(yaml)84 85 86def merge_repeated_keys(items: Iterable[tuple[str, str]]) -> dict[str, str | list[str]]:87    """88    Helper function that takes a list of pairs and merges repeated keys.89    """90    ret: dict[str, str | list[str]] = {}91    for key, value in items:92        if existing := ret.get(key):93            if isinstance(existing, list):94                existing.append(value)95            else:96                ret[key] = [existing, value]97        else:98            ret[key] = value99    return ret100 101 102def byte_pairs_to_str_pairs(103    items: Iterable[tuple[bytes, bytes]],104) -> Iterable[tuple[str, str]]:105    for key, value in items:106        yield (strutils.bytes_to_escaped_str(key), strutils.bytes_to_escaped_str(value))107