feat: make heatmap layers extensible via LayerRegistry (OCP)
This commit is contained in:
@@ -453,19 +453,71 @@ def build_campaign_layer(
|
||||
return layer
|
||||
|
||||
|
||||
# ── Layer dispatch (for Navigator export) ────────────────────────────
|
||||
# ── Layer registry (OCP-compliant dispatch) ──────────────────────────
|
||||
#
|
||||
# To add a new layer type:
|
||||
# 1. Write a builder function: ``def build_X_layer(db, *, platforms, tactics, min_score) -> dict``
|
||||
# 2. Call ``register_layer("x", build_X_layer)`` (or ``register_layer("x", fn, requires_id=True)``)
|
||||
# 3. Optionally add a convenience endpoint in the router
|
||||
#
|
||||
# The ``/export-navigator?layer=x`` endpoint picks up new layers automatically.
|
||||
|
||||
_LAYER_BUILDERS = {
|
||||
"coverage": lambda db, **kw: build_coverage_layer(db, **kw),
|
||||
"detection-rules": lambda db, **kw: build_detection_rules_layer(db, **kw),
|
||||
}
|
||||
|
||||
_LAYER_BUILDERS_WITH_ID = {
|
||||
"threat-actor": lambda db, lid, **kw: build_threat_actor_layer(db, lid, **kw),
|
||||
"campaign": lambda db, lid, **kw: build_campaign_layer(db, lid, **kw),
|
||||
}
|
||||
class _LayerRegistry:
|
||||
"""Extensible registry that maps layer type names to builder functions."""
|
||||
|
||||
SUPPORTED_LAYER_TYPES = set(_LAYER_BUILDERS) | set(_LAYER_BUILDERS_WITH_ID)
|
||||
__slots__ = ("_simple", "_with_id")
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._simple: dict[str, object] = {}
|
||||
self._with_id: dict[str, object] = {}
|
||||
|
||||
def register(self, name: str, builder, *, requires_id: bool = False) -> None:
|
||||
target = self._with_id if requires_id else self._simple
|
||||
target[name] = builder
|
||||
|
||||
@property
|
||||
def supported_types(self) -> set[str]:
|
||||
return set(self._simple) | set(self._with_id)
|
||||
|
||||
def build(
|
||||
self,
|
||||
db: Session,
|
||||
layer_type: str,
|
||||
*,
|
||||
layer_id: str | None = None,
|
||||
platforms: str | None = None,
|
||||
tactics: str | None = None,
|
||||
min_score: int = 0,
|
||||
) -> dict:
|
||||
kwargs = dict(platforms=platforms, tactics=tactics, min_score=min_score)
|
||||
|
||||
if layer_type in self._simple:
|
||||
return self._simple[layer_type](db, **kwargs)
|
||||
|
||||
if layer_type in self._with_id:
|
||||
if not layer_id:
|
||||
raise BusinessRuleViolation(
|
||||
f"layer_id is required for '{layer_type}' layer"
|
||||
)
|
||||
return self._with_id[layer_type](db, layer_id, **kwargs)
|
||||
|
||||
raise BusinessRuleViolation(f"Unknown layer type: {layer_type}")
|
||||
|
||||
|
||||
LAYER_REGISTRY = _LayerRegistry()
|
||||
|
||||
LAYER_REGISTRY.register("coverage", build_coverage_layer)
|
||||
LAYER_REGISTRY.register("detection-rules", build_detection_rules_layer)
|
||||
LAYER_REGISTRY.register("threat-actor", build_threat_actor_layer, requires_id=True)
|
||||
LAYER_REGISTRY.register("campaign", build_campaign_layer, requires_id=True)
|
||||
|
||||
SUPPORTED_LAYER_TYPES = LAYER_REGISTRY.supported_types # snapshot of built-in types
|
||||
|
||||
|
||||
def register_layer(name: str, builder, *, requires_id: bool = False) -> None:
|
||||
"""Public API to register a new heatmap layer type at import time."""
|
||||
LAYER_REGISTRY.register(name, builder, requires_id=requires_id)
|
||||
|
||||
|
||||
def build_navigator_export(
|
||||
@@ -484,16 +536,7 @@ def build_navigator_export(
|
||||
an entity-bound layer (threat-actor, campaign) references a
|
||||
non-existent record.
|
||||
"""
|
||||
kwargs = dict(platforms=platforms, tactics=tactics, min_score=min_score)
|
||||
|
||||
if layer_type in _LAYER_BUILDERS:
|
||||
return _LAYER_BUILDERS[layer_type](db, **kwargs)
|
||||
|
||||
if layer_type in _LAYER_BUILDERS_WITH_ID:
|
||||
if not layer_id:
|
||||
raise BusinessRuleViolation(
|
||||
f"layer_id is required for '{layer_type}' layer"
|
||||
)
|
||||
return _LAYER_BUILDERS_WITH_ID[layer_type](db, layer_id, **kwargs)
|
||||
|
||||
raise BusinessRuleViolation(f"Unknown layer type: {layer_type}")
|
||||
return LAYER_REGISTRY.build(
|
||||
db, layer_type,
|
||||
layer_id=layer_id, platforms=platforms, tactics=tactics, min_score=min_score,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user