MAI/configurable_udp_capture.py

import socket

import cv2

from delimited_frame_capture import decode_frame_data, raw_frame_size
from udp_dump_capture import UdpDumpCapture


DEFAULT_PACKET_SCHEMA = {
    "assembly": "fragmented",
    "payload_format": "frame",
    "header_size": 8,
    "byte_order": "little",
    "flags_offset": 1,
    "flags_size": 1,
    "start_mask": 0x02,
    "end_mask": 0x01,
    "sequence_offset": 2,
    "sequence_size": 1,
    "packet_number_offset": 3,
    "packet_number_size": 1,
    "value_offset": 4,
    "value_size": 4,
    "value_mode": "total_then_offset",
    "read_fields": [],
}

ASSEMBLIES = {"fragmented", "datagram", "stream"}
PAYLOAD_FORMATS = {"frame", "mik"}
BYTE_ORDERS = {"little", "big"}
VALUE_MODES = {"total_then_offset", "total_size", "offset", "unused"}
PACKET_FIELD_ROLES = {"skip", "field", "flags", "sequence", "packet_number", "value"}
DEFAULT_PACKET_LAYOUT = [
    {"role": "skip", "size": 1, "label": "Version"},
    {"role": "flags", "size": 1, "label": "Flags"},
    {"role": "sequence", "size": 1, "label": "Sequence"},
    {"role": "packet_number", "size": 1, "label": "Packet"},
    {"role": "value", "size": 4, "label": "Value"},
]
MAX_ARRAY_SIZE = 256 * 1024 * 1024


def _bounded_int(value, default, minimum, maximum):
    try:
        number = int(str(value).strip(), 0)
    except (TypeError, ValueError):
        number = int(default)
    return max(minimum, min(maximum, number))


def normalize_packet_schema(value=None):
    schema = DEFAULT_PACKET_SCHEMA.copy()
    if isinstance(value, dict):
        schema.update({key: item for key, item in value.items() if key in schema})

    schema["assembly"] = str(schema["assembly"]).lower()
    if schema["assembly"] not in ASSEMBLIES:
        schema["assembly"] = DEFAULT_PACKET_SCHEMA["assembly"]
    schema["payload_format"] = str(schema["payload_format"]).lower()
    if schema["payload_format"] not in PAYLOAD_FORMATS:
        schema["payload_format"] = DEFAULT_PACKET_SCHEMA["payload_format"]
    schema["byte_order"] = str(schema["byte_order"]).lower()
    if schema["byte_order"] not in BYTE_ORDERS:
        schema["byte_order"] = DEFAULT_PACKET_SCHEMA["byte_order"]
    schema["value_mode"] = str(schema["value_mode"]).lower()
    if schema["value_mode"] not in VALUE_MODES:
        schema["value_mode"] = DEFAULT_PACKET_SCHEMA["value_mode"]

    schema["header_size"] = _bounded_int(schema["header_size"], 8, 0, 1024)
    for name in ("flags_offset", "sequence_offset", "packet_number_offset", "value_offset"):
        schema[name] = _bounded_int(schema[name], -1, -1, 1023)
    for name in ("flags_size", "sequence_size", "packet_number_size", "value_size"):
        schema[name] = _bounded_int(schema[name], 1, 1, 8)
    for name in ("start_mask", "end_mask"):
        maximum = (1 << (8 * schema["flags_size"])) - 1
        schema[name] = _bounded_int(schema[name], DEFAULT_PACKET_SCHEMA[name], 0, maximum)

    fields = (
        ("flags_offset", schema["flags_size"]),
        ("sequence_offset", schema["sequence_size"]),
        ("packet_number_offset", schema["packet_number_size"]),
        ("value_offset", schema["value_size"]),
    )
    for offset_name, size in fields:
        offset = schema[offset_name]
        if offset >= 0 and offset + size > schema["header_size"]:
            schema[offset_name] = -1
    read_fields = []
    for index, field in enumerate(schema.get("read_fields") if isinstance(schema.get("read_fields"), list) else []):
        if not isinstance(field, dict):
            continue
        offset = _bounded_int(field.get("offset"), -1, -1, 1023)
        size = _bounded_int(field.get("size"), 1, 1, 8)
        if offset < 0 or offset + size > schema["header_size"]:
            continue
        name = str(field.get("name") or f"field_{index}").strip()[:64] or f"field_{index}"
        read_fields.append({"name": name, "offset": offset, "size": size})
    schema["read_fields"] = read_fields[:64]
    return schema


def normalize_packet_layout(value=None):
    source = value if isinstance(value, list) else DEFAULT_PACKET_LAYOUT
    layout = []
    remaining = 1024
    for item in source[:128]:
        if not isinstance(item, dict) or remaining <= 0:
            continue
        role = str(item.get("role") or "skip").strip().lower()
        if role not in PACKET_FIELD_ROLES:
            role = "skip"
        maximum = min(remaining, 8 if role != "skip" else 1024)
        size = _bounded_int(item.get("size"), 1, 1, maximum)
        label = str(item.get("label") or "").strip()[:64]
        layout.append({"role": role, "size": size, "label": label})
        remaining -= size
    return layout


