oscura 0.0.1__py3-none-any.whl → 0.1.1__py3-none-any.whl

oscura/analyzers/packet/parser.py

@@ -0,0 +1,327 @@
+"""Binary parsing utilities for packet analysis.
+
+This module provides fast binary parsing using struct.Struct
+pre-compilation and TLV record support.
+
+
+Example:
+    >>> from oscura.analyzers.packet.parser import BinaryParser, parse_tlv
+    >>> parser = BinaryParser(">HBB")  # big-endian: ushort, 2 ubytes
+    >>> values = parser.unpack(data)
+
+References:
+    Python struct module documentation
+"""
+
+from __future__ import annotations
+
+import struct
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+
+@dataclass
+class TLVRecord:
+    """Type-Length-Value record.
+
+    Attributes:
+        type_id: Record type identifier.
+        length: Data length in bytes.
+        value: Record data.
+        offset: Byte offset in source data.
+    """
+
+    type_id: int
+    length: int
+    value: bytes
+    offset: int = 0
+
+
+class BinaryParser:
+    """Fast binary parser using pre-compiled struct format.
+
+    Uses struct.Struct for pre-compiled format strings to achieve
+    high parsing throughput (>10K packets/second).
+
+    Args:
+        format_string: struct format string (e.g., ">HBB").
+
+    Attributes:
+        format: The format string.
+        size: Size of the packed structure in bytes.
+
+    Example:
+        >>> parser = BinaryParser(">HHBBI")  # Header format
+        >>> header = parser.unpack(data)
+        >>> print(f"Fields: {header}")
+
+        >>> # Unpack from offset
+        >>> payload = parser.unpack_from(data, offset=8)
+    """
+
+    def __init__(self, format_string: str) -> None:
+        """Initialize parser with format string.
+
+        Args:
+            format_string: struct format string. Use standard prefixes:
+                - ">": Big-endian
+                - "<": Little-endian
+                - "=": Native byte order
+                - "!": Network byte order (big-endian)
+        """
+        self._struct = struct.Struct(format_string)
+        self._format = format_string
+
+    @property
+    def format(self) -> str:
+        """Get format string."""
+        return self._format
+
+    @property
+    def size(self) -> int:
+        """Get packed size in bytes."""
+        return self._struct.size
+
+    def unpack(self, buffer: bytes) -> tuple[Any, ...]:
+        """Unpack data from bytes.
+
+        Args:
+            buffer: Bytes to unpack. Must be at least self.size bytes.
+
+        Returns:
+            Tuple of unpacked values.
+        """
+        return self._struct.unpack(buffer[: self.size])
+
+    def unpack_from(self, buffer: bytes, offset: int = 0) -> tuple[Any, ...]:
+        """Unpack data from offset in buffer.
+
+        More efficient than slicing when parsing multiple fields
+        from a single buffer.
+
+        Args:
+            buffer: Source buffer.
+            offset: Byte offset to start unpacking.
+
+        Returns:
+            Tuple of unpacked values.
+        """
+        return self._struct.unpack_from(buffer, offset)
+
+    def pack(self, *values: Any) -> bytes:
+        """Pack values to bytes.
+
+        Args:
+            *values: Values matching format string types.
+
+        Returns:
+            Packed bytes.
+        """
+        return self._struct.pack(*values)
+
+    def iter_unpack(self, buffer: bytes) -> Iterator[tuple[Any, ...]]:
+        """Iterate over repeated structures in buffer.
+
+        Args:
+            buffer: Source buffer.
+
+        Returns:
+            Iterator yielding tuples of unpacked values.
+
+        Example:
+            >>> parser = BinaryParser(">HH")
+            >>> for a, b in parser.iter_unpack(data):
+            ...     print(f"Pair: {a}, {b}")
+        """
+        return self._struct.iter_unpack(buffer)
+
+
+class PacketParser:
+    """Multi-field packet parser.
+
+    Parses packets with multiple fields using named field definitions.
+
+    Example:
+        >>> fields = [
+        ...     ("sync", "H"),
+        ...     ("length", "H"),
+        ...     ("type", "B"),
+        ...     ("flags", "B"),
+        ... ]
+        >>> parser = PacketParser(fields, byte_order=">")
+        >>> packet = parser.parse(data)
+        >>> print(f"Type: {packet['type']}")
+    """
+
+    def __init__(
+        self,
+        fields: list[tuple[str, str]],
+        byte_order: str = ">",
+    ) -> None:
+        """Initialize packet parser.
+
+        Args:
+            fields: List of (name, format_char) tuples.
+            byte_order: Byte order prefix (">", "<", "=", "!").
+        """
+        self._field_names = [f[0] for f in fields]
+        format_chars = "".join(f[1] for f in fields)
+        self._parser = BinaryParser(byte_order + format_chars)
+
+    @property
+    def size(self) -> int:
+        """Get packet header size in bytes."""
+        return self._parser.size
+
+    def parse(self, buffer: bytes, offset: int = 0) -> dict[str, Any]:
+        """Parse packet fields.
+
+        Args:
+            buffer: Source buffer.
+            offset: Byte offset.
+
+        Returns:
+            Dictionary mapping field names to values.
+        """
+        values = self._parser.unpack_from(buffer, offset)
+        return dict(zip(self._field_names, values, strict=False))
+
+    def pack(self, **fields: Any) -> bytes:
+        """Pack fields to bytes.
+
+        Args:
+            **fields: Field values by name.
+
+        Returns:
+            Packed bytes.
+        """
+        values = [fields[name] for name in self._field_names]
+        return self._parser.pack(*values)
+
+
+def parse_tlv(
+    buffer: bytes,
+    *,
+    type_size: int = 1,
+    length_size: int = 1,
+    big_endian: bool = True,
+    include_length_in_length: bool = False,
+    type_map: dict[int, str] | None = None,
+) -> list[TLVRecord]:
+    """Parse Type-Length-Value records.
+
+    Args:
+        buffer: Source buffer containing TLV records.
+        type_size: Size of type field in bytes (1, 2, or 4).
+        length_size: Size of length field in bytes (1, 2, or 4).
+        big_endian: True for big-endian byte order.
+        include_length_in_length: True if length includes type+length fields.
+        type_map: Optional mapping of type IDs to names.
+
+    Returns:
+        List of TLVRecord objects.
+
+    Example:
+        >>> records = parse_tlv(data, type_size=2, length_size=2)
+        >>> for rec in records:
+        ...     print(f"Type {rec.type_id}: {rec.length} bytes")
+    """
+    records: list[TLVRecord] = []
+    offset = 0
+    header_size = type_size + length_size
+
+    # Determine struct format
+    byte_order = ">" if big_endian else "<"
+    type_fmt = {1: "B", 2: "H", 4: "I"}[type_size]
+    length_fmt = {1: "B", 2: "H", 4: "I"}[length_size]
+    header_parser = BinaryParser(byte_order + type_fmt + length_fmt)
+
+    while offset + header_size <= len(buffer):
+        type_id, length = header_parser.unpack_from(buffer, offset)
+
+        # Adjust length if it includes header
+        data_length = length - header_size if include_length_in_length else length
+
+        if data_length < 0:
+            break
+
+        # Extract value
+        value_start = offset + header_size
+        value_end = value_start + data_length
+
+        if value_end > len(buffer):
+            break
+
+        value = buffer[value_start:value_end]
+
+        records.append(
+            TLVRecord(
+                type_id=type_id,
+                length=data_length,
+                value=value,
+                offset=offset,
+            )
+        )
+
+        offset = value_end
+
+    return records
+
+
+def parse_tlv_nested(
+    buffer: bytes,
+    *,
+    type_size: int = 1,
+    length_size: int = 1,
+    big_endian: bool = True,
+    container_types: set[int] | None = None,
+) -> dict[int, Any]:
+    """Parse nested TLV structure.
+
+    Args:
+        buffer: Source buffer.
+        type_size: Size of type field.
+        length_size: Size of length field.
+        big_endian: Byte order.
+        container_types: Set of type IDs that contain nested TLV.
+
+    Returns:
+        Dictionary with type_id keys and either bytes or nested dict values.
+    """
+    container_types = container_types or set()
+    result: dict[int, Any] = {}
+
+    records = parse_tlv(
+        buffer,
+        type_size=type_size,
+        length_size=length_size,
+        big_endian=big_endian,
+    )
+
+    for rec in records:
+        if rec.type_id in container_types:
+            # Recursively parse nested structure
+            nested = parse_tlv_nested(
+                rec.value,
+                type_size=type_size,
+                length_size=length_size,
+                big_endian=big_endian,
+                container_types=container_types,
+            )
+            result[rec.type_id] = nested
+        else:
+            result[rec.type_id] = rec.value
+
+    return result
+
+
+__all__ = [
+    "BinaryParser",
+    "PacketParser",
+    "TLVRecord",
+    "parse_tlv",
+    "parse_tlv_nested",
+]