Files

132 lines
4.8 KiB
Python
Raw Permalink Normal View History

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