| from __future__ import annotations |
|
|
| from collections.abc import Iterable |
|
|
| from .models import Classification, ProcessGroup, ProcessNode, ProcessSample |
|
|
|
|
| def build_nodes( |
| samples: list[ProcessSample], |
| classifications: dict[str, Classification], |
| ) -> dict[str, ProcessNode]: |
| nodes = { |
| sample.process_key: ProcessNode( |
| sample=sample, |
| classification=classifications.get(sample.process_key, Classification("unknown", 0.0)), |
| ) |
| for sample in samples |
| } |
| pid_to_key = {sample.pid: sample.process_key for sample in samples} |
|
|
| for node in nodes.values(): |
| parent_key = pid_to_key.get(node.sample.ppid) |
| if parent_key and parent_key in nodes and parent_key != node.process_key: |
| node.parent_key = parent_key |
| nodes[parent_key].children.append(node) |
|
|
| for node in nodes.values(): |
| node.children.sort(key=lambda item: item.sample.pid) |
|
|
| for root in _roots(nodes.values()): |
| _propagate_agent_context(root, inherited=False) |
|
|
| return nodes |
|
|
|
|
| def _roots(nodes: Iterable[ProcessNode]) -> list[ProcessNode]: |
| node_list = list(nodes) |
| keys = {node.process_key for node in node_list} |
| return [node for node in node_list if not node.parent_key or node.parent_key not in keys] |
|
|
|
|
| def _propagate_agent_context(node: ProcessNode, *, inherited: bool) -> None: |
| current = inherited or node.classification.is_agent or node.classification.is_agent_related |
| node.inherited_agent = current |
| for child in node.children: |
| _propagate_agent_context(child, inherited=current) |
|
|
|
|
| def iter_descendants(node: ProcessNode, *, include_self: bool = True) -> Iterable[ProcessNode]: |
| if include_self: |
| yield node |
| for child in node.children: |
| yield from iter_descendants(child, include_self=True) |
|
|
|
|
| def _has_agent_ancestor(node: ProcessNode, nodes: dict[str, ProcessNode]) -> bool: |
| parent_key = node.parent_key |
| while parent_key and parent_key in nodes: |
| parent = nodes[parent_key] |
| if parent.classification.is_agent: |
| return True |
| parent_key = parent.parent_key |
| return False |
|
|
|
|
| def build_groups(nodes: dict[str, ProcessNode]) -> list[ProcessGroup]: |
| group_roots: list[ProcessNode] = [] |
| for node in nodes.values(): |
| classification = node.classification |
| if classification.is_agent and not _has_agent_ancestor(node, nodes): |
| group_roots.append(node) |
| elif ( |
| classification.is_mcp or classification.is_bridge |
| ) and not _has_agent_ancestor(node, nodes): |
| group_roots.append(node) |
|
|
| seen: set[str] = set() |
| groups: list[ProcessGroup] = [] |
| for root in sorted(group_roots, key=lambda item: item.sample.pid): |
| if root.process_key in seen: |
| continue |
| seen.add(root.process_key) |
| members = list(iter_descendants(root)) |
| bridge_count = sum(1 for node in members if node.classification.is_bridge) |
| mcp_count = sum(1 for node in members if node.classification.is_mcp) |
| groups.append( |
| ProcessGroup( |
| group_id=root.process_key, |
| root=root, |
| nodes=members, |
| label=_group_label(root), |
| child_count=max(0, len(members) - 1), |
| bridge_count=bridge_count, |
| mcp_count=mcp_count, |
| total_rss_bytes=sum(node.sample.rss_bytes for node in members), |
| max_age_seconds=max((node.sample.age_seconds for node in members), default=0.0), |
| ) |
| ) |
| return groups |
|
|
|
|
| def _group_label(root: ProcessNode) -> str: |
| argv0 = root.sample.argv[0] if root.sample.argv else root.sample.comm |
| name = argv0.rsplit("/", 1)[-1] or root.sample.comm |
| return f"{name} - PID {root.sample.pid}" |
|
|
