execution_backends/function_call.py

"""Function-Call backend — the advisory/reasoning execution backend.

WHAT THIS IS
------------
This backend stands in for the "tool calling" that an agent framework normally
does with plain local functions. Where the Direct-API backend *changes things*
(it can create a ticket), this backend only **thinks out loud**: it looks things
up, scores risk, and drafts text for a human to read. It produces *analysis and
recommendations* — never an action against a real system.

    * LOOK-UPS / SCORING — cve_lookup, risk_score      (deterministic facts)
    * DRAFTS             — draft_remediation,           (generated text, for a
                           draft_incident                human to review)
    * RECOMMENDATIONS    — recommend_containment        (suggested next steps;
                                                         it suggests, never does)

THE ONE PROPERTY THAT DEFINES THIS BACKEND: NO SIDE EFFECTS
-----------------------------------------------------------
Every handler here is **pure**: given the same arguments it returns the same
result, and it changes *nothing* — not the mock data, not any external system,
not even an internal store (unlike Direct-API, this backend has no writable
store at all). That is what lets the policy allow these actions at the low
"recommend-only" autonomy tier (L1): a recommendation to contain a host is just
words until a human acts on it. "recommend_containment" returns a *plan*; it
does not contain anything.

WHERE IT SITS IN THE TRUST MODEL (same as every backend)
--------------------------------------------------------
This backend never decides *whether* an action is allowed — that already
happened upstream. By the time ``run`` is called, the gate has said ``ALLOW``
and the dispatcher has verified the gate's signed grant.

    gate decides ─► dispatcher verifies grant ─► THIS backend executes

It still keeps the same defense-in-depth guard as the other backends: an
**allow-list of handlers**. An action this backend doesn't know how to perform
(anything not in the table below) is refused, not guessed at.
"""

from __future__ import annotations

from typing import Any, Callable

import data
from control_plane.schema import ProposedAction


class FunctionBackendError(Exception):
    """Base error for anything this backend refuses or cannot do."""


class FunctionUnsupportedActionError(FunctionBackendError):
    """The requested action is not in this backend's allow-list of handlers."""


class FunctionRecordNotFoundError(FunctionBackendError):
    """A look-up referenced an ID that does not exist in the mock data."""


class FunctionMissingArgumentError(FunctionBackendError):
    """A required argument was absent from the action's ``arguments`` map."""


# How much an asset's business criticality amplifies a raw vulnerability score.
# Kept as a fixed table (not a formula) so risk scores are deterministic and easy
# to explain: the same CVE on a more critical asset scores higher, by a known amount.
_CRITICALITY_WEIGHT: dict[str, float] = {
    "low": 0.6,
    "medium": 0.8,
    "high": 1.0,
}