def packet_schema_from_layout(layout, base=None):
    fields = normalize_packet_layout(layout)
    schema = normalize_packet_schema(base)
    for name in ("flags_offset", "sequence_offset", "packet_number_offset", "value_offset"):
        schema[name] = -1
    schema["read_fields"] = []

    role_fields = {
        "flags": ("flags_offset", "flags_size"),
        "sequence": ("sequence_offset", "sequence_size"),
        "packet_number": ("packet_number_offset", "packet_number_size"),
        "value": ("value_offset", "value_size"),
    }
    used = set()
    offset = 0
    for field in fields:
        role = field["role"]
        if role in role_fields and role not in used:
            offset_name, size_name = role_fields[role]
            schema[offset_name] = offset
            schema[size_name] = field["size"]
            used.add(role)
        if role == "field":
            schema["read_fields"].append({
                "name": field["label"] or f"field_{offset}",
                "offset": offset,
                "size": field["size"],
            })
        offset += field["size"]
    schema["header_size"] = offset
    return normalize_packet_schema(schema)


def packet_layout_from_schema(value=None):
    schema = normalize_packet_schema(value)
    fields = []
    for role, offset_name, size_name in (
        ("flags", "flags_offset", "flags_size"),
        ("sequence", "sequence_offset", "sequence_size"),
        ("packet_number", "packet_number_offset", "packet_number_size"),
        ("value", "value_offset", "value_size"),
    ):
        offset = schema[offset_name]
        if offset >= 0:
            label = {
                "flags": "Flags",
                "sequence": "Sequence",
                "packet_number": "Packet",
                "value": "Value",
            }[role]
            fields.append((offset, schema[size_name], role, label))
    for field in schema["read_fields"]:
        fields.append((field["offset"], field["size"], "field", field["name"]))
    fields.sort()

    layout = []
    cursor = 0
    for offset, size, role, label in fields:
        if offset < cursor:
            continue
        if offset > cursor:
            layout.append({"role": "skip", "size": offset - cursor, "label": ""})
        layout.append({"role": role, "size": size, "label": label})
        cursor = offset + size
    if cursor < schema["header_size"]:
        layout.append({"role": "skip", "size": schema["header_size"] - cursor, "label": ""})
    return normalize_packet_layout(layout)


class ConfigurablePacketAssembler:
    def __init__(self, schema=None):
        self.schema = normalize_packet_schema(schema)
        self.current = None
        self.expected_packet = None
        self.expected_total = None
        self.dropped_arrays = 0
        self.last_fields = {}

    def _field(self, payload, offset_name, size_name=None):
        offset = self.schema[offset_name]
        if offset < 0:
            return None
        size = self.schema[size_name] if size_name else 1
        if offset + size > len(payload):
            raise ValueError(f"packet too short for {offset_name}")
        return int.from_bytes(payload[offset:offset + size], self.schema["byte_order"])

    def _drop(self):
        if self.current is not None:
            self.dropped_arrays += 1
        self.current = None
        self.expected_packet = None
        self.expected_total = None

    def _finish(self):
        result = bytes(self.current["data"])
        self.current = None
        self.expected_packet = None
        self.expected_total = None
        return result

    def push(self, payload):
        schema = self.schema
        if len(payload) < schema["header_size"]:
            self._drop()
            raise ValueError("UDP payload is shorter than configured header")
        self.last_fields = {
            field["name"]: int.from_bytes(
                payload[field["offset"]:field["offset"] + field["size"]],
                schema["byte_order"],
            )
            for field in schema["read_fields"]
        }
        packet_data = payload[schema["header_size"]:]
        if schema["assembly"] == "datagram":
            return packet_data
        if schema["assembly"] != "fragmented":
            return None

        flags = self._field(payload, "flags_offset", "flags_size") or 0
        sequence = self._field(payload, "sequence_offset", "sequence_size")
        packet_number = self._field(payload, "packet_number_offset", "packet_number_size")
        value = self._field(payload, "value_offset", "value_size")
        is_start = bool(flags & schema["start_mask"]) if schema["start_mask"] else self.current is None
        is_end = bool(flags & schema["end_mask"]) if schema["end_mask"] else False

        if is_start:
            if self.current is not None:
                self._drop()
            total = value if schema["value_mode"] in {"total_then_offset", "total_size"} else None
            if total is not None and (total <= 0 or total > MAX_ARRAY_SIZE):
                return None
            if len(packet_data) > MAX_ARRAY_SIZE or (total is not None and len(packet_data) > total):
                return None
            self.current = {"sequence": sequence, "data": bytearray(packet_data)}
            self.expected_total = total
            if packet_number is not None:
                modulo = 1 << (8 * schema["packet_number_size"])
                self.expected_packet = (packet_number + 1) % modulo
        elif self.current is None:
            return None
        else:
            if sequence is not None and self.current["sequence"] is not None and sequence != self.current["sequence"]:
                self._drop()
                return None
            if packet_number is not None and self.expected_packet is not None and packet_number != self.expected_packet:
                self._drop()
                return None
            if value is not None:
                if schema["value_mode"] in {"total_then_offset", "offset"} and value != len(self.current["data"]):
                    self._drop()
                    return None
                if schema["value_mode"] == "total_size" and value != self.expected_total:
                    self._drop()
                    return None
            self.current["data"].extend(packet_data)
            if packet_number is not None:
                modulo = 1 << (8 * schema["packet_number_size"])
                self.expected_packet = (packet_number + 1) % modulo

        size = len(self.current["data"])
        if size > MAX_ARRAY_SIZE or (self.expected_total is not None and size > self.expected_total):
            self._drop()
            return None
        if is_end:
            if self.expected_total is not None and size != self.expected_total:
                self._drop()
                return None
            return self._finish()
        if not schema["end_mask"] and self.expected_total is not None and size == self.expected_total:
            return self._finish()
        return None


