graph_builder.py

"""Build a node/edge graph from an Azure Application Gateway / WAF JSON export.

The builder is intentionally schema-driven rather than hard-coded to specific
property names, because export files evolve between builds:

* Any object that carries a globally-unique identifier (an ARM resource id or
  a GUID) is promoted to a graph node. The "self id" key of an object is
  detected structurally: an ``id`` key, ``<SingularOfParentKey>Id``,
  ``ItemId``, or the id-like key whose values are unique across the sibling
  array elements.
* Everything that is not promoted stays attached to its nearest ancestor node
  as detail payload (rendered in the UI detail panel).
* Edges come from three sources:
  - containment (the JSON hierarchy),
  - shared values: any string anywhere in a node's details that equals
    another node's id, plus "external" nodes synthesized for ARM ids that
    are referenced by two or more nodes but described nowhere in the file
    (destination VNets, resource groups, sites, ...),
  - resource-path prefixes (ARM ids that extend another node's id path).
"""

import re

GUID_RE = re.compile(
    r"^[0-9a-f]{8}-(?:[0-9a-f]{4}-){3}[0-9a-f]{12}$", re.IGNORECASE
)
ARM_PREFIX = "/subscriptions/"

# The nil UUID is what GUID-scrubbing tools substitute for real ids. It
# identifies nothing: it must never become a node id and never match as a
# reference, otherwise every scrubbed id would "relate" to every other.
NIL_UUID = "00000000-0000-0000-0000-000000000000"

ROOT_ID = "__root__"


def _is_arm_id(value):
    if not isinstance(value, str):
        return False
    v = value.strip()
    return v.lower().startswith(ARM_PREFIX) and v.count("/") >= 2


def _is_guid(value):
    if not isinstance(value, str):
        return False
    v = value.strip().lower()
    return bool(GUID_RE.match(v)) and v != NIL_UUID


def _is_global_id(value):
    return _is_arm_id(value) or _is_guid(value)


def _norm(value):
    return value.strip().rstrip("/").lower()


def _singular(key):
    if not key:
        return key
    if key.endswith("ies") and len(key) > 3:
        return key[:-3] + "y"
    # double-s + es (addresses -> address, classes -> class, processes -> process)
    if key.endswith("sses") and len(key) > 4:
        return key[:-2]
    if key.endswith("s") and not key.endswith("ss"):
        return key[:-1]
    return key


def _last_segment(value):
    return value.strip().rstrip("/").rsplit("/", 1)[-1]


def _slug(text):
    return re.sub(r"[^a-z0-9]+", "-", text.lower()).strip("-")


def _tokens(text):
    """Split camelCase / snake-case into a set of lowercase word tokens."""
    spaced = re.sub(r"([a-z0-9])([A-Z])", r"\1 \2", text)
    return {t.lower() for t in re.findall(r"[A-Za-z][a-z0-9]*", spaced) if len(t) > 1}


# Generic qualifiers that carry no entity meaning in a *Name property.
NAME_STOP_TOKENS = {"name", "names", "id", "friendly", "display", "primary", "hydrated"}


def _classify_arm(raw):
    """Derive a (type, label) for an ARM id that has no backing object."""
    segments = [s for s in raw.strip().split("/") if s]
    label = segments[-1] if segments else raw
    if len(segments) >= 2:
        type_segment = segments[-2]
    else:
        type_segment = "Resource"
    type_name = _singular(type_segment)
    type_name = type_name[:1].upper() + type_name[1:]
    return type_name, label


