MAI/udp_probe.py

import hashlib
import json
import socket
import time
from collections import Counter
from pathlib import Path

from udp_dump_capture import MikPacketAssembler


MAX_UDP_PAYLOAD = 65535
MAX_MIK_ARRAY = 256 * 1024 * 1024
MIK_PIXEL_BYTES = {
    0x01: 1,
    0x02: 2,
    0x03: 3,
    0x0A: 2,
    0x12: 2,
}
RAW_ENCODINGS = {
    "gray8": 1,
    "gray16": 2,
    "yuyv422": 2,
    "bgr24": 3,
    "rgb24": 3,
}


def inspect_mik_array(data):
    if len(data) < 12:
        return None
    label_count = int.from_bytes(data[:4], "little")
    video_offset = 4 + label_count * 40
    if label_count > 1_000_000 or video_offset + 8 > len(data):
        return None
    header = data[video_offset:video_offset + 8]
    width = int.from_bytes(header[:2], "little")
    height = int.from_bytes(header[2:4], "little")
    pixel_id = header[4]
    padding = header[6]
    bytes_per_pixel = MIK_PIXEL_BYTES.get(pixel_id)
    if not bytes_per_pixel or not (1 <= width <= 8192 and 1 <= height <= 8192):
        return None
    expected = video_offset + 8 + (width * bytes_per_pixel + padding) * height
    if expected > len(data):
        return None
    return {
        "width": width,
        "height": height,
        "pixel_id": pixel_id,
        "row_padding": padding,
        "labels": label_count,
        "array_bytes": len(data),
        "expected_bytes": expected,
    }


def _mik_candidate(payloads):
    assembler = MikPacketAssembler()
    header_count = 0
    starts = 0
    ends = 0
    arrays = []
    for payload in payloads:
        if len(payload) < 8:
            continue
        flags = payload[1]
        packet_number = payload[3]
        value = int.from_bytes(payload[4:8], "little")
        if flags & ~0x03:
            continue
        if flags & 0x02:
            if packet_number != 0 or value <= 0 or value > MAX_MIK_ARRAY:
                continue
            starts += 1
        elif value > MAX_MIK_ARRAY:
            continue
        header_count += 1
        ends += int(bool(flags & 0x01))
        try:
            array = assembler.push(payload)
        except ValueError:
            continue
        if array is not None:
            arrays.append(array)

    frames = [frame for frame in map(inspect_mik_array, arrays) if frame]
    if frames:
        frame = frames[0]
        return {
            "kind": "mik_video",
            "confidence": 100,
            "evidence": {
                "matching_headers": header_count,
                "start_packets": starts,
                "end_packets": ends,
                "complete_arrays": len(arrays),
                "valid_video_arrays": len(frames),
                "dropped_arrays": assembler.dropped_arrays,
            },
            "frame": frame,
            "recommended": {
                "source_mode": "udp_mik_live",
                "quality": f"{frame['width']}x{frame['height']}",
            },
        }
    ratio = header_count / max(1, len(payloads))
    if header_count >= 3 and ratio >= 0.7 and starts:
        return {
            "kind": "mik_fragments",
            "confidence": min(92, round(65 + ratio * 25)),
            "evidence": {
                "matching_headers": header_count,
                "start_packets": starts,
                "end_packets": ends,
                "complete_arrays": len(arrays),
                "dropped_arrays": assembler.dropped_arrays,
            },
            "recommended": {"source_mode": "udp_mik_live"},
        }
    return None


def _rtp_parts(payload):
    if len(payload) < 12 or payload[0] >> 6 != 2:
        return None
    cc = payload[0] & 0x0F
    offset = 12 + cc * 4
    if offset > len(payload):
        return None
    if payload[0] & 0x10:
        if offset + 4 > len(payload):
            return None
        words = int.from_bytes(payload[offset + 2:offset + 4], "big")
        offset += 4 + words * 4
        if offset > len(payload):
            return None
    return {
        "payload_type": payload[1] & 0x7F,
        "sequence": int.from_bytes(payload[2:4], "big"),
        "timestamp": int.from_bytes(payload[4:8], "big"),
        "ssrc": int.from_bytes(payload[8:12], "big"),
        "payload": payload[offset:],
    }


