oscura 0.5.0__py3-none-any.whl → 0.6.0__py3-none-any.whl

oscura/loaders/pcap.py

@@ -27,7 +27,7 @@ if TYPE_CHECKING:
 
 # Try to import dpkt for full PCAP support
 try:
-    import dpkt  # type: ignore[import-not-found]
+    import dpkt
 
     DPKT_AVAILABLE = True
 except ImportError:
@@ -160,6 +160,135 @@ def load_pcap(
         )
 
 
+def _create_pcap_reader(f: Any, path: Path) -> Any:
+    """Create appropriate PCAP reader based on file format.
+
+    Args:
+        f: File handle.
+        path: Path to PCAP file.
+
+    Returns:
+        dpkt PCAP or PCAPNG reader.
+
+    Raises:
+        LoaderError: If PCAPNG support is unavailable.
+    """
+    magic = f.read(4)
+    f.seek(0)
+    magic_int = struct.unpack("<I", magic)[0]
+
+    if magic_int == PCAPNG_MAGIC:
+        try:
+            return dpkt.pcapng.Reader(f)
+        except AttributeError:
+            raise LoaderError(
+                "PCAPNG support requires newer dpkt version",
+                file_path=str(path),
+                fix_hint="Install dpkt >= 1.9: pip install dpkt>=1.9",
+            )
+    else:
+        return dpkt.pcap.Reader(f)
+
+
+def _parse_transport_layer(ip: Any, annotations: dict[str, Any]) -> str:
+    """Parse TCP/UDP/ICMP transport layer from IP packet.
+
+    Args:
+        ip: dpkt IP object.
+        annotations: Annotations dictionary to populate.
+
+    Returns:
+        Protocol name ("TCP", "UDP", "ICMP", or "IP").
+    """
+    if isinstance(ip.data, dpkt.tcp.TCP):
+        tcp = ip.data
+        annotations["src_port"] = tcp.sport
+        annotations["dst_port"] = tcp.dport
+        annotations["layer4_protocol"] = "TCP"
+        annotations["tcp_flags"] = tcp.flags
+        return "TCP"
+
+    elif isinstance(ip.data, dpkt.udp.UDP):
+        udp = ip.data
+        annotations["src_port"] = udp.sport
+        annotations["dst_port"] = udp.dport
+        annotations["layer4_protocol"] = "UDP"
+        return "UDP"
+
+    elif isinstance(ip.data, dpkt.icmp.ICMP):
+        annotations["layer4_protocol"] = "ICMP"
+        return "ICMP"
+
+    return "IP"
+
+
+def _parse_ethernet_frame(raw_data: bytes, link_type: int) -> tuple[str, dict[str, Any]]:
+    """Parse Ethernet frame and extract protocol information.
+
+    Args:
+        raw_data: Raw packet bytes.
+        link_type: Link layer type.
+
+    Returns:
+        Tuple of (protocol_name, annotations_dict).
+    """
+    annotations: dict[str, Any] = {}
+    protocol = "RAW"
+
+    try:
+        if link_type != 1:  # Not Ethernet
+            return protocol, annotations
+
+        eth = dpkt.ethernet.Ethernet(raw_data)
+        annotations["src_mac"] = _format_mac(eth.src)
+        annotations["dst_mac"] = _format_mac(eth.dst)
+
+        # Parse network layer
+        if isinstance(eth.data, dpkt.ip.IP):
+            ip = eth.data
+            annotations["src_ip"] = _format_ip(ip.src)
+            annotations["dst_ip"] = _format_ip(ip.dst)
+            annotations["layer3_protocol"] = "IP"
+            protocol = _parse_transport_layer(ip, annotations)
+
+        elif isinstance(eth.data, dpkt.ip6.IP6):
+            protocol = "IPv6"
+            annotations["layer3_protocol"] = "IPv6"
+
+        elif isinstance(eth.data, dpkt.arp.ARP):
+            protocol = "ARP"
+            annotations["layer3_protocol"] = "ARP"
+
+    except Exception:
+        # If parsing fails, return defaults
+        pass
+
+    return protocol, annotations
+
+
+def _matches_protocol_filter(
+    protocol: str, annotations: dict[str, Any], protocol_filter: str | None
+) -> bool:
+    """Check if packet matches protocol filter.
+
+    Args:
+        protocol: Packet protocol name.
+        annotations: Packet annotations.
+        protocol_filter: Filter string.
+
+    Returns:
+        True if packet matches filter (or no filter set).
+    """
+    if protocol_filter is None:
+        return True
+
+    return (
+        annotations.get("layer3_protocol") == protocol_filter
+        or annotations.get("layer4_protocol") == protocol_filter
+        or protocol == protocol_filter
+    )
+
+
 def _load_with_dpkt(
     path: Path,
     *,
@@ -180,27 +309,8 @@ def _load_with_dpkt(
         LoaderError: If file cannot be read or dpkt version is incompatible.
     """
     try:
-        with open(path, "rb") as f:
-            # Detect file format
-            magic = f.read(4)
-            f.seek(0)
-
-            magic_int = struct.unpack("<I", magic)[0]
-
-            if magic_int == PCAPNG_MAGIC:
-                # PCAPNG format
-                try:
-                    pcap_reader = dpkt.pcapng.Reader(f)
-                except AttributeError:
-                    raise LoaderError(  # noqa: B904
-                        "PCAPNG support requires newer dpkt version",
-                        file_path=str(path),
-                        fix_hint="Install dpkt >= 1.9: pip install dpkt>=1.9",
-                    )
-            else:
-                # Standard PCAP format
-                pcap_reader = dpkt.pcap.Reader(f)
-
+        with open(path, "rb", buffering=65536) as f:
+            pcap_reader = _create_pcap_reader(f, path)
             packets: list[ProtocolPacket] = []
             link_type = getattr(pcap_reader, "datalink", lambda: 1)()
 
@@ -208,62 +318,9 @@ def _load_with_dpkt(
                 if max_packets is not None and len(packets) >= max_packets:
                     break
 
-                # Parse Ethernet frame
-                annotations: dict[str, Any] = {}
-                protocol = "RAW"
-
-                try:
-                    if link_type == 1:  # Ethernet
-                        eth = dpkt.ethernet.Ethernet(raw_data)
-                        annotations["src_mac"] = _format_mac(eth.src)
-                        annotations["dst_mac"] = _format_mac(eth.dst)
-
-                        # Parse IP layer
-                        if isinstance(eth.data, dpkt.ip.IP):
-                            ip = eth.data
-                            protocol = "IP"
-                            annotations["src_ip"] = _format_ip(ip.src)
-                            annotations["dst_ip"] = _format_ip(ip.dst)
-                            annotations["layer3_protocol"] = "IP"
-
-                            # Parse transport layer
-                            if isinstance(ip.data, dpkt.tcp.TCP):
-                                tcp = ip.data
-                                protocol = "TCP"
-                                annotations["src_port"] = tcp.sport
-                                annotations["dst_port"] = tcp.dport
-                                annotations["layer4_protocol"] = "TCP"
-                                annotations["tcp_flags"] = tcp.flags
-
-                            elif isinstance(ip.data, dpkt.udp.UDP):
-                                udp = ip.data
-                                protocol = "UDP"
-                                annotations["src_port"] = udp.sport
-                                annotations["dst_port"] = udp.dport
-                                annotations["layer4_protocol"] = "UDP"
-
-                            elif isinstance(ip.data, dpkt.icmp.ICMP):
-                                protocol = "ICMP"
-                                annotations["layer4_protocol"] = "ICMP"
-
-                        elif isinstance(eth.data, dpkt.ip6.IP6):
-                            protocol = "IPv6"
-                            annotations["layer3_protocol"] = "IPv6"
-
-                        elif isinstance(eth.data, dpkt.arp.ARP):
-                            protocol = "ARP"
-                            annotations["layer3_protocol"] = "ARP"
-
-                except Exception:
-                    # If parsing fails, store raw data
-                    pass
-
-                # Apply protocol filter
-                if protocol_filter is not None and (
-                    annotations.get("layer3_protocol") != protocol_filter
-                    and annotations.get("layer4_protocol") != protocol_filter
-                    and protocol != protocol_filter
-                ):
+                protocol, annotations = _parse_ethernet_frame(raw_data, link_type)
+
+                if not _matches_protocol_filter(protocol, annotations, protocol_filter):
                     continue
 
                 packet = ProtocolPacket(
@@ -312,75 +369,9 @@ def _load_basic(
         LoaderError: If file cannot be read.
     """
     try:
-        with open(path, "rb") as f:
-            # Read global header (24 bytes)
-            header = f.read(24)
-            if len(header) < 24:
-                raise FormatError(
-                    "File too small to be a valid PCAP",
-                    file_path=str(path),
-                    expected="At least 24 bytes",
-                    got=f"{len(header)} bytes",
-                )
-
-            # Parse magic number
-            magic = struct.unpack("<I", header[:4])[0]
-
-            if magic in (PCAP_MAGIC_LE, PCAP_MAGIC_NS_LE):
-                byte_order = "<"
-                nanosecond = magic == PCAP_MAGIC_NS_LE
-            elif magic in (PCAP_MAGIC_BE, PCAP_MAGIC_NS_BE):
-                byte_order = ">"
-                nanosecond = magic == PCAP_MAGIC_NS_BE
-            elif magic == PCAPNG_MAGIC:
-                raise LoaderError(
-                    "PCAPNG format requires dpkt library",
-                    file_path=str(path),
-                    fix_hint="Install dpkt: pip install dpkt",
-                )
-            else:
-                raise FormatError(
-                    "Invalid PCAP magic number",
-                    file_path=str(path),
-                    expected="PCAP magic (0xa1b2c3d4)",
-                    got=f"0x{magic:08x}",
-                )
-
-            # Parse rest of header (version_major, version_minor, thiszone, sigfigs, snaplen, network)
-            _, _, _, _, snaplen, link_type = struct.unpack(f"{byte_order}HHiIII", header[4:])
-
-            packets: list[ProtocolPacket] = []
-
-            # Read packets
-            while True:
-                if max_packets is not None and len(packets) >= max_packets:
-                    break
-
-                # Read packet header (16 bytes)
-                pkt_header = f.read(16)
-                if len(pkt_header) < 16:
-                    break
-
-                ts_sec, ts_usec, incl_len, orig_len = struct.unpack(f"{byte_order}IIII", pkt_header)
-
-                # Calculate timestamp
-                if nanosecond:
-                    timestamp = ts_sec + ts_usec / 1e9
-                else:
-                    timestamp = ts_sec + ts_usec / 1e6
-
-                # Read packet data
-                pkt_data = f.read(incl_len)
-                if len(pkt_data) < incl_len:
-                    break
-
-                packet = ProtocolPacket(
-                    timestamp=timestamp,
-                    protocol="RAW",
-                    data=bytes(pkt_data),
-                    annotations={"original_length": orig_len},
-                )
-                packets.append(packet)
+        with open(path, "rb", buffering=65536) as f:
+            byte_order, nanosecond, snaplen, link_type = _parse_pcap_header(f, path)
+            packets = _read_pcap_packets(f, byte_order, nanosecond, max_packets)
 
             return PcapPacketList(
                 packets=packets,
@@ -390,10 +381,7 @@ def _load_basic(
             )
 
     except struct.error as e:
-        raise FormatError(
-            "Corrupted PCAP file",
-            file_path=str(path),
-        ) from e
+        raise FormatError("Corrupted PCAP file", file_path=str(path)) from e
     except Exception as e:
         if isinstance(e, LoaderError | FormatError):
             raise
@@ -405,6 +393,93 @@ def _load_basic(
         ) from e
 
 
+def _parse_pcap_header(f: Any, path: Path) -> tuple[str, bool, int, int]:
+    """Parse PCAP global header and return format info."""
+    header = f.read(24)
+    if len(header) < 24:
+        raise FormatError(
+            "File too small to be a valid PCAP",
+            file_path=str(path),
+            expected="At least 24 bytes",
+            got=f"{len(header)} bytes",
+        )
+
+    # Parse magic number
+    magic = struct.unpack("<I", header[:4])[0]
+    byte_order, nanosecond = _determine_byte_order(magic, path)
+
+    # Parse rest of header
+    _, _, _, _, snaplen, link_type = struct.unpack(f"{byte_order}HHiIII", header[4:])
+    return byte_order, nanosecond, snaplen, link_type
+
+
+def _determine_byte_order(magic: int, path: Path) -> tuple[str, bool]:
+    """Determine byte order and timestamp precision from magic number."""
+    if magic in (PCAP_MAGIC_LE, PCAP_MAGIC_NS_LE):
+        return "<", magic == PCAP_MAGIC_NS_LE
+    elif magic in (PCAP_MAGIC_BE, PCAP_MAGIC_NS_BE):
+        return ">", magic == PCAP_MAGIC_NS_BE
+    elif magic == PCAPNG_MAGIC:
+        raise LoaderError(
+            "PCAPNG format requires dpkt library",
+            file_path=str(path),
+            fix_hint="Install dpkt: pip install dpkt",
+        )
+    else:
+        raise FormatError(
+            "Invalid PCAP magic number",
+            file_path=str(path),
+            expected="PCAP magic (0xa1b2c3d4)",
+            got=f"0x{magic:08x}",
+        )
+
+
+def _read_pcap_packets(
+    f: Any, byte_order: str, nanosecond: bool, max_packets: int | None
+) -> list[ProtocolPacket]:
+    """Read all packets from PCAP file."""
+    DEFAULT_MAX_PACKETS = 1_000_000  # Prevent memory exhaustion on unbounded files
+
+    packets: list[ProtocolPacket] = []
+    effective_max = max_packets if max_packets is not None else DEFAULT_MAX_PACKETS
+
+    while True:
+        if len(packets) >= effective_max:
+            break
+
+        packet = _read_one_packet(f, byte_order, nanosecond)
+        if packet is None:
+            break
+
+        packets.append(packet)
+
+    return packets
+
+
+def _read_one_packet(f: Any, byte_order: str, nanosecond: bool) -> ProtocolPacket | None:
+    """Read one packet from PCAP file."""
+    pkt_header = f.read(16)
+    if len(pkt_header) < 16:
+        return None
+
+    ts_sec, ts_usec, incl_len, orig_len = struct.unpack(f"{byte_order}IIII", pkt_header)
+
+    # Calculate timestamp
+    timestamp = ts_sec + (ts_usec / 1e9 if nanosecond else ts_usec / 1e6)
+
+    # Read packet data
+    pkt_data = f.read(incl_len)
+    if len(pkt_data) < incl_len:
+        return None
+
+    return ProtocolPacket(
+        timestamp=timestamp,
+        protocol="RAW",
+        data=bytes(pkt_data),
+        annotations={"original_length": orig_len},
+    )
+
+
 def _format_mac(mac_bytes: bytes) -> str:
     """Format MAC address bytes to string.
 

oscura/loaders/rigol.py

@@ -16,6 +16,7 @@ from pathlib import Path
 from typing import TYPE_CHECKING, Any
 
 import numpy as np
+from numpy.typing import NDArray
 
 from oscura.core.exceptions import FormatError, LoaderError
 from oscura.core.types import TraceMetadata, WaveformTrace
@@ -25,7 +26,7 @@ if TYPE_CHECKING:
 
 # Try to import RigolWFM for full Rigol support
 try:
-    import RigolWFM.wfm as rigol_wfm  # type: ignore[import-not-found, import-untyped]
+    import RigolWFM.wfm as rigol_wfm  # type: ignore[import-untyped]  # Optional third-party library
 
     RIGOL_WFM_AVAILABLE = True
 except ImportError:
@@ -100,49 +101,13 @@ def _load_with_rigolwfm(
         LoaderError: If the file cannot be loaded.
     """
     try:
-        # Try to auto-detect model from filename (e.g., DS1054Z)
-        model = None
-        filename_upper = path.name.upper()
-        if "DS1" in filename_upper or "MSO1" in filename_upper or "DHO" in filename_upper:
-            if "Z" in filename_upper or "MSO" in filename_upper or "DHO" in filename_upper:
-                model = "Z"
-            elif "E" in filename_upper:
-                model = "E"
-
-        # Try model detection, fallback to trying both models
-        last_error = None
-        for try_model in [model] if model else ["Z", "E"]:
-            try:
-                wfm = rigol_wfm.Wfm.from_file(str(path), model=try_model)
-                break
-            except Exception as e:
-                last_error = e
-                continue
-        else:
-            # None of the models worked
-            raise last_error if last_error else RuntimeError("Failed to load WFM file")
-
-        # Get channel data
-        if hasattr(wfm, "channels") and len(wfm.channels) > channel:
-            ch = wfm.channels[channel]
-            data = np.array(ch.volts, dtype=np.float64)
-            sample_rate = wfm.sample_rate if hasattr(wfm, "sample_rate") else 1e6
-            vertical_scale = ch.volts_per_div if hasattr(ch, "volts_per_div") else None
-            vertical_offset = ch.volt_offset if hasattr(ch, "volt_offset") else None
-            channel_name = f"CH{channel + 1}"
-        elif hasattr(wfm, "volts"):
-            # Single channel format
-            data = np.array(wfm.volts, dtype=np.float64)
-            sample_rate = wfm.sample_rate if hasattr(wfm, "sample_rate") else 1e6
-            vertical_scale = wfm.volts_per_div if hasattr(wfm, "volts_per_div") else None
-            vertical_offset = wfm.volt_offset if hasattr(wfm, "volt_offset") else None
-            channel_name = "CH1"
-        else:
-            raise FormatError(
-                "No waveform data found in Rigol file",
-                file_path=str(path),
-                expected="Rigol channel data",
-            )
+        # Auto-detect model and load waveform
+        wfm = _load_rigol_with_model_detection(path)
+
+        # Extract channel data
+        data, sample_rate, vertical_scale, vertical_offset, channel_name = (
+            _extract_rigol_channel_data(wfm, channel, str(path))
+        )
 
         # Build metadata
         metadata = TraceMetadata(
@@ -157,12 +122,8 @@ def _load_with_rigolwfm(
         return WaveformTrace(data=data, metadata=metadata)
 
     except Exception as e:
-        # Re-raise FormatError as-is for tests that expect it
-        # All other exceptions (including kaitaistruct errors) get wrapped
         if isinstance(e, FormatError):
             raise
-        # Wrap other exceptions in LoaderError
-        # The outer load_rigol_wfm() will catch LoaderError and fall back to basic loader
         raise LoaderError(
             "Failed to load Rigol WFM file with RigolWFM library",
             file_path=str(path),
@@ -171,6 +132,106 @@ def _load_with_rigolwfm(
         ) from e
 
 
+def _detect_rigol_model_from_filename(path: Path) -> str | None:
+    """Auto-detect Rigol oscilloscope model from filename.
+
+    Args:
+        path: Path to WFM file.
+
+    Returns:
+        Model string ("Z" or "E") or None if detection fails.
+    """
+    filename_upper = path.name.upper()
+
+    # Check for Rigol model indicators in filename
+    if "DS1" not in filename_upper and "MSO1" not in filename_upper and "DHO" not in filename_upper:
+        return None
+
+    # Z-series (DS1000Z, MSO1000Z, DHO series)
+    if "Z" in filename_upper or "MSO" in filename_upper or "DHO" in filename_upper:
+        return "Z"
+
+    # E-series (DS1000E)
+    if "E" in filename_upper:
+        return "E"
+
+    return None
+
+
+def _load_rigol_with_model_detection(path: Path) -> Any:
+    """Load Rigol WFM file with automatic model detection.
+
+    Args:
+        path: Path to WFM file.
+
+    Returns:
+        Loaded waveform object from RigolWFM.
+
+    Raises:
+        RuntimeError: If loading fails with all models.
+    """
+    model = _detect_rigol_model_from_filename(path)
+
+    # Try detected model first, then fallback to both models
+    models_to_try = [model] if model else ["Z", "E"]
+
+    last_error = None
+    for try_model in models_to_try:
+        try:
+            return rigol_wfm.Wfm.from_file(str(path), model=try_model)
+        except Exception as e:
+            last_error = e
+            continue
+
+    # None of the models worked
+    raise last_error if last_error else RuntimeError("Failed to load WFM file")
+
+
+def _extract_rigol_channel_data(
+    wfm: Any,
+    channel: int,
+    file_path: str,
+) -> tuple[NDArray[np.float64], float, float | None, float | None, str]:
+    """Extract channel data from Rigol waveform object.
+
+    Args:
+        wfm: Rigol waveform object from RigolWFM.
+        channel: Channel index.
+        file_path: File path for error messages.
+
+    Returns:
+        Tuple of (data, sample_rate, vertical_scale, vertical_offset, channel_name).
+
+    Raises:
+        FormatError: If no waveform data found.
+    """
+    # Multi-channel format
+    if hasattr(wfm, "channels") and len(wfm.channels) > channel:
+        ch = wfm.channels[channel]
+        data = np.array(ch.volts, dtype=np.float64)
+        sample_rate = wfm.sample_rate if hasattr(wfm, "sample_rate") else 1e6
+        vertical_scale = ch.volts_per_div if hasattr(ch, "volts_per_div") else None
+        vertical_offset = ch.volt_offset if hasattr(ch, "volt_offset") else None
+        channel_name = f"CH{channel + 1}"
+        return data, sample_rate, vertical_scale, vertical_offset, channel_name
+
+    # Single channel format
+    if hasattr(wfm, "volts"):
+        data = np.array(wfm.volts, dtype=np.float64)
+        sample_rate = wfm.sample_rate if hasattr(wfm, "sample_rate") else 1e6
+        vertical_scale = wfm.volts_per_div if hasattr(wfm, "volts_per_div") else None
+        vertical_offset = wfm.volt_offset if hasattr(wfm, "volt_offset") else None
+        channel_name = "CH1"
+        return data, sample_rate, vertical_scale, vertical_offset, channel_name
+
+    # No recognized format
+    raise FormatError(
+        "No waveform data found in Rigol file",
+        file_path=file_path,
+        expected="Rigol channel data",
+    )
+
+
 def _load_basic(
     path: Path,
     *,
@@ -193,7 +254,7 @@ def _load_basic(
         LoaderError: If the file cannot be read or parsed.
     """
     try:
-        with open(path, "rb") as f:
+        with open(path, "rb", buffering=65536) as f:
             # Read header
             header = f.read(256)