class _Builder:
    def __init__(self):
        self.nodes = {}       # normalized id -> node dict
        self.node_order = []  # normalized ids in creation order
        self.edges = {}       # (source, target) -> edge dict

    # ------------------------------------------------------------- nodes

    def add_node(self, norm_id, raw_id, node_type, label, external=False):
        node = {
            "id": norm_id,
            "rawId": raw_id,
            "type": node_type,
            "label": label,
            "external": external,
            "details": None,
        }
        self.nodes[norm_id] = node
        self.node_order.append(norm_id)
        return node

    # ------------------------------------------------------------- edges

    def add_edge(self, source, target, kind, label=None):
        if source == target:
            return
        for key in ((source, target), (target, source)):
            existing = self.edges.get(key)
            if existing is not None:
                # Merge reference labels; never demote a containment edge.
                if (
                    kind == "ref"
                    and existing["kind"] == "ref"
                    and label
                    and label not in existing["labels"]
                ):
                    existing["labels"].append(label)
                return
        self.edges[(source, target)] = {
            "source": source,
            "target": target,
            "kind": kind,
            "labels": [label] if label else [],
        }

    def connected(self, a, b):
        return (a, b) in self.edges or (b, a) in self.edges

    # ---------------------------------------------------- self-id detection

    @staticmethod
    def _id_key_candidates(d):
        return [
            k
            for k, v in d.items()
            if isinstance(v, str) and v.strip() and k.lower().endswith("id")
        ]

    def _self_id_key(self, d, parent_key, sibling_unique):
        """Return (key, strong) for the object's own id, or None.

        ``strong`` marks the structural rules (<Singular>Id and the
        <Stem>Id/<Stem>Name pair) that are confident enough to promote the
        object to a node even when the id value itself is unusable
        (empty or scrubbed to the nil UUID)."""
        candidates = self._id_key_candidates(d)
        if not candidates:
            return None
        for k in candidates:
            if k.lower() == "id":
                return (k, False)
        singular = _singular(parent_key).lower() if parent_key else ""
        if singular:
            for k in candidates:
                if k.lower() == singular + "id":
                    return (k, True)
        # <Stem>Id with a matching <Stem>Name sibling marks self-identity,
        # but only when the stem relates to the parent key: otherwise a
        # dict that merely mentions another entity's id/name pair (e.g.
        # processServerId/processServerName inside providerSpecificDetails)
        # would wrongly become a node.
        parent_tokens = _tokens(_singular(parent_key)) if parent_key else set()
        for k in candidates:
            stem = k[:-2]
            if (
                stem
                and (_tokens(stem) & parent_tokens)
                and any(
                    sk.lower() == stem.lower() + "name"
                    and isinstance(sv, str)
                    and sv.strip()
                    for sk, sv in d.items()
                )
            ):
                return (k, True)
        for k in candidates:
            if k.lower() == "itemid":
                return (k, False)
        if sibling_unique:
            for k in candidates:
                if k in sibling_unique and _is_global_id(d[k]):
                    return (k, False)
        return None

    @staticmethod
    def _unique_id_keys(elements):
        """Id-like keys whose values are distinct across all sibling dicts."""
        values = {}
        for element in elements:
            for k, v in element.items():
                if isinstance(v, str) and v.strip() and k.lower().endswith("id"):
                    values.setdefault(k, []).append(_norm(v))
        return {
            k
            for k, vs in values.items()
            if len(vs) == len(elements) and len(set(vs)) == len(vs)
        }

    # ----------------------------------------------------------- labeling

    def _find_label(self, d, parent_key):
        """Best display name in the subtree: displayName > friendlyName >
        <Singular>Name > name > any *Name, nearest-first."""
        singular = _singular(parent_key).lower() if parent_key else ""
        best = None  # (priority, depth, sequence, value)
        queue = [(d, 0)]
        sequence = 0
        while queue:
            current, depth = queue.pop(0)
            for k, v in current.items():
                sequence += 1
                if isinstance(v, str) and v.strip():
                    kl = k.lower()
                    if kl == "displayname":
                        priority = 0
                    elif kl == "friendlyname":
                        priority = 1
                    elif singular and kl == singular + "name":
                        priority = 2
                    elif kl == "name":
                        priority = 3
                    elif kl.endswith("name"):
                        priority = 4
                    else:
                        continue
                    candidate = (priority, depth, sequence, v.strip())
                    if best is None or candidate[:3] < best[:3]:
                        best = candidate
                elif isinstance(v, dict) and depth < 3:
                    queue.append((v, depth + 1))
        return best[3] if best else None

    # ------------------------------------------------------------- walking

    def prune(self, obj, parent_node_id, key, sibling_unique=None):
        if isinstance(obj, dict):
            return self.prune_dict(obj, parent_node_id, key, sibling_unique)
        if isinstance(obj, list):
            if obj and all(isinstance(el, dict) for el in obj):
                unique = self._unique_id_keys(obj)
                return [
                    self.prune_dict(el, parent_node_id, key, unique)
                    for el in obj
                ]
            return [self.prune(el, parent_node_id, key) for el in obj]
        return obj

    def prune_dict(self, d, parent_node_id, key, sibling_unique=None):
        found = self._self_id_key(d, key, sibling_unique)
        if found:
            id_key, strong = found
            raw = d[id_key].strip()
            if _is_global_id(raw):
                norm_id = _norm(raw)
                if norm_id == parent_node_id:
                    # The same entity restated inside itself (e.g. a
                    # SourceResource whose id equals the item's ItemId):
                    # keep it inline as detail payload.
                    return {
                        k: self.prune(v, parent_node_id, k)
                        for k, v in d.items()
                    }
                if norm_id in self.nodes:
                    # A second description of an already-known entity.
                    self.add_edge(parent_node_id, norm_id, "ref", key)
                    return {"__node__": norm_id}
                node_type = _singular(key) if key else "Item"
                node = self.add_node(norm_id, raw, node_type, None)
                self.add_edge(parent_node_id, norm_id, "contains", key)
                node["details"] = {
                    k: self.prune(v, norm_id, k) for k, v in d.items()
                }
                node["label"] = self._find_label(d, key) or _last_segment(raw)
                return {"__node__": norm_id}
            if strong:
                # Structurally an entity, but its id was scrubbed (nil
                # UUID) or left empty. Promote it anyway under a synthetic
                # id derived from its position and name, so e.g. the
                # ReplicationAppliances under a vault stay first-class
                # nodes in anonymized exports.
                label = self._find_label(d, key)
                if label:
                    node_type = _singular(key) if key else "Item"
                    node_id = "{}|{}|{}".format(
                        parent_node_id, _slug(node_type), _slug(label)
                    )
                    if node_id in self.nodes:
                        self.add_edge(parent_node_id, node_id, "ref", key)
                        return {"__node__": node_id}
                    node = self.add_node(
                        node_id,
                        raw or "(no globally-unique id in file)",
                        node_type,
                        label,
                    )
                    self.add_edge(parent_node_id, node_id, "contains", key)
                    node["details"] = {
                        k: self.prune(v, node_id, k) for k, v in d.items()
                    }
                    return {"__node__": node_id}
        return {k: self.prune(v, parent_node_id, k) for k, v in d.items()}

    # ----------------------------------------------------- reference edges

    def _scan_strings(self, obj, node_id, key, out):
        if isinstance(obj, dict):
            if "__node__" in obj and len(obj) == 1:
                return
            for k, v in obj.items():
                self._scan_strings(v, node_id, k, out)
        elif isinstance(obj, list):
            for v in obj:
                self._scan_strings(v, node_id, key, out)
        elif isinstance(obj, str):
            s = obj.strip()
            if _is_global_id(s):
                out.append((node_id, key or "value", _norm(s), s))

    def add_reference_edges(self):
        refs = []
        for node_id in list(self.node_order):
            self._scan_strings(self.nodes[node_id]["details"], node_id, None, refs)

        unresolved = {}
        for node_id, key, norm_value, raw in refs:
            if norm_value == node_id:
                continue
            if norm_value in self.nodes:
                self.add_edge(node_id, norm_value, "ref", key)
            elif _is_arm_id(raw):
                # Paths ending in "None" or the nil UUID identify nothing
                # (the nil tail means the id was scrubbed, so distinct
                # resources would collapse into one misleading node).
                last = _last_segment(raw).lower()
                if last != "none" and last != NIL_UUID:
                    entry = unresolved.setdefault(
                        norm_value, {"raw": raw, "sources": {}}
                    )
                    entry["sources"].setdefault(node_id, set()).add(key)

        # An ARM id described nowhere in the file becomes an "external"
        # node only when two or more nodes share it - shared values are
        # the relations worth showing; one-off ids are just noise.
        for norm_value, entry in unresolved.items():
            if len(entry["sources"]) < 2:
                continue
            node_type, label = _classify_arm(entry["raw"])
            self.add_node(norm_value, entry["raw"], node_type, label, external=True)
            for source, keys in entry["sources"].items():
                for k in sorted(keys):
                    self.add_edge(source, norm_value, "ref", k)

    # ------------------------------------------------- name-based ref edges

    def _scan_names(self, obj, key, out):
        if isinstance(obj, dict):
            if "__node__" in obj and len(obj) == 1:
                return
            for k, v in obj.items():
                self._scan_names(v, k, out)
        elif isinstance(obj, list):
            for v in obj:
                self._scan_names(v, key, out)
        elif isinstance(obj, str) and key:
            kl = key.lower()
            if (kl.endswith("name") or kl.endswith("names")) and obj.strip():
                out.append((key, obj))

    def add_name_reference_edges(self):
        """Link nodes through *Name properties whose value equals another
        node's label. This is what still connects a MigrationItem to its
        ReplicationAppliance (applianceNames / processServerName) when the
        GUID ids have been scrubbed and can no longer match. Guarded by:
        the label must belong to exactly one node, and the property name
        must share a meaningful token with the target's type or keys."""
        label_owners = {}
        for node_id in self.node_order:
            label = (self.nodes[node_id]["label"] or "").strip().lower()
            if label:
                label_owners.setdefault(label, []).append(node_id)

        for node_id in list(self.node_order):
            found = []
            self._scan_names(self.nodes[node_id]["details"], None, found)
            for key, value in found:
                owners = label_owners.get(value.strip().lower())
                if not owners or len(owners) != 1 or owners[0] == node_id:
                    continue
                key_tokens = _tokens(key) - NAME_STOP_TOKENS
                if not key_tokens:
                    continue
                target = self.nodes[owners[0]]
                target_tokens = _tokens(target["type"])
                if isinstance(target["details"], dict):
                    for k in target["details"]:
                        target_tokens |= _tokens(k)
                if key_tokens & target_tokens:
                    self.add_edge(node_id, owners[0], "ref", key)

    # ---------------------------------------------------------- path edges

    def add_path_edges(self):
        arm_ids = [
            node_id
            for node_id in self.node_order
            if node_id.startswith(ARM_PREFIX)
        ]
        for node_id in arm_ids:
            best = None
            for other in arm_ids:
                if other != node_id and node_id.startswith(other + "/"):
                    if best is None or len(other) > len(best):
                        best = other
            if best and not self.connected(node_id, best):
                self.add_edge(best, node_id, "path", "resource path")

    # --------------------------------------------------------------- output

    def result(self, source_name):
        types = []
        for node_id in self.node_order:
            t = self.nodes[node_id]["type"]
            if t not in types:
                types.append(t)
        return {
            "source": source_name,
            "types": types,
            "nodes": [self.nodes[node_id] for node_id in self.node_order],
            "edges": list(self.edges.values()),
        }


def build_graph(data, source_name="JSON export"):
    builder = _Builder()

    if isinstance(data, dict):
        root_source = data
    elif isinstance(data, list):
        root_source = {"items": data}
    else:
        root_source = {"value": data}

    label = builder._find_label(root_source, "Export") or source_name
    root = builder.add_node(ROOT_ID, ROOT_ID, "Export", label)
    root["details"] = {
        k: builder.prune(v, ROOT_ID, k) for k, v in root_source.items()
    }

    builder.add_reference_edges()
    builder.add_name_reference_edges()
    builder.add_path_edges()
    return builder.result(source_name)