feat(graph): evidence graph node and edge population

alembic/versions/e2c4a6b8d013_graph_current_scope.py

@@ -0,0 +1,48 @@
+"""Track the explicit current snapshot for each tenant/subscription graph.
+
+Revision ID: e2c4a6b8d013
+Revises: e1f2a3b4c5d6
+"""
+
+from typing import Sequence, Union
+from alembic import op
+
+revision: str = "e2c4a6b8d013"
+down_revision: str = "e1f2a3b4c5d6"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    op.execute("""
+        CREATE TABLE graph_snapshot_scopes (
+            tenant_id text NOT NULL,
+            subscription_id text NOT NULL,
+            snapshot_id text NOT NULL,
+            status text NOT NULL CHECK (status IN ('COMPLETE', 'PARTIAL')),
+            collected_at timestamptz NOT NULL,
+            PRIMARY KEY (tenant_id, subscription_id)
+        )
+    """)
+    op.execute("""
+        CREATE VIEW current_graph_nodes AS
+        SELECT n.* FROM graph_nodes n
+        JOIN graph_snapshot_scopes s
+          ON s.tenant_id = n.tenant_id AND s.subscription_id = n.subscription_id
+         AND s.snapshot_id = n.snapshot_id
+    """)
+    op.execute("""
+        CREATE VIEW current_graph_edges AS
+        SELECT e.* FROM graph_edges e
+        JOIN current_graph_nodes src ON src.node_id = e.source_node_id
+        JOIN current_graph_nodes tgt ON tgt.node_id = e.target_node_id
+        WHERE src.tenant_id = tgt.tenant_id
+          AND e.evidence_snapshot_id = src.snapshot_id
+          AND e.evidence_snapshot_id = tgt.snapshot_id
+    """)
+
+
+def downgrade() -> None:
+    op.execute("DROP VIEW current_graph_edges")
+    op.execute("DROP VIEW current_graph_nodes")
+    op.execute("DROP TABLE graph_snapshot_scopes")

scanner/arg_inventory.py

@@ -17,6 +17,7 @@
 
 DEFAULT_QUERY = """
 Resources
+| extend properties = bag_merge(properties, pack('identity', identity))
 | project id, name, type, location, subscriptionId, resourceGroup, tenantId, tags, properties
 | order by id asc
 """.strip()

scanner/engine.py

@@ -61,6 +61,7 @@ def __init__(self, subscription_id: str) -> None:
         self.subscription_id = subscription_id
         self.client = AzureClient(subscription_id)
         self.rules: List[Any] = []
+        self.snapshot: Optional[Any] = None
         self.load_rules()
 
     # ------------------------------------------------------------------ #
@@ -118,6 +119,7 @@ def run_scan(self, scan_id: Optional[str] = None) -> Dict[str, Any]:
         # Collect an ARG inventory snapshot for graph population and rule enrichment.
         # Failure is non-fatal: rules fall back to direct SDK calls.
         snapshot = collect_snapshot(self.client, self.subscription_id)
