diff --git a/backend/src/utilities/tshark_manager.py b/backend/src/utilities/tshark_manager.py index 1316229..73af08e 100644 --- a/backend/src/utilities/tshark_manager.py +++ b/backend/src/utilities/tshark_manager.py @@ -288,6 +288,10 @@ def _eth_type_from_protocol_stack(protocol_stack: List[str]) -> Optional[int]: return None +def _is_generic_stream_protocol(protocol_name: Optional[str]) -> bool: + return str(protocol_name or "").replace("-", "_").replace(".", "_").lower() in _GENERIC_PROTOCOLS + + def _http_metadata(layers: Dict[str, Any]) -> Dict[str, Any]: http_layer = layers.get("http") if not isinstance(http_layer, dict): @@ -532,15 +536,36 @@ def _build_enrichment(event: Dict[str, Any], flow_context: Optional[Dict[str, An derived_context = _derive_flow_context(event) protocol_stack = list(event.get("protocol_stack") or []) - app_protocol = derived_context.get("app_protocol") or flow_context.get("app_protocol") - app_category = derived_context.get("app_category") or flow_context.get("app_category") + derived_protocol = _safe_text(derived_context.get("app_protocol")) + flow_protocol = _safe_text(flow_context.get("app_protocol")) + prefer_flow_context = ( + flow_protocol is not None + and not _is_generic_stream_protocol(flow_protocol) + and derived_protocol is not None + and _is_generic_stream_protocol(derived_protocol) + ) + + app_protocol = flow_protocol if prefer_flow_context else (derived_protocol or flow_protocol) + app_category = ( + flow_context.get("app_category") or derived_context.get("app_category") + if prefer_flow_context + else (derived_context.get("app_category") or flow_context.get("app_category")) + ) app_hostname = derived_context.get("app_hostname") or flow_context.get("app_hostname") app_is_encrypted = ( - derived_context.get("app_is_encrypted") - if derived_context.get("app_is_encrypted") is not None - else flow_context.get("app_is_encrypted") + flow_context.get("app_is_encrypted") + if prefer_flow_context and flow_context.get("app_is_encrypted") is not None + else ( + derived_context.get("app_is_encrypted") + if derived_context.get("app_is_encrypted") is not None + else flow_context.get("app_is_encrypted") + ) + ) + app_confidence = ( + flow_context.get("app_confidence") or derived_context.get("app_confidence") + if prefer_flow_context + else (derived_context.get("app_confidence") or flow_context.get("app_confidence")) ) - app_confidence = derived_context.get("app_confidence") or flow_context.get("app_confidence") tcp_flags = _safe_int(event.get("tcp_flags")) tcp_len = int(event.get("tcp_len") or 0) @@ -631,8 +656,34 @@ def _build_stream_enrichment(flow_context: Dict[str, Any]) -> Dict[str, Any]: } -def _is_generic_stream_protocol(protocol_name: Optional[str]) -> bool: - return str(protocol_name or "").replace("-", "_").replace(".", "_").lower() in _GENERIC_PROTOCOLS +def _merge_flow_context(existing: Optional[Dict[str, Any]], incoming: Optional[Dict[str, Any]]) -> Dict[str, Any]: + merged = dict(existing or {}) + incoming_context = dict(incoming or {}) + existing_protocol = _safe_text(merged.get("app_protocol")) + incoming_protocol = _safe_text(incoming_context.get("app_protocol")) + preserve_existing_protocol = ( + existing_protocol is not None + and not _is_generic_stream_protocol(existing_protocol) + and incoming_protocol is not None + and _is_generic_stream_protocol(incoming_protocol) + ) + + for key, value in incoming_context.items(): + if value is None: + continue + if preserve_existing_protocol and key in { + "app_protocol", + "app_master_protocol", + "app_category", + "app_confidence", + "app_is_encrypted", + }: + continue + merged[key] = value + + if merged.get("app_master_protocol") is None and merged.get("app_protocol") is not None: + merged["app_master_protocol"] = merged["app_protocol"] + return merged class TsharkManager: @@ -864,8 +915,7 @@ class TsharkManager: flow_context = dict(self._flow_context.get(stream_key) or {}) derived_context = _derive_flow_context(event) if stream_key is not None and derived_context.get("app_protocol") is not None: - merged_context = dict(flow_context or {}) - merged_context.update({key: value for key, value in derived_context.items() if value is not None}) + merged_context = _merge_flow_context(flow_context, derived_context) self._flow_context[stream_key] = merged_context flow_context = merged_context entries = self._cache.setdefault(signature, [])