import json from collections import defaultdict from sqlalchemy.orm import Session from app.models import NetObjectModel, RelationModel from app.schemas import InventoryTreeNode, NetObject, NodeContext, NodeContextRelations, Relation def _to_net_object(m: NetObjectModel) -> NetObject: return NetObject( id=m.id, name=m.name, slug=m.slug, type=m.type, parent_id=m.parent_id, status=m.status, vendor=m.vendor, model=m.model, serial=m.serial, notes=m.notes, created_at=m.created_at, updated_at=m.updated_at, ) def _to_relation(m: RelationModel) -> Relation: return Relation( id=m.id, type=m.type, layer=m.layer, source_id=m.source_id, target_id=m.target_id, direction=m.direction, source=m.source, confidence=m.confidence, status=m.status, attributes=json.loads(m.attributes_json or "{}"), created_at=m.created_at, updated_at=m.updated_at, ) def get_object_map(db: Session) -> dict[str, NetObject]: return {r.id: _to_net_object(r) for r in db.query(NetObjectModel).all()} def get_children_map(db: Session) -> dict[str, list[NetObject]]: object_map = get_object_map(db) rels = db.query(RelationModel).filter( RelationModel.type == "contains", RelationModel.status == "active", ).all() children: dict[str, list[NetObject]] = defaultdict(list) for rel in rels: if rel.target_id in object_map: children[rel.source_id].append(object_map[rel.target_id]) return children def get_badges(db: Session, object_id: str) -> list[str]: obj = get_object_map(db).get(object_id) if not obj: return [] badges: list[str] = [] if obj.type == "physical_host": badges.append("virtualization-host") if obj.type in {"vm", "router", "firewall"}: badges.append("network-node") outgoing = db.query(RelationModel).filter( RelationModel.source_id == object_id, RelationModel.status == "active", ).all() 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(db: Session, object_id: str, depth: int | None = None) -> InventoryTreeNode: om = get_object_map(db) cm = get_children_map(db) obj = om[object_id] ch = cm.get(object_id, []) child_nodes = [] if depth is None or depth > 0: nd = None if depth is None else depth - 1 child_nodes = [build_tree_node(db, c.id, nd) for c in ch] return InventoryTreeNode( id=obj.id, name=obj.name, type=obj.type, status=obj.status, badges=get_badges(db, obj.id), children_count=len(ch), children=child_nodes, ) def build_inventory_tree( db: Session, root_id: str | None = None, depth: int | None = None ) -> list[InventoryTreeNode]: om = get_object_map(db) if root_id: return [build_tree_node(db, root_id, depth)] if root_id in om else [] roots = db.query(NetObjectModel).filter(NetObjectModel.parent_id == None).all() # noqa: E711 return [build_tree_node(db, r.id, depth) for r in roots] def build_node_context(db: Session, item_id: str) -> NodeContext | None: row = db.query(NetObjectModel).filter(NetObjectModel.id == item_id).first() if not row: return None item = _to_net_object(row) parent = None if item.parent_id: p = db.query(NetObjectModel).filter(NetObjectModel.id == item.parent_id).first() parent = _to_net_object(p) if p else None children = get_children_map(db).get(item_id, []) inc = [_to_relation(r) for r in db.query(RelationModel).filter( RelationModel.target_id == item_id, RelationModel.status == "active").all()] out = [_to_relation(r) for r in db.query(RelationModel).filter( RelationModel.source_id == item_id, RelationModel.status == "active").all()] return NodeContext( item=item, parent=parent, children=children, relations=NodeContextRelations(incoming=inc, outgoing=out), ) def get_node_relations( db: Session, item_id: str, layer: str | None = None, rel_type: str | None = None, direction: str = "both", ) -> list[Relation]: result: list[Relation] = [] base = db.query(RelationModel).filter(RelationModel.status == "active") if direction in {"incoming", "both"}: result += [_to_relation(r) for r in base.filter(RelationModel.target_id == item_id).all()] if direction in {"outgoing", "both"}: result += [_to_relation(r) for r in base.filter(RelationModel.source_id == item_id).all()] 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