+        self.snapshot = snapshot
 
         logger.info(
             "Scan %s starting against subscription %s — %d rules loaded",

scanner/graph/edge_detector.py

@@ -0,0 +1,197 @@
+"""Typed edge detectors that infer relationships between Azure resources in an InventorySnapshot."""
+
+from __future__ import annotations
+
+import logging
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from scanner.arg_inventory import InventorySnapshot
+
+logger = logging.getLogger(__name__)
+
+_NSG_TYPE = "microsoft.network/networksecuritygroups"
+_PUBLIC_IP_TYPE = "microsoft.network/publicipaddresses"
+_STORAGE_TYPE = "microsoft.storage/storageaccounts"
+
+
+@dataclass
+class GraphEdge:
+    """A directed relationship between two Azure resources."""
+
+    source_resource_id: str
+    target_resource_id: str
+    relationship_type: str
+    evidence_source: str
+    confidence: float
+
+
+class EdgeDetector(ABC):
+    """Base class for relationship detectors."""
+
+    @abstractmethod
+    def detect(self, snapshot: InventorySnapshot) -> list[GraphEdge]:
+        """Return edges inferred from this snapshot."""
+
+
+class NsgToSubnetDetector(EdgeDetector):
+    """NSG -> Subnet: PROTECTS (ARG-confirmed, confidence 1.0)."""
+
+    def detect(self, snapshot: InventorySnapshot) -> list[GraphEdge]:
+        edges = []
+        for resource in snapshot.resources:
+            if resource.resource_type.lower() != _NSG_TYPE:
+                continue
+            subnets = resource.properties.get("subnets") or []
+            for subnet in subnets:
+                subnet_id = subnet.get("id") if isinstance(subnet, dict) else None
+                if not subnet_id:
+                    continue
+                edges.append(
+                    GraphEdge(
+                        source_resource_id=resource.resource_id,
+                        target_resource_id=subnet_id,
+                        relationship_type="PROTECTS",
+                        evidence_source="arg:properties.subnets",
+                        confidence=1.0,
+                    )
+                )
+        return edges
+
+
+class SubnetToResourceDetector(EdgeDetector):
+    """Resource -> Subnet: MEMBER_OF (inferred from properties, confidence 0.8)."""
+
+    def detect(self, snapshot: InventorySnapshot) -> list[GraphEdge]:
+        edges = []
+        resource_ids = {r.resource_id.lower() for r in snapshot.resources}
+        for resource in snapshot.resources:
+            subnet_ids: list[str] = []
+
+            # Direct subnet property (VMs attached directly)
+            direct = (resource.properties.get("subnet") or {}).get("id")
+            if direct:
+                subnet_ids.append(direct)
+
+            # NIC ipConfigurations pattern (most common path for NICs)
+            for ip_cfg in resource.properties.get("ipConfigurations", []):
+                sid = (ip_cfg.get("properties", {}).get("subnet") or {}).get("id")
+                if sid:
+                    subnet_ids.append(sid)
+
+            for subnet_id in subnet_ids:
+                if not subnet_id or subnet_id.lower() not in resource_ids:
+                    continue
+                edges.append(
+                    GraphEdge(
+                        source_resource_id=resource.resource_id,
+                        target_resource_id=subnet_id,
+                        relationship_type="MEMBER_OF",
+                        evidence_source="arg:properties.subnet.id",
+                        confidence=0.8,
+                    )
+                )
+        return edges
+
+
+class PublicIpToResourceDetector(EdgeDetector):
+    """PublicIP -> Resource: EXPOSES (ARG-confirmed, confidence 1.0)."""
+
+    def detect(self, snapshot: InventorySnapshot) -> list[GraphEdge]:
+        edges = []
+        for resource in snapshot.resources:
+            if resource.resource_type.lower() != _PUBLIC_IP_TYPE:
+                continue
+            ip_config = resource.properties.get("ipConfiguration") or {}
+            target_id = ip_config.get("id")
+            if not target_id:
+                continue
+            # Trim to the NIC resource ID (namespace/type/name = 3 segments).
+            parts = target_id.split("/providers/")
+            if len(parts) >= 2:
+                provider_path = parts[-1].split("/")
+                if len(provider_path) >= 3:
+                    target_id = "/providers/".join(parts[:-1]) + "/providers/" + "/".join(provider_path[:3])
+            edges.append(
+                GraphEdge(
+                    source_resource_id=resource.resource_id,
+                    target_resource_id=target_id,
+                    relationship_type="EXPOSES",
+                    evidence_source="arg:properties.ipConfiguration.id",
+                    confidence=1.0,
+                )
+            )
+        return edges
+
+
+class IdentityToResourceDetector(EdgeDetector):
+    """Identity -> Resource: HAS_IDENTITY (inferred from properties, confidence 0.8)."""
+
+    def detect(self, snapshot: InventorySnapshot) -> list[GraphEdge]:
+        edges = []
+        for resource in snapshot.resources:
+            identity = resource.properties.get("identity") or {}
+            user_assigned = identity.get("userAssignedIdentities") or {}
+            for identity_id in user_assigned:
+                if not identity_id:
+                    continue
+                edges.append(
+                    GraphEdge(
+                        source_resource_id=identity_id,
+                        target_resource_id=resource.resource_id,
+                        relationship_type="HAS_IDENTITY",
+                        evidence_source="arg:properties.identity.userAssignedIdentities",
+                        confidence=0.8,
+                    )
+                )
+        return edges
+
+
+class StoragePrivateEndpointDetector(EdgeDetector):
+    """Storage -> PrivateEndpoint: REACHABLE_VIA (ARG-confirmed, confidence 1.0)."""
+
+    def detect(self, snapshot: InventorySnapshot) -> list[GraphEdge]:
+        edges = []
+        for resource in snapshot.resources:
+            if resource.resource_type.lower() != _STORAGE_TYPE:
+                continue
+            connections = resource.properties.get("privateEndpointConnections") or []
+            for conn in connections:
+                pe_id = ((conn.get("properties") or {}).get("privateEndpoint") or {}).get("id")
+                if not pe_id:
+                    continue
+                edges.append(
+                    GraphEdge(
+                        source_resource_id=resource.resource_id,
+                        target_resource_id=pe_id,
+                        relationship_type="REACHABLE_VIA",
+                        evidence_source="arg:properties.privateEndpointConnections",
+                        confidence=1.0,
+                    )
+                )
+        return edges
+
+
+def detect_all_edges(snapshot: InventorySnapshot) -> list[GraphEdge]:
+    """Run all detectors and return the combined edge list.
+
+    A detector failure aborts graph publication. Publishing a partial edge list
+    as complete would remove valid relationships from the preceding snapshot.
+    """
+    detectors: list[EdgeDetector] = [
+        NsgToSubnetDetector(),
+        SubnetToResourceDetector(),
+        PublicIpToResourceDetector(),
+        IdentityToResourceDetector(),
+        StoragePrivateEndpointDetector(),
+    ]
+    edges: list[GraphEdge] = []
+    for detector in detectors:
+        try:
+            edges.extend(detector.detect(snapshot))
+        except Exception as exc:
+            logger.warning("edge_detector: %s failed: %s", type(detector).__name__, exc)
+            raise
+    return edges