def _is_mpeg_ts(data):
    return len(data) >= 188 and len(data) % 188 == 0 and all(
        data[index] == 0x47 for index in range(0, len(data), 188)
    )


def _rtp_candidate(payloads):
    headers = [header for header in map(_rtp_parts, payloads) if header]
    if len(headers) < 2 or len(headers) / max(1, len(payloads)) < 0.8:
        return None
    ssrc, ssrc_count = Counter(item["ssrc"] for item in headers).most_common(1)[0]
    payload_type, type_count = Counter(item["payload_type"] for item in headers).most_common(1)[0]
    sequential = sum(
        ((current["sequence"] - previous["sequence"]) & 0xFFFF) == 1
        for previous, current in zip(headers, headers[1:])
    )
    ts_packets = sum(_is_mpeg_ts(item["payload"]) for item in headers)
    confidence = 75
    if ssrc_count / len(headers) >= 0.9 and type_count / len(headers) >= 0.9:
        confidence += 10
    if sequential / max(1, len(headers) - 1) >= 0.7:
        confidence += 10
    return {
        "kind": "rtp_mpeg_ts" if ts_packets else "rtp",
        "confidence": min(98, confidence),
        "evidence": {
            "rtp_packets": len(headers),
            "payload_type": payload_type,
            "ssrc": f"0x{ssrc:08x}",
            "sequential_pairs": sequential,
            "mpeg_ts_payloads": ts_packets,
        },
    }


def _encoded_candidate(payloads):
    ts_packets = sum(_is_mpeg_ts(payload) for payload in payloads)
    if ts_packets and ts_packets / len(payloads) >= 0.7:
        return {
            "kind": "mpeg_ts",
            "confidence": 99,
            "evidence": {"mpeg_ts_datagrams": ts_packets},
        }

    jpeg = sum(
        payload.startswith(b"\xff\xd8\xff") and payload.rstrip().endswith(b"\xff\xd9")
        for payload in payloads
    )
    png = sum(
        payload.startswith(b"\x89PNG\r\n\x1a\n") and b"IEND" in payload[-32:]
        for payload in payloads
    )
    if jpeg or png:
        kind = "jpeg" if jpeg >= png else "png"
        count = max(jpeg, png)
        return {
            "kind": kind,
            "confidence": 100,
            "evidence": {"complete_images": count},
            "recommended": {
                "source_mode": "udp_delimited_live",
                "frame_encoding": "auto",
            },
        }

    start_code_packets = 0
    h264_packets = 0
    h265_packets = 0
    for payload in payloads:
        offset = 4 if payload.startswith(b"\x00\x00\x00\x01") else 3
        if offset == 3 and not payload.startswith(b"\x00\x00\x01"):
            continue
        if len(payload) <= offset:
            continue
        start_code_packets += 1
        h264_packets += int(1 <= (payload[offset] & 0x1F) <= 23)
        h265_packets += int(((payload[offset] >> 1) & 0x3F) <= 40)
    if start_code_packets:
        kind = "h264_annex_b" if h264_packets >= h265_packets else "h265_annex_b"
        return {
            "kind": kind,
            "confidence": 90,
            "evidence": {"start_code_datagrams": start_code_packets},
        }
    return None


