oscura 0.5.1__py3-none-any.whl → 0.7.0__py3-none-any.whl

oscura/{pipeline → utils/pipeline}/reverse_engineering.py

@@ -17,7 +17,7 @@ from collections.abc import Callable, Sequence
 from dataclasses import dataclass, field
 from datetime import datetime
 from pathlib import Path
-from typing import Any, ClassVar, Literal
+from typing import Any, ClassVar, Literal, cast
 
 logger = logging.getLogger(__name__)
 
@@ -206,6 +206,111 @@ class REPipeline:
         self._checkpoint_path: str | None = None
         self._checkpoint_data: dict[str, Any] = {}
 
+    def _initialize_analysis_context(
+        self, data: bytes | Sequence[dict[str, Any]] | Sequence[bytes]
+    ) -> dict[str, Any]:
+        """Initialize analysis context with empty containers.
+
+        Args:
+            data: Input data to analyze.
+
+        Returns:
+            Initialized context dictionary.
+        """
+        return {
+            "raw_data": data,
+            "flows": [],
+            "payloads": [],
+            "messages": [],
+            "patterns": [],
+            "clusters": [],
+            "schemas": {},
+            "protocol_candidates": [],
+            "state_machine": None,
+            "warnings": [],
+            "statistics": {},
+        }
+
+    def _execute_stage(
+        self,
+        stage_name: str,
+        context: dict[str, Any],
+        checkpoint: str | None,
+    ) -> StageResult:
+        """Execute single pipeline stage.
+
+        Args:
+            stage_name: Name of stage to execute.
+            context: Analysis context.
+            checkpoint: Checkpoint path.
+
+        Returns:
+            StageResult with execution outcome.
+        """
+        handler = self._stage_handlers.get(stage_name)
+        if not handler:
+            return StageResult(
+                stage_name=stage_name, success=False, duration=0, output=None, error="No handler"
+            )
+
+        try:
+            stage_start = time.time()
+            output = handler(context)
+            stage_duration = time.time() - stage_start
+
+            if output:
+                context.update(output)
+
+            if checkpoint:
+                self._save_checkpoint(checkpoint, stage_name, context)
+
+            return StageResult(
+                stage_name=stage_name,
+                success=True,
+                duration=stage_duration,
+                output=output,
+            )
+
+        except Exception as e:
+            warnings_list: list[str] = context.get("warnings", [])
+            warnings_list.append(f"Stage {stage_name} failed: {e}")
+            context["warnings"] = warnings_list
+
+            return StageResult(
+                stage_name=stage_name,
+                success=False,
+                duration=0,
+                output=None,
+                error=str(e),
+            )
+
+    def _execute_all_stages(
+        self, context: dict[str, Any], checkpoint: str | None
+    ) -> list[StageResult]:
+        """Execute all pipeline stages.
+
+        Args:
+            context: Analysis context.
+            checkpoint: Checkpoint path.
+
+        Returns:
+            List of stage results.
+        """
+        stage_results = []
+        total_stages = len(self.stages)
+
+        for i, stage_name in enumerate(self.stages):
+            if stage_name in self._checkpoint_data:
+                context.update(self._checkpoint_data[stage_name])
+                continue
+
+            self._report_progress(stage_name, (i / total_stages) * 100)
+            stage_result = self._execute_stage(stage_name, context, checkpoint)
+            stage_results.append(stage_result)
+
+        self._report_progress("complete", 100)
+        return stage_results
+
     def analyze(
         self,
         data: bytes | Sequence[dict[str, Any]] | Sequence[bytes],
@@ -229,85 +334,22 @@ class REPipeline:
             >>> for msg_type in results.message_types:
             ...     print(f"{msg_type.name}: {msg_type.sample_count} samples")
         """
+        # Setup: initialize state and load checkpoint
         start_time = time.time()
         self._progress_callback = progress_callback
         self._checkpoint_path = checkpoint
         self._checkpoint_data = {}
 
-        # Load checkpoint if available
         if checkpoint and os.path.exists(checkpoint):
             self._load_checkpoint(checkpoint)
 
-        # Initialize context
-        context: dict[str, Any] = {
-            "raw_data": data,
-            "flows": [],
-            "payloads": [],
-            "messages": [],
-            "patterns": [],
-            "clusters": [],
-            "schemas": {},
-            "protocol_candidates": [],
-            "state_machine": None,
-            "warnings": [],
-            "statistics": {},
-        }
-
-        # Execute stages
-        stage_results = []
-        total_stages = len(self.stages)
-
-        for i, stage_name in enumerate(self.stages):
-            if stage_name in self._checkpoint_data:
-                # Skip completed stages
-                context.update(self._checkpoint_data[stage_name])
-                continue
-
-            self._report_progress(stage_name, (i / total_stages) * 100)
+        context = self._initialize_analysis_context(data)
 
-            handler = self._stage_handlers.get(stage_name)
-            if handler:
-                try:
-                    stage_start = time.time()
-                    output = handler(context)
-                    stage_duration = time.time() - stage_start
-
-                    stage_results.append(
-                        StageResult(
-                            stage_name=stage_name,
-                            success=True,
-                            duration=stage_duration,
-                            output=output,
-                        )
-                    )
+        # Processing: execute pipeline stages
+        stage_results = self._execute_all_stages(context, checkpoint)
 
-                    # Update context with stage output
-                    if output:
-                        context.update(output)
-
-                    # Checkpoint after each stage
-                    if checkpoint:
-                        self._save_checkpoint(checkpoint, stage_name, context)
-
-                except Exception as e:
-                    stage_results.append(
-                        StageResult(
-                            stage_name=stage_name,
-                            success=False,
-                            duration=0,
-                            output=None,
-                            error=str(e),
-                        )
-                    )
-                    warnings_list: list[str] = context.get("warnings", [])
-                    warnings_list.append(f"Stage {stage_name} failed: {e}")
-                    context["warnings"] = warnings_list
-
-        self._report_progress("complete", 100)
-
-        # Build result
+        # Result building: construct final result
         duration = time.time() - start_time
-
         flows_list: list[Any] = context.get("flows", [])
         messages_list: list[Any] = context.get("messages", [])
         protocol_candidates_list: list[ProtocolCandidate] = context.get("protocol_candidates", [])
@@ -416,115 +458,205 @@ class REPipeline:
             context: Pipeline context.
 
         Returns:
-            Updated context with flows.
+            Updated context with flows and payloads.
         """
         data = context["raw_data"]
-        flows = []
-        payloads = []
 
         if isinstance(data, bytes):
-            # Raw binary - treat as single payload
-            payloads.append(data)
+            flows, payloads = self._extract_from_raw_bytes(data)
+        elif isinstance(data, list | tuple):
+            flows, payloads = self._extract_from_packet_list(data)
+        else:
+            flows, payloads = [], []
+
+        self._update_flow_statistics(context, flows, payloads)
+        return {"flows": flows, "payloads": payloads}
+
+    def _extract_from_raw_bytes(self, data: bytes) -> tuple[list[FlowInfo], list[bytes]]:
+        """Extract flow from raw binary data.
+
+        Args:
+            data: Raw binary data.
+
+        Returns:
+            Tuple of (flows, payloads).
+        """
+        flow = FlowInfo(
+            flow_id="flow_0",
+            src_ip="unknown",
+            dst_ip="unknown",
+            src_port=0,
+            dst_port=0,
+            protocol="unknown",
+            packet_count=1,
+            byte_count=len(data),
+            start_time=0,
+            end_time=0,
+        )
+        return [flow], [data]
+
+    def _extract_from_packet_list(
+        self, packets: Sequence[dict[str, Any] | bytes]
+    ) -> tuple[list[FlowInfo], list[bytes]]:
+        """Extract flows from list of packets.
+
+        Args:
+            packets: List of packet dicts or raw bytes.
+
+        Returns:
+            Tuple of (flows, payloads).
+        """
+        flow_map: dict[str, dict[str, Any]] = {}
+        payloads: list[bytes] = []
+        raw_bytes_payloads: list[bytes] = []
+
+        for pkt in packets:
+            if isinstance(pkt, dict):
+                self._process_packet_dict(pkt, flow_map, payloads)
+            else:
+                payload = bytes(pkt) if not isinstance(pkt, bytes) else pkt
+                payloads.append(payload)
+                raw_bytes_payloads.append(payload)
+
+        flows = self._build_flows_from_map(flow_map)
+
+        # Create default flow for raw bytes if needed
+        if raw_bytes_payloads and not flows:
+            flows.append(self._create_default_flow(raw_bytes_payloads))
+
+        return flows, payloads
+
+    def _process_packet_dict(
+        self,
+        pkt: dict[str, Any],
+        flow_map: dict[str, dict[str, Any]],
+        payloads: list[bytes],
+    ) -> None:
+        """Process a packet dictionary and update flow map.
+
+        Args:
+            pkt: Packet dictionary with metadata.
+            flow_map: Flow mapping to update.
+            payloads: Payloads list to append to.
+        """
+        # Extract payload
+        payload_raw = pkt.get("data", pkt.get("payload", b""))
+        if isinstance(payload_raw, list | tuple):
+            payload = bytes(payload_raw)
+        else:
+            payload = payload_raw if isinstance(payload_raw, bytes) else b""
+
+        # Create flow key
+        flow_key = self._create_flow_key(pkt)
+
+        # Initialize flow if new
+        if flow_key not in flow_map:
+            flow_map[flow_key] = self._create_flow_entry(pkt)
+
+        # Update flow data
+        flow_map[flow_key]["packets"].append(pkt)
+        flow_map[flow_key]["payloads"].append(payload)
+        if "timestamp" in pkt:
+            flow_map[flow_key]["timestamps"].append(pkt["timestamp"])
+
+        payloads.append(payload)
+
+    def _create_flow_key(self, pkt: dict[str, Any]) -> str:
+        """Create flow identifier key from packet.
+
+        Args:
+            pkt: Packet dictionary.
+
+        Returns:
+            Flow key string.
+        """
+        src_ip = pkt.get("src_ip", "0.0.0.0")
+        dst_ip = pkt.get("dst_ip", "0.0.0.0")
+        src_port = pkt.get("src_port", 0)
+        dst_port = pkt.get("dst_port", 0)
+        protocol = pkt.get("protocol", "unknown")
+        return f"{src_ip}:{src_port}-{dst_ip}:{dst_port}-{protocol}"
+
+    def _create_flow_entry(self, pkt: dict[str, Any]) -> dict[str, Any]:
+        """Create new flow entry from packet.
+
+        Args:
+            pkt: Packet dictionary.
+
+        Returns:
+            Flow entry dictionary.
+        """
+        return {
+            "src_ip": pkt.get("src_ip", "0.0.0.0"),
+            "dst_ip": pkt.get("dst_ip", "0.0.0.0"),
+            "src_port": pkt.get("src_port", 0),
+            "dst_port": pkt.get("dst_port", 0),
+            "protocol": pkt.get("protocol", "unknown"),
+            "packets": [],
+            "payloads": [],
+            "timestamps": [],
+        }
+
+    def _build_flows_from_map(self, flow_map: dict[str, dict[str, Any]]) -> list[FlowInfo]:
+        """Build FlowInfo objects from flow map.
+
+        Args:
+            flow_map: Mapping of flow keys to flow data.
+
+        Returns:
+            List of FlowInfo objects.
+        """
+        flows = []
+        for flow_id, flow_data in flow_map.items():
+            timestamps = flow_data.get("timestamps", [0])
             flows.append(
                 FlowInfo(
-                    flow_id="flow_0",
-                    src_ip="unknown",
-                    dst_ip="unknown",
-                    src_port=0,
-                    dst_port=0,
-                    protocol="unknown",
-                    packet_count=1,
-                    byte_count=len(data),
-                    start_time=0,
-                    end_time=0,
+                    flow_id=flow_id,
+                    src_ip=flow_data["src_ip"],
+                    dst_ip=flow_data["dst_ip"],
+                    src_port=flow_data["src_port"],
+                    dst_port=flow_data["dst_port"],
+                    protocol=flow_data["protocol"],
+                    packet_count=len(flow_data["packets"]),
+                    byte_count=sum(len(p) for p in flow_data["payloads"]),
+                    start_time=min(timestamps) if timestamps else 0,
+                    end_time=max(timestamps) if timestamps else 0,
                 )
             )
+        return flows
 
-        elif isinstance(data, list | tuple):
-            # List of packets
-            flow_map: dict[str, dict[str, Any]] = {}
-            raw_bytes_payloads: list[bytes] = []
-
-            for _i, pkt in enumerate(data):
-                if isinstance(pkt, dict):
-                    # Packet with metadata
-                    payload_raw = pkt.get("data", pkt.get("payload", b""))
-                    if isinstance(payload_raw, list | tuple):
-                        payload = bytes(payload_raw)
-                    else:
-                        payload = payload_raw if isinstance(payload_raw, bytes) else b""
-
-                    # Create flow key
-                    src_ip = pkt.get("src_ip", "0.0.0.0")
-                    dst_ip = pkt.get("dst_ip", "0.0.0.0")
-                    src_port = pkt.get("src_port", 0)
-                    dst_port = pkt.get("dst_port", 0)
-                    protocol = pkt.get("protocol", "unknown")
-
-                    flow_key = f"{src_ip}:{src_port}-{dst_ip}:{dst_port}-{protocol}"
-
-                    if flow_key not in flow_map:
-                        flow_map[flow_key] = {
-                            "src_ip": src_ip,
-                            "dst_ip": dst_ip,
-                            "src_port": src_port,
-                            "dst_port": dst_port,
-                            "protocol": protocol,
-                            "packets": [],
-                            "payloads": [],
-                            "timestamps": [],
-                        }
-
-                    flow_map[flow_key]["packets"].append(pkt)
-                    flow_map[flow_key]["payloads"].append(payload)
-                    if "timestamp" in pkt:
-                        flow_map[flow_key]["timestamps"].append(pkt["timestamp"])
-
-                    payloads.append(payload)
-
-                else:
-                    # Raw bytes - collect for default flow
-                    raw_payload = bytes(pkt) if not isinstance(pkt, bytes) else pkt
-                    payloads.append(raw_payload)
-                    raw_bytes_payloads.append(raw_payload)
-
-            # Build flow objects from flow_map
-            for flow_id, flow_data in flow_map.items():
-                timestamps = flow_data.get("timestamps", [0])
-                flows.append(
-                    FlowInfo(
-                        flow_id=flow_id,
-                        src_ip=flow_data["src_ip"],
-                        dst_ip=flow_data["dst_ip"],
-                        src_port=flow_data["src_port"],
-                        dst_port=flow_data["dst_port"],
-                        protocol=flow_data["protocol"],
-                        packet_count=len(flow_data["packets"]),
-                        byte_count=sum(len(p) for p in flow_data["payloads"]),
-                        start_time=min(timestamps) if timestamps else 0,
-                        end_time=max(timestamps) if timestamps else 0,
-                    )
-                )
+    def _create_default_flow(self, payloads: list[bytes]) -> FlowInfo:
+        """Create default flow for raw bytes.
 
-            # Create a default flow for raw bytes if we have any
-            # This ensures flow_count >= 1 when analyzing raw byte sequences
-            if raw_bytes_payloads and not flows:
-                flows.append(
-                    FlowInfo(
-                        flow_id="flow_default",
-                        src_ip="unknown",
-                        dst_ip="unknown",
-                        src_port=0,
-                        dst_port=0,
-                        protocol="unknown",
-                        packet_count=len(raw_bytes_payloads),
-                        byte_count=sum(len(p) for p in raw_bytes_payloads),
-                        start_time=0,
-                        end_time=0,
-                    )
-                )
+        Args:
+            payloads: List of raw byte payloads.
 
-        # Initialize statistics dict if it doesn't exist
+        Returns:
+            Default FlowInfo object.
+        """
+        return FlowInfo(
+            flow_id="flow_default",
+            src_ip="unknown",
+            dst_ip="unknown",
+            src_port=0,
+            dst_port=0,
+            protocol="unknown",
+            packet_count=len(payloads),
+            byte_count=sum(len(p) for p in payloads),
+            start_time=0,
+            end_time=0,
+        )
+
+    def _update_flow_statistics(
+        self, context: dict[str, Any], flows: list[FlowInfo], payloads: list[bytes]
+    ) -> None:
+        """Update context statistics with flow extraction results.
+
+        Args:
+            context: Pipeline context to update.
+            flows: Extracted flows.
+            payloads: Extracted payloads.
+        """
         if "statistics" not in context:
             context["statistics"] = {}
 
@@ -534,8 +666,6 @@ class REPipeline:
             "total_bytes": sum(len(p) for p in payloads),
         }
 