class ConfigurableUdpCapture(UdpDumpCapture):
    """VideoCapture-compatible live UDP reader driven by a UI packet schema."""

    MAX_FRAME_SIZE = 64 * 1024 * 1024

    def __init__(
        self,
        host="0.0.0.0",
        port=59005,
        fps=30.0,
        width=1280,
        height=720,
        encoding="auto",
        separator=0,
        schema=None,
    ):
        self.path = None
        self.host = str(host)
        self.port = int(port)
        self.encoding = str(encoding or "auto").lower()
        self.separator = bytes((int(separator) & 0xFF,))
        self.schema = normalize_packet_schema(schema)
        self._init_decoder(fps, width, height)
        self._file = None
        self._next_frame = None
        self._assembler = ConfigurablePacketAssembler(self.schema)
        self._socket = None
        self._stream_buffer = bytearray()
        self._frames = []
        self.last_packet_fields = {}
        try:
            self._socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            self._socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            self._socket.bind((self.host, self.port))
            self.port = int(self._socket.getsockname()[1])
            self._socket.settimeout(0.5)
        except OSError as exc:
            self.last_error = str(exc)
            self.release()

    def isOpened(self):
        return self._socket is not None

    def _stream_frames(self, payload):
        self.last_packet_fields = {
            field["name"]: int.from_bytes(
                payload[field["offset"]:field["offset"] + field["size"]],
                self.schema["byte_order"],
            )
            for field in self.schema["read_fields"]
        }
        data = payload[self.schema["header_size"]:]
        expected = raw_frame_size(self.encoding, self.width, self.height)
        if data == self.separator:
            if self._stream_buffer and not expected:
                self._frames.append(bytes(self._stream_buffer))
                self._stream_buffer.clear()
            elif self._stream_buffer and expected:
                self._stream_buffer.clear()
                self.dropped_arrays += 1
            return
        self._stream_buffer.extend(data)
        if expected:
            while len(self._stream_buffer) >= expected:
                self._frames.append(bytes(self._stream_buffer[:expected]))
                del self._stream_buffer[:expected]
        if len(self._stream_buffer) > self.MAX_FRAME_SIZE:
            self._stream_buffer.clear()
            self.dropped_arrays += 1

    def _decode(self, data):
        if self.schema["payload_format"] == "mik":
            return self._decode_array(data)
        frame = decode_frame_data(data, self.encoding, self.width, self.height)
        self.height, self.width = frame.shape[:2]
        return frame

    def read(self):
        while self._socket is not None:
            if not self._frames:
                try:
                    payload, _address = self._socket.recvfrom(65535)
                except socket.timeout:
                    continue
                except OSError as exc:
                    self.last_error = str(exc)
                    return False, None
                try:
                    if self.schema["assembly"] == "stream":
                        if len(payload) < self.schema["header_size"]:
                            raise ValueError("UDP payload is shorter than configured header")
                        self._stream_frames(payload)
                    else:
                        dropped_before = self._assembler.dropped_arrays
                        data = self._assembler.push(payload)
                        self.last_packet_fields = self._assembler.last_fields.copy()
                        self.dropped_arrays += self._assembler.dropped_arrays - dropped_before
                        if data is not None:
                            self._frames.append(data)
                except ValueError as exc:
                    self.last_error = str(exc)
                    dropped_before = self._assembler.dropped_arrays
                    self._assembler._drop()
                    dropped = self._assembler.dropped_arrays - dropped_before
                    self.dropped_arrays += max(1, dropped)
                continue

            data = self._frames.pop(0)
            try:
                frame = self._decode(data)
            except ValueError as exc:
                self.last_error = str(exc)
                self.dropped_arrays += 1
                continue
            self.frames_read += 1
            return True, frame
        return False, None

    def release(self):
        sock, self._socket = getattr(self, "_socket", None), None
        if sock is not None:
            sock.close()