104 lines
3.8 KiB
Python
104 lines
3.8 KiB
Python
|
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
|