def _raw_candidate(payloads, width, height, separator, configured_encoding):
    if width <= 0 or height <= 0:
        return None
    separator_payload = bytes((separator & 0xFF,))
    groups = []
    current = 0
    separator_packets = 0
    for payload in payloads:
        if payload == separator_payload:
            separator_packets += 1
            if current:
                groups.append(current)
                current = 0
        else:
            current += len(payload)
    if current:
        groups.append(current)

    encodings = (
        {configured_encoding: RAW_ENCODINGS[configured_encoding]}
        if configured_encoding in RAW_ENCODINGS
        else RAW_ENCODINGS
    )
    matches = []
    total = sum(len(payload) for payload in payloads if payload != separator_payload)
    for encoding, bytes_per_pixel in encodings.items():
        expected = width * height * bytes_per_pixel
        exact_groups = sum(size == expected for size in groups)
        complete_frames = total // expected
        remainder = total % expected
        if exact_groups:
            confidence = 99
        elif separator_packets and complete_frames and remainder <= max(map(len, payloads)):
            confidence = 78
        elif not separator_packets and complete_frames:
            confidence = 55
        else:
            continue
        matches.append((confidence, encoding, expected, exact_groups, complete_frames, remainder))
    if not matches:
        return None

    matches.sort(reverse=True)
    confidence, encoding, expected, exact_groups, complete_frames, remainder = matches[0]
    same_size = sorted({
        candidate_encoding
        for _, candidate_encoding, candidate_size, *_ in matches
        if candidate_size == expected
    })
    ambiguous = len(same_size) > 1 and configured_encoding not in RAW_ENCODINGS
    recommended = {
        "source_mode": "udp_delimited_live",
        "quality": f"{width}x{height}",
        "separator_byte": separator,
    }
    if not ambiguous:
        recommended["frame_encoding"] = encoding
    return {
        "kind": "raw_delimited" if separator_packets else "raw_stream",
        "confidence": confidence,
        "evidence": {
            "separator_packets": separator_packets,
            "expected_frame_bytes": expected,
            "exact_frame_groups": exact_groups,
            "complete_frame_equivalents": complete_frames,
            "trailing_bytes": remainder,
            "possible_encodings": same_size,
        },
        "frame": {"width": width, "height": height, "encoding": encoding},
        "recommended": recommended,
    }


def _text_candidate(payloads):
    if not payloads:
        return None
    sample = payloads[0][:8192]
    try:
        text = sample.decode("utf-8")
    except UnicodeDecodeError:
        return None
    printable = sum(character.isprintable() or character in "\r\n\t" for character in text)
    if not text or printable / len(text) < 0.9:
        return None
    try:
        json.loads(text)
        kind = "json"
        confidence = 100
    except json.JSONDecodeError:
        kind = "text"
        confidence = 85
    return {
        "kind": kind,
        "confidence": confidence,
        "evidence": {"preview": text[:160]},
    }


def analyze_udp_records(records, width=0, height=0, separator=0, frame_encoding="auto"):
    payloads = [record["payload"] for record in records]
    candidates = [
        candidate
        for candidate in (
            _mik_candidate(payloads),
            _rtp_candidate(payloads),
            _encoded_candidate(payloads),
            _raw_candidate(payloads, int(width), int(height), int(separator), frame_encoding),
            _text_candidate(payloads),
        )
        if candidate is not None
    ]
    candidates.sort(key=lambda candidate: candidate["confidence"], reverse=True)
    detected = candidates[0] if candidates else {
        "kind": "unknown",
        "confidence": 0,
        "evidence": {"reason": "no known structure matched"},
    }
    return {"detected": detected, "candidates": candidates}


def capture_udp_records(host, port, duration=3.0, max_packets=4096, max_bytes=32 * 1024 * 1024):
    host = str(host or "0.0.0.0").strip() or "0.0.0.0"
    port = max(1, min(65535, int(port)))
    duration = max(0.2, min(15.0, float(duration)))
    max_packets = max(1, min(32768, int(max_packets)))
    max_bytes = max(MAX_UDP_PAYLOAD, min(256 * 1024 * 1024, int(max_bytes)))
    records = []
    total = 0
    truncated = False
    started = time.monotonic()
    deadline = started + duration
    sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    try:
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, 16 * 1024 * 1024)
        sock.bind((host, port))
        bound_host, bound_port = sock.getsockname()
        while len(records) < max_packets:
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                break
            sock.settimeout(min(0.25, remaining))
            try:
                payload, address = sock.recvfrom(MAX_UDP_PAYLOAD)
            except socket.timeout:
                continue
            if total + len(payload) > max_bytes:
                truncated = True
                break
            records.append({
                "timestamp_ns": time.time_ns(),
                "address": (str(address[0]), int(address[1])),
                "payload": payload,
            })
            total += len(payload)
        truncated = truncated or len(records) >= max_packets
    finally:
        sock.close()
    return records, {
        "listen_host": bound_host,
        "listen_port": bound_port,
        "elapsed_sec": round(time.monotonic() - started, 3),
        "truncated": truncated,
    }


