from collections import defaultdict from app.schemas import InventoryTreeNode, NetObject, NodeContext, NodeContextRelations, Relation from app.services.store import OBJECTS, RELATIONS def get_object_map() -> dict[str, NetObject]: return {item.id: item for item in OBJECTS} def get_children_map() -> dict[str, list[NetObject]]: object_map = get_object_map() children: dict[str, list[NetObject]] = defaultdict(list) for rel in RELATIONS: if rel.type == 'contains' and rel.status == 'active' and rel.target_id in object_map: children[rel.source_id].append(object_map[rel.target_id]) return children def get_badges(object_id: str) -> list[str]: obj = get_object_map()[object_id] badges: list[str] = [] if obj.type == 'physical_host': badges.append('virtualization-host') if obj.type in {'vm', 'router', 'firewall'}: badges.append('network-node') outgoing = [r for r in RELATIONS if r.source_id == object_id and r.status == 'active'] if any(r.type == 'routes_via' for r in outgoing): badges.append('gateway-dependent') if any(r.type in {'routes_via', 'bridges_to'} for r in outgoing) or obj.name.lower().startswith('openwrt'): badges.append('gateway') return badges def build_tree_node(object_id: str, depth: int | None = None) -> InventoryTreeNode: object_map = get_object_map() children_map = get_children_map() obj = object_map[object_id] child_objects = children_map.get(object_id, []) child_nodes = [] if depth is None or depth > 0: next_depth = None if depth is None else depth - 1 child_nodes = [build_tree_node(child.id, next_depth) for child in child_objects] return InventoryTreeNode( id=obj.id, name=obj.name, type=obj.type, status=obj.status, badges=get_badges(obj.id), children_count=len(child_objects), children=child_nodes, ) def build_inventory_tree(root_id: str | None = None, depth: int | None = None) -> list[InventoryTreeNode]: object_map = get_object_map() if root_id: return [build_tree_node(root_id, depth)] if root_id in object_map else [] roots = [obj for obj in OBJECTS if obj.parent_id is None] return [build_tree_node(obj.id, depth) for obj in roots] def get_parent(item: NetObject) -> NetObject | None: object_map = get_object_map() if not item.parent_id: return None return object_map.get(item.parent_id) def get_children(item_id: str) -> list[NetObject]: return get_children_map().get(item_id, []) def get_relations(item_id: str) -> NodeContextRelations: incoming = [rel for rel in RELATIONS if rel.target_id == item_id and rel.status == 'active'] outgoing = [rel for rel in RELATIONS if rel.source_id == item_id and rel.status == 'active'] return NodeContextRelations(incoming=incoming, outgoing=outgoing) def build_node_context(item_id: str) -> NodeContext | None: object_map = get_object_map() item = object_map.get(item_id) if not item: return None return NodeContext( item=item, parent=get_parent(item), children=get_children(item_id), relations=get_relations(item_id), ) def get_node_relations(item_id: str, layer: str | None = None, rel_type: str | None = None, direction: str = 'both') -> list[Relation]: result: list[Relation] = [] if direction in {'incoming', 'both'}: result.extend([r for r in RELATIONS if r.target_id == item_id and r.status == 'active']) if direction in {'outgoing', 'both'}: result.extend([r for r in RELATIONS if r.source_id == item_id and r.status == 'active']) if layer: result = [r for r in result if r.layer == layer] if rel_type: result = [r for r in result if r.type == rel_type] return result