-        return {"flows": flows, "payloads": payloads}
-
     def _stage_payload_analysis(self, context: dict[str, Any]) -> dict[str, Any]:
         """Analyze payloads for structure.
 
@@ -713,9 +843,27 @@ class REPipeline:
         """
         messages = context.get("messages", [])
         flows = context.get("flows", [])
-        candidates = []
+        candidates: list[ProtocolCandidate] = []
+
+        # Detect protocols from multiple sources
+        candidates.extend(self._detect_by_port(flows))
+        candidates.extend(self._detect_by_magic_bytes(messages))
+        candidates.extend(self._detect_from_library(messages))
 
-        # Check well-known port mappings
+        # Deduplicate candidates
+        unique_candidates = self._deduplicate_candidates(candidates)
+
+        return {"protocol_candidates": unique_candidates}
+
+    def _detect_by_port(self, flows: list[FlowInfo]) -> list[ProtocolCandidate]:
+        """Detect protocols based on well-known port numbers.
+
+        Args:
+            flows: List of network flows.
+
+        Returns:
+            List of protocol candidates.
+        """
         port_protocols = {
             53: "dns",
             80: "http",
@@ -726,6 +874,7 @@ class REPipeline:
             47808: "bacnet",
         }
 
+        candidates = []
         for flow in flows:
             port = flow.dst_port or flow.src_port
             if port in port_protocols:
