File size: 3,817 Bytes
06357e7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 | 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}"
|