class FunctionCallBackend:
    """Advisory backend satisfying the dispatcher's ``BackendAdapter`` contract.

    Construct it once and let the dispatcher call :meth:`run` per authorized
    action. It holds only **read-only** references to the mock data — there is
    deliberately no mutable state here, because none of its actions may change
    anything.
    """

    def __init__(self) -> None:
        # Read-only lookups by ID. Indexed once at construction for O(1) access.
        # Nothing in this class ever writes back to these.
        self._cves = data.index_by_id("cves")
        self._assets = data.index_by_id("assets")
        self._alerts = data.load_alerts()  # scanned by incident_id when drafting

        # The allow-list: action_name -> handler. Membership here is what makes an
        # action executable by this backend; the names match the policy exactly
        # (policies/agt_policy.yaml, function_call section). Anything not listed is
        # refused by :meth:`run`.
        self._handlers: dict[str, Callable[[dict[str, Any]], Any]] = {
            # deterministic facts
            "cve_lookup": self._cve_lookup,
            "risk_score": self._risk_score,
            # generated text for a human to review
            "draft_remediation": self._draft_remediation,
            "draft_incident": self._draft_incident,
            # suggested next steps (advisory only)
            "recommend_containment": self._recommend_containment,
        }

    # -- dispatcher entry point ------------------------------------------------

    def run(self, action: ProposedAction) -> Any:
        """Execute *action* and return its advisory result.

        Routes on ``action.action_name`` to a handler in the allow-list. An action
        this backend doesn't handle raises :class:`FunctionUnsupportedActionError` — it is
        never silently ignored or guessed at.
        """
        handler = self._handlers.get(action.action_name)
        if handler is None:
            raise FunctionUnsupportedActionError(
                f"function_call backend cannot perform {action.action_name!r} "
                f"(not in its allow-list of advisory actions)"
            )
        return handler(action.arguments)

    # -- look-ups / scoring (deterministic facts) ------------------------------

    def _cve_lookup(self, args: dict[str, Any]) -> dict[str, Any]:
        """Return the mock CVE record for ``cve_id`` (a pure read of the fixtures)."""
        cve_id = _require(args, "cve_id")
        return _found(self._cves.get(cve_id), "cve", cve_id)

    def _risk_score(self, args: dict[str, Any]) -> dict[str, Any]:
        """Score a CVE, optionally amplified by the criticality of a named asset.

        The score is deterministic: it combines the CVE's published CVSS score
        with a fixed per-asset weight (see :data:`_CRITICALITY_WEIGHT`). The same
        inputs always yield the same score and band, which is what makes this
        usable as an advisory signal a human can sanity-check.
        """
        cve = self._cve_lookup(args)  # reuses the look-up (and its not-found guard)
        cvss = float(cve.get("cvss_score", 0.0))

        # Asset is optional: if given, a more business-critical asset raises the
        # effective risk; if absent, we score the raw vulnerability on its own.
        asset_id = args.get("asset_id")
        if asset_id:
            asset = _found(self._assets.get(asset_id), "asset", asset_id)
            criticality = asset.get("criticality", "medium")
            weight = _CRITICALITY_WEIGHT.get(criticality, 0.8)
        else:
            criticality, weight = None, 1.0

        score = round(cvss * weight, 1)
        return {
            "cve_id": cve["cve_id"],
            "asset_id": asset_id,
            "cvss_score": cvss,
            "asset_criticality": criticality,
            "risk_score": score,
            "risk_band": _band(score),
            "rationale": (
                f"CVSS {cvss} weighted by asset criticality "
                f"{criticality or 'n/a'} (×{weight}) → {score}"
            ),
        }

    # -- drafts (generated text, no side effects) ------------------------------

    def _draft_remediation(self, args: dict[str, Any]) -> dict[str, Any]:
        """Draft a remediation write-up for a CVE — text only, nothing is applied.

        This is the "L1 draft" action. It assembles human-readable
        guidance from the mock CVE record; it does **not** patch, restart, or
        change anything. The ``is_draft`` flag makes that explicit to any consumer.
        """
        cve = self._cve_lookup(args)
        content = (
            f"Remediation draft for {cve['cve_id']} — {cve['title']}\n"
            f"Severity: {cve.get('severity', 'unknown')} "
            f"(CVSS {cve.get('cvss_score', 'n/a')})\n\n"
            f"Summary: {cve.get('summary', '')}\n\n"
            f"Recommended steps: {cve.get('remediation', 'No guidance on file.')}\n\n"
            f"NOTE: This is a draft for analyst review. No change has been applied."
        )
        return {
            "kind": "remediation_draft",
            "cve_id": cve["cve_id"],
            "is_draft": True,
            "content": content,
        }

    def _draft_incident(self, args: dict[str, Any]) -> dict[str, Any]:
        """Draft an incident summary from the alerts attached to an incident_id.

        Pulls every mock alert sharing the incident_id and folds them into a short
        narrative a human could paste into an incident record. Read-only: it never
        creates or edits the incident itself.
        """
        incident_id = _require(args, "incident_id")
        alerts = [a for a in self._alerts if a.get("incident_id") == incident_id]
        if not alerts:
            raise FunctionRecordNotFoundError(f"no alerts found for incident {incident_id!r}")

        lines = [f"Incident draft for {incident_id}", ""]
        for alert in alerts:
            lines.append(
                f"- [{alert.get('severity', '?')}] {alert['alert_id']}: "
                f"{alert.get('title', '')} (asset {alert.get('asset_id', 'n/a')})"
            )
        lines += ["", "NOTE: Draft for analyst review. The incident record is unchanged."]
        return {
            "kind": "incident_draft",
            "incident_id": incident_id,
            "alert_ids": [a["alert_id"] for a in alerts],
            "is_draft": True,
            "content": "\n".join(lines),
        }

    def _recommend_containment(self, args: dict[str, Any]) -> dict[str, Any]:
        """Recommend containment steps for an asset — a suggestion, not an action.

        Returns an ordered list of *suggested* steps (e.g. isolate the host, stop a
        suspicious container). Crucially this backend cannot perform any of them;
        executing a containment step is a high-tier action that would have to go
        back through the gate on a separate, approval-gated request.
        """
        asset_id = _require(args, "asset_id")
        asset = _found(self._assets.get(asset_id), "asset", asset_id)

        steps = [f"Isolate host {asset.get('hostname', asset_id)} from the network."]
        # If a specific container was named, recommend stopping just that one;
        # otherwise call out any container the mock data already flags suspicious.
        container_id = args.get("container_id")
        suspicious = [
            c for c in asset.get("containers", [])
            if c.get("suspicious") and (container_id is None or c["container_id"] == container_id)
        ]
        for c in suspicious:
            steps.append(f"Stop suspicious container {c['name']} ({c['container_id']}).")
        steps.append("Preserve volatile evidence (memory, logs) before remediation.")

        return {
            "kind": "containment_recommendation",
            "asset_id": asset_id,
            "is_recommendation": True,  # advisory only — nothing here is executed
            "steps": steps,
        }


# Module-level helpers keep the handlers above short and uniform. (They mirror the
# small guards in the Direct-API backend; each backend keeps its own copy so the
# backends stay independent of one another.)


def _require(args: dict[str, Any], key: str) -> Any:
    """Return ``args[key]`` or raise a clear error naming the missing argument."""
    if key not in args or args[key] in (None, ""):
        raise FunctionMissingArgumentError(f"missing required argument {key!r}")
    return args[key]


def _found(record: dict[str, Any] | None, kind: str, record_id: str) -> dict[str, Any]:
    """Return *record* if present, else raise a clear not-found error."""
    if record is None:
        raise FunctionRecordNotFoundError(f"no {kind} with id {record_id!r}")
    return record


def _band(score: float) -> str:
    """Map a 0–10 risk score to a human-readable band (fixed, deterministic cut-offs)."""
    if score >= 9.0:
        return "critical"
    if score >= 7.0:
        return "high"
    if score >= 4.0:
        return "medium"
    return "low"