@@ -737,69 +886,110 @@ class REPipeline:
                     )
                 )
 
-        # Check magic byte signatures
-        if messages:
-            try:
-                from oscura.inference.binary import MagicByteDetector
-
-                detector = MagicByteDetector()
-                sample = messages[0] if messages else b""
-
-                if len(sample) >= 2:
-                    result = detector.detect(sample)
-                    if result and result.known_format:
-                        candidates.append(
-                            ProtocolCandidate(
-                                name=result.known_format,
-                                confidence=result.confidence,
-                                header_match=True,
-                            )
+        return candidates
+
+    def _detect_by_magic_bytes(self, messages: list[bytes]) -> list[ProtocolCandidate]:
+        """Detect protocols by magic byte signatures.
+
+        Args:
+            messages: List of message bytes.
+
+        Returns:
+            List of protocol candidates.
+        """
+        if not messages:
+            return []
+
+        try:
+            from oscura.inference.binary import MagicByteDetector
+
+            detector = MagicByteDetector()
+            sample = messages[0]
+
+            if len(sample) >= 2:
+                result = detector.detect(sample)
+                if result and result.known_format:
+                    return [
+                        ProtocolCandidate(
+                            name=result.known_format,
+                            confidence=result.confidence,
+                            header_match=True,
                         )
+                    ]
 
-            except Exception as e:
-                logger.debug("Magic byte detection failed (non-critical): %s", e)
+        except Exception as e:
+            logger.debug("Magic byte detection failed (non-critical): %s", e)
+
+        return []
+
+    def _detect_from_library(self, messages: list[bytes]) -> list[ProtocolCandidate]:
+        """Detect protocols from protocol library.
 
-        # Check protocol library
+        Args:
+            messages: List of message bytes.
+
+        Returns:
+            List of protocol candidates.
+        """
         try:
             from oscura.inference.protocol_library import get_library
 
             library = get_library()
+            candidates = []
 
             for protocol in library.list_protocols():
-                if protocol.definition:
-                    # Check if first bytes match protocol header
-                    for msg in messages[:10]:
-                        if len(msg) >= 4:
-                            # Simple header matching
-                            first_field = (
-                                protocol.definition.fields[0]
-                                if protocol.definition.fields
-                                else None
-                            )
-                            if first_field and hasattr(first_field, "value"):
-                                # Has expected value
-                                candidates.append(
-                                    ProtocolCandidate(
-                                        name=protocol.name,
-                                        confidence=0.4,
-                                        matched_patterns=["header_value"],
-                                    )
-                                )
-                                break
+                if self._matches_protocol_header(protocol, messages):
+                    candidates.append(
+                        ProtocolCandidate(
+                            name=protocol.name,
+                            confidence=0.4,
+                            matched_patterns=["header_value"],
+                        )
+                    )
+
+            return candidates
 
         except Exception as e:
             logger.debug("Protocol library matching failed (non-critical): %s", e)
+            return []
+
+    def _matches_protocol_header(self, protocol: Any, messages: list[bytes]) -> bool:
+        """Check if messages match protocol header.
 
-        # Deduplicate by name, keeping highest confidence
-        unique_candidates: dict[str, ProtocolCandidate] = {}
+        Args:
+            protocol: Protocol definition.
+            messages: List of message bytes.
+
+        Returns:
+            True if matches.
+        """
+        if not protocol.definition or not protocol.definition.fields:
+            return False
+
+        first_field = protocol.definition.fields[0]
+        if not hasattr(first_field, "value"):
+            return False
+
+        # Check first 10 messages
+        return any(len(msg) >= 4 for msg in messages[:10])
+
+    def _deduplicate_candidates(
+        self, candidates: list[ProtocolCandidate]
+    ) -> list[ProtocolCandidate]:
+        """Deduplicate candidates, keeping highest confidence.
+
+        Args:
+            candidates: List of candidates.
+
+        Returns:
+            Deduplicated list.
+        """
+        unique: dict[str, ProtocolCandidate] = {}
         for c in candidates:
-            if (
-                c.name not in unique_candidates
-                or c.confidence > unique_candidates[c.name].confidence
-            ):
-                unique_candidates[c.name] = c
+            if c.name not in unique or c.confidence > unique[c.name].confidence:
+                unique[c.name] = c
 
-        return {"protocol_candidates": list(unique_candidates.values())}
+        return list(unique.values())
 
     def _stage_state_machine(self, context: dict[str, Any]) -> dict[str, Any]:
         """Infer protocol state machine.
@@ -825,7 +1015,7 @@ class REPipeline:
                     message_to_cluster[idx] = getattr(cluster, "cluster_id", 0)
 
             # Build observation sequence
-            sequence = [
+            sequence: list[str] = [
                 f"type_{message_to_cluster.get(i, 0)}"
                 for i in range(len(messages))
                 if i in message_to_cluster
@@ -835,7 +1025,8 @@ class REPipeline:
                 from oscura.inference.state_machine import StateMachineInferrer
 
                 inferrer = StateMachineInferrer()
-                automaton = inferrer.infer_rpni([sequence])
+                # Cast list[str] to list[str | int] for API compatibility
+                automaton = inferrer.infer_rpni([cast("list[str | int]", sequence)])
 
                 return {
                     "state_machine": {