def save_udp_records(records, destination_port, output_dir):
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    stamp = time.strftime("%Y%m%d_%H%M%S")
    suffix = f"{time.time_ns() % 1_000_000_000:09d}"
    dump_path = output_dir / f"udp_probe_{stamp}_{suffix}.udp"
    report_path = dump_path.with_suffix(".json")
    digest = hashlib.sha256()
    packet_meta = []
    offset = 0
    with dump_path.open("wb") as stream:
        for index, record in enumerate(records):
            payload = record["payload"]
            envelope = int(destination_port).to_bytes(2, "little") + len(payload).to_bytes(2, "little")
            stream.write(envelope)
            stream.write(payload)
            digest.update(envelope)
            digest.update(payload)
            packet_meta.append({
                "index": index,
                "timestamp_ns": record["timestamp_ns"],
                "source_ip": record["address"][0],
                "source_port": record["address"][1],
                "payload_size": len(payload),
                "payload_sha256": hashlib.sha256(payload).hexdigest(),
                "dump_offset": offset,
            })
            offset += len(envelope) + len(payload)
    report = {
        "format": "uint16_le destination_port, uint16_le payload_size, payload bytes",
        "destination_port": int(destination_port),
        "packets": packet_meta,
        "dump_name": dump_path.name,
        "dump_size": dump_path.stat().st_size,
        "dump_sha256": digest.hexdigest(),
    }
    report_path.write_text(json.dumps(report, indent=2), encoding="utf-8")
    return dump_path, report_path, report


def run_udp_probe(
    host,
    port,
    output_dir,
    width=0,
    height=0,
    separator=0,
    frame_encoding="auto",
    duration=3.0,
):
    records, capture = capture_udp_records(host, port, duration=duration)
    analysis = analyze_udp_records(records, width, height, separator, frame_encoding)
    sizes = Counter(len(record["payload"]) for record in records)
    sources = Counter(f"{record['address'][0]}:{record['address'][1]}" for record in records)
    if len(records) > 1:
        span = (records[-1]["timestamp_ns"] - records[0]["timestamp_ns"]) / 1e9
        packet_rate = (len(records) - 1) / max(span, 1e-9)
    else:
        packet_rate = 0.0

    sample_indexes = sorted(set(
        list(range(min(3, len(records))))
        + ([len(records) - 1] if records else [])
    ))
    samples = []
    for index in sample_indexes:
        payload = records[index]["payload"]
        samples.append({
            "index": index,
            "source": f"{records[index]['address'][0]}:{records[index]['address'][1]}",
            "size": len(payload),
            "sha256": hashlib.sha256(payload).hexdigest(),
            "head_hex": payload[:64].hex(" "),
            "tail_hex": payload[-32:].hex(" ") if len(payload) > 64 else "",
            "ascii": "".join(chr(byte) if 32 <= byte < 127 else "." for byte in payload[:64]),
        })

    exact_capture = None
    if records:
        dump_path, report_path, report = save_udp_records(records, port, output_dir)
        exact_capture = {
            "dump_name": dump_path.name,
            "report_name": report_path.name,
            "bytes": report["dump_size"],
            "sha256": report["dump_sha256"],
        }
    return {
        **capture,
        "packets": len(records),
        "payload_bytes": sum(len(record["payload"]) for record in records),
        "packets_per_sec": round(packet_rate, 1),
        "sources": [{"address": address, "packets": count} for address, count in sources.most_common()],
        "sizes": [{"bytes": size, "packets": count} for size, count in sizes.most_common(12)],
        "samples": samples,
        "exact_capture": exact_capture,
        **analysis,
    }