diff --git a/app/capabilities/cti_context.py b/app/capabilities/cti_context.py new file mode 100644 index 0000000..04bd212 --- /dev/null +++ b/app/capabilities/cti_context.py @@ -0,0 +1,206 @@ +"""CTI grounding: hand the agent the intel the operator attached. + +A bundle from this pipeline is small. One report yields tens of +attack-patterns, not thousands, so the whole of it fits inside a modern +context window several times over. Embedding it to retrieve a handful of +chunks discarded the rest for no gain, and the query it retrieved against was +the operator's raw prompt, which rarely contains technique language, so the +selection was close to arbitrary. Published long-context evaluations put the +point where retrieval starts to pay at the reader's effective context window, +far above a single report. + +The default is therefore everything, in document order. A selection that +outgrows the deployment's budget is trimmed rather than sampled, and the +prompt says what was left out so the agent can call a CTI tool for the rest. +The budget lives on the `cti` profile in conf/default.yml; nothing here +decides how much a deployment can afford. +""" + +from __future__ import annotations + +import json +import logging +from dataclasses import dataclass +from pathlib import Path + +from plugins.mcp.app.capabilities.base import Capability +from plugins.mcp.app.config import profile_defaults +from plugins.mcp.app.utilities.paths import get_mcp_data_dir + +log = logging.getLogger("plugins.mcp") + +# The two object types a bundle carries that say something about adversary +# behaviour. Everything else in STIX 2.1 describes the previous victim's +# estate, which is not ours to act on. +_INTEL_TYPES = {"attack-pattern": "technique", "threat-actor": "actor"} + + +@dataclass(frozen=True) +class IntelObject: + """One grounded fact from a bundle, ready to render into a prompt.""" + + kind: str + attack_id: str + name: str + description: str + evidence: str + + def render(self) -> str: + heading = " ".join(part for part in (self.attack_id, self.name) if part) + lines = [f"### {heading or self.kind}"] + # The report's own words first: they are what distinguishes this report + # from every other one citing the same technique. + if self.evidence: + lines.append(f"Observed in report: {self.evidence}") + if self.description: + lines.append(self.description) + return "\n".join(lines) + + +def _attack_id(obj: dict) -> str: + for ref in obj.get("external_references") or []: + if ref.get("source_name") == "mitre-attack" and ref.get("external_id"): + return str(ref["external_id"]).strip() + return "" + + +def read_intel(bundle: dict) -> list[IntelObject]: + """Every behavioural object in a bundle, in the order the bundle lists them.""" + intel: list[IntelObject] = [] + for obj in bundle.get("objects") or []: + kind = _INTEL_TYPES.get(obj.get("type")) + if kind is None: + continue + name = (obj.get("name") or "").strip() + description = (obj.get("description") or "").strip() + if not name and not description: + continue + intel.append( + IntelObject( + kind=kind, + attack_id=_attack_id(obj), + name=name, + description=description, + evidence=(obj.get("x_cti_evidence") or "").strip(), + ) + ) + return intel + + +def estimate_tokens(text: str) -> int: + """Character-based proxy. The budget only needs to be right to the nearest + object, and this avoids a tokeniser that would have to match whichever + model the deployment points at.""" + return len(text) // 4 + + +def render(intel: list[IntelObject], *, budget_tokens: int) -> str: + """Render as much intel as the budget allows, in document order. + + Objects are kept whole: half a technique reads like a complete one and + would ground the agent in a truncated sentence. + + Trimming takes from the end, so with several bundles attached the last one + loses most. That is fine for the single-report case this is built for; if + operators start attaching many at once, spread the budget across bundles + instead. + """ + if not intel: + return "" + + kept: list[str] = [] + used = 0 + for obj in intel: + block = obj.render() + cost = estimate_tokens(block) + if kept and used + cost > budget_tokens: + break + kept.append(block) + used += cost + + body = "\n\n".join(kept) + omitted = len(intel) - len(kept) + if omitted: + # An agent that knows its intel is partial can call a CTI tool for the + # rest. One that does not will reason from a silent gap. + body += ( + f"\n\n[{omitted} further object(s) omitted to stay within the " + f"context budget. Use the CTI tools to read the full bundle.]" + ) + return body + + +def _resolve_bundle(name: str) -> Path: + """Resolve an attached filename to a bundle on disk. + + Pipeline output is searched first, so a bundle regenerated by a later run + wins over an older copy of the same name uploaded by hand. + """ + raw = Path(str(name)) + if raw.is_absolute(): + return raw + + data_dir = get_mcp_data_dir() + candidates = [data_dir / "outputs_stix" / raw] if len(raw.parts) == 1 else [] + candidates.append(data_dir / raw) + for path in candidates: + if path.exists(): + return path + return candidates[0] + + +def _context_budget() -> int: + return int(profile_defaults("cti")["context_budget_tokens"]) + + +async def _enrich(prompt: str, settings: dict) -> dict: + """Read the attached bundles and return them as one cti_context string. + + Returns {} when nothing is attached, so the workflow falls through to its + ungrounded signature rather than being handed an empty string that reads + like intelligence. + + Raises when a bundle cannot be read or carries no behavioural objects. The + operator attached it deliberately; a run that silently proceeds without it + looks identical to one that worked. + """ + attached = settings.get("cti_files") or [] + if not attached: + return {} + + intel: list[IntelObject] = [] + for name in attached: + path = _resolve_bundle(name) + if not path.exists(): + raise FileNotFoundError(f"attached CTI bundle not found: {path}") + with open(path, "r", encoding="utf-8") as handle: + bundle = json.load(handle) + found = read_intel(bundle) + if not found: + raise ValueError( + f"{path.name} carries no attack-pattern or threat-actor " + f"objects, so there is nothing to ground on" + ) + intel.extend(found) + + context = render(intel, budget_tokens=_context_budget()) + log.info( + "[cti] grounding on %d object(s) from %d bundle(s), ~%d tokens", + len(intel), len(attached), estimate_tokens(context), + ) + return {"cti_context": context} + + +CAPABILITIES = [ + Capability( + id="cti", + display_name="CTI grounding", + description=( + "Put the attached intel in front of the agent: every technique and " + "actor the bundle carries, with the report's own wording where the " + "pipeline captured it." + ), + enrich=_enrich, + contributes_fields=["cti_context"], + ), +] diff --git a/app/capabilities/rag.py b/app/capabilities/rag.py deleted file mode 100644 index f7f6ee6..0000000 --- a/app/capabilities/rag.py +++ /dev/null @@ -1,317 +0,0 @@ -import dspy -import json -from typing import List, Dict, Optional -import logging - -from plugins.mcp.app.dspy_env import apply_litellm_ssl_verify - -class RAGService: - """RAG service for CTI (Cyber Threat Intelligence) data retrieval using STIX bundles.""" - - def __init__( - self, - stix_bundle_path: Optional[str] = None, - api_key: Optional[str] = None, - api_base: Optional[str] = None, - ssl_verify: Optional[bool] = None, - log: Optional[logging.Logger] = None, - ): - self.max_characters = 6000 - self.topk_objects_to_retrieve = 5 - self.corpus = [] - self.adv_step = {} - self.search = None - # Embedding-side credentials. Default to the chat LLM key/base when - # the orchestrator does not supply a separate embedding provider. - self.api_key = api_key - self.api_base = api_base - self.ssl_verify = ssl_verify - self.log = log or logging.getLogger("plugins.mcp") - - self.log.info(f"Loading STIX bundle from: {stix_bundle_path}") - - # Initialize with STIX bundle if provided (single file) - if stix_bundle_path: - self.load_stix_bundle(stix_bundle_path) - - def load_stix_bundle(self, stix_bundle_path: str, embed_model: str | None = None): - """Load STIX bundle from file path and build embeddings.""" - try: - with open(stix_bundle_path, 'r') as f: - stix_bundle = json.load(f) - self.initialize_from_bundles([stix_bundle], embed_model=embed_model) - except FileNotFoundError: - raise FileNotFoundError(f"STIX bundle not found at: {stix_bundle_path}") - except json.JSONDecodeError: - raise ValueError(f"Invalid JSON in STIX bundle: {stix_bundle_path}") - - def initialize_from_bundles(self, stix_bundles: List[dict], embed_model: str | None = None): - """Initialize the RAG service with multiple STIX bundles and create retriever.""" - all_corpus = [] - all_adv_step = {} - for bundle in stix_bundles: - corpus, adv_step = self.extract_text_chunks(bundle) - all_corpus.extend(corpus) - all_adv_step.update(adv_step) - - self.corpus = all_corpus - self.adv_step = all_adv_step - - if not self.corpus: - # Embeddings() raises numpy AxisError on an empty corpus, and the - # caller swallows it and silently runs without CTI grounding. - # Bundles with no attack-pattern and no threat-actor are the - # normal way to get here. search_cti_title handles search=None. - self.search = None - self.log.warning( - "[RAG] No attack-pattern or threat-actor objects in the " - "selected STIX; CTI retrieval is unavailable for this run" - ) - return - - self.log.info("Initializing embeddings and retriever for STIX corpus") - embedder_kwargs = {"api_key": self.api_key} - if self.api_base: - embedder_kwargs["api_base"] = self.api_base - embedder_kwargs["custom_llm_provider"] = "custom_openai" - apply_litellm_ssl_verify(self.ssl_verify) - embedder = dspy.Embedder(embed_model, **embedder_kwargs) - self.search = dspy.retrievers.Embeddings( - corpus=self.corpus, - embedder=embedder, - k=self.topk_objects_to_retrieve, - ) - self.log.info(f"[RAG] Initialized search retriever: {self.search} with model {embed_model} and k={self.topk_objects_to_retrieve}") - - def extract_text_chunks(self, stix_bundle: dict) -> tuple[List[str], Dict[str, str]]: - """Extract text chunks from STIX bundle objects.""" - text_chunks = [] - adv_step = {} - - for obj in stix_bundle.get("objects", []): - if obj.get("type") in ["attack-pattern", "threat-actor"]: - name = obj.get("name", "") - description = obj.get("description", "") - - if name or description: - adv_step[name] = description - text_chunks.append(f"{name} | {description}") - - return text_chunks, adv_step - - def search_cti_title(self, query: str) -> List[str]: - """Returns top-5 results and then the names of the top-5 to top-30 results.""" - self.log.info(f"Searching CTI title with query: {query}") - - if not self.search: - self.log.warning("Search attempted but RAG service not initialized with STIX data") - return ["RAG service not initialized with STIX data"] - - self.log.debug(f"Using search retriever: {self.search}") - - topK = self.search(query) - self.log.debug(f" Retrieved top {len(topK)} results") - self.log.info(f"topK: {topK}") - names = [] - if len(topK) > 5: - names = [f"{x.split(' | ')[0]}" for x in topK.passages[5:30]] - topK = topK[:5] - self.log.info(f"names: {names}") - else: - names = [f"{x.split(' | ')[0]}" for x in topK.passages] - topK = topK - self.log.info(f"names: {names}") - return topK.passages - - def search_cti_data_by_title(self, name: str) -> str: - """Returns the full CTI data for a given name.""" - self.log.info(f"Searching CTI data for title: {name}") - - if name in self.adv_step: - self.log.debug("Found title in adv_step cache") - return self.adv_step[name] - - if not self.search: - self.log.warning("Search attempted but RAG service not initialized with STIX data") - return "RAG service not initialized with STIX data" - - results = [x for x in self.search(name, 10) if x.startswith(name + " | ")] - if not results: - self.log.warning(f"No CTI data found for name: {name}") - return f"No CTI data found for name: {name}" - - self.log.debug(f"Found {len(results)} matching results") - return results[0] - - def get_context_for_task(self, task: str) -> Dict[str, any]: - thoughts = [] - - thoughts.append(f"Getting context for task: {task}") - cti_results = self.search_cti_title(task) - thoughts.append(f"Retrieved {len(cti_results)} CTI results") - - detailed_context = [] - for result in cti_results[:3]: - if " | " in result: - name = result.split(" | ")[0] - detail = self.search_cti_data_by_title(name) - detailed_context.append({ - "name": name, - "description": detail - }) - thoughts.append(f"Retrieved detail for: {name}") - - return { - "search_results": cti_results, - "detailed_context": detailed_context, - "query": task, - "thoughts": thoughts # <-- used for Stage/Thoughts display - } - - - - - -# --------------------------------------------------------------------------- -# Capability registration -# --------------------------------------------------------------------------- -# Wraps the RAGService above as a Capability the workflow registry can discover -# and invoke. Workflows that accept "rag" in their accepted_capabilities list -# will receive a cti_context string in their context dict when this capability -# is enabled by the user. Until the orchestrator switches to consuming the -# capability registry, this declaration is dormant; the legacy code path in -# mcp_svc still wires RAG into runs directly. - -import asyncio -from pathlib import Path - -from plugins.mcp.app.capabilities.base import Capability - - -def _resolve_rag_file(base_dir: Path, name: str) -> Path: - """Resolve either uploaded RAG paths or generated STIX filenames.""" - raw = Path(str(name)) - if raw.is_absolute(): - return raw - - candidates = [base_dir / raw] - if len(raw.parts) == 1: - candidates.append(base_dir / "outputs_stix" / raw) - - for path in candidates: - if path.exists(): - return path - return candidates[0] - - -def _format_rag_context(rag_context: dict) -> str: - """Format a RAGService.get_context_for_task() result into a prompt string. - - Mirrors the formatter that lives inline in workflows/author.py and - workflows/plan_execute.py so the capability can stand on its own once the - orchestrator switches over. - """ - if not rag_context: - return "No CTI context available." - - parts = [] - if rag_context.get("search_results"): - parts.append("Relevant CTI findings:") - for i, result in enumerate(rag_context["search_results"][:3], 1): - parts.append(f"{i}. {result}") - if rag_context.get("detailed_context"): - parts.append("\nDetailed CTI Information:") - for ctx in rag_context["detailed_context"]: - parts.append(f"\n{ctx['name']}:") - parts.append(f"{ctx['description']}") - return "\n".join(parts) - - -def _build_rag_context_sync( - prompt: str, - bundles: list[dict], - *, - api_key: str, - api_base: str | None, - ssl_verify: bool | None, - embed_model: str, - topk: int, -) -> str: - rag = RAGService(api_key=api_key, api_base=api_base, ssl_verify=ssl_verify) - rag.topk_objects_to_retrieve = topk - rag.initialize_from_bundles(bundles, embed_model=embed_model) - raw_context = rag.get_context_for_task(prompt) - return _format_rag_context(raw_context) - - -async def _enrich(prompt: str, settings: dict) -> dict: - """Build a RAGService from selected files and return cti_context string. - - settings keys (all optional): - rag_files list of filenames under plugins/mcp/data/ - topk int, defaults to 5 - embed_model str, falls back to the chat model when unset - api_key str, embedding key fallback (mcp_svc fills this with - the resolved chat-LLM key when the caller does not - override it via embed_api_key) - embed_api_key str, optional override pointing embeddings at a - different provider than the chat LLM - embed_api_base str, optional override paired with embed_api_key - - Returns an empty dict when no files are selected so workflow signatures - that accept cti_context fall back to their default (empty string). - """ - rag_files = settings.get("rag_files") or [] - if not rag_files: - return {} - - embed_api_key = settings.get("embed_api_key") or settings.get("api_key") or "" - embed_api_base = settings.get("embed_api_base") or settings.get("api_base") or None - ssl_verify = settings.get("ssl_verify") - # An undeclared embedding model falls back to the chat model. A deployment - # pointing at its own gateway has no reason to reach for an OpenAI model - # name it never chose, and that name would be sent to that gateway. - embed_model = settings.get("embed_model") or settings.get("model") - if not embed_model: - raise ValueError( - "No embedding model: set embed_model, or a model on the global " - "LLM profile for it to inherit." - ) - topk = int(settings.get("topk") or 5) - - base_dir = Path(__file__).resolve().parent.parent.parent / "data" - bundles = [] - for name in rag_files: - path = _resolve_rag_file(base_dir, name) - if not path.exists(): - raise FileNotFoundError(f"RAG file not found: {path}") - with open(path, "r", encoding="utf-8") as f: - bundles.append(json.load(f)) - - cti_context = await asyncio.to_thread( - _build_rag_context_sync, - prompt, - bundles, - api_key=embed_api_key, - api_base=embed_api_base, - ssl_verify=ssl_verify, - embed_model=embed_model, - topk=topk, - ) - return {"cti_context": cti_context} - - -CAPABILITIES = [ - Capability( - id="rag", - display_name="CTI Ingestion (RAG)", - description=( - "Embed uploaded CTI bundles and inject the most relevant chunks " - "as context for the workflow. Useful for grounding both ability " - "authoring grounded in real threat reports." - ), - enrich=_enrich, - ui_settings_component="components/RAGSettings.vue", - contributes_fields=["cti_context"], - ), -] diff --git a/app/mcp_api.py b/app/mcp_api.py index f108cb3..7b74a23 100644 --- a/app/mcp_api.py +++ b/app/mcp_api.py @@ -28,22 +28,10 @@ from plugins.mcp.app.utilities.paths import get_mcp_data_dir, get_mcp_root from plugins.mcp.app.cti_ingest_svc import CTIIngestService -# Fallbacks for UI fields the yaml doesn't model. The numeric LM tunables -# come back from llm_defaults() which already applies its own fallbacks; -# the RAG-specific defaults stay here because they belong to the capability -# rather than the LLM provider. -# -# rag_embed_model has no literal default: an undeclared embedding model falls -# back to the configured chat model, since a deployment pointing at its own -# gateway has no reason to reach for an OpenAI model name it never chose. -_RAG_DEFAULTS = { - "rag_topk": 5, -} - # conf/local.yml is plaintext on disk beside tracked config, so credentials # never go in it. They come from .env or the per-session UI. # Matched by name rather than listed: the UI keeps growing key fields -# (embed_api_key, plan_api_key, cti_rag_api_key) and a list keeps missing them. +# and a list keeps missing them. # api_key_env names a variable, not a value, so it is kept. # Widened past api_key: a gateway credential also arrives as an Authorization # or x-api-key header, and the llm section used to accept any key at all, so @@ -197,14 +185,14 @@ async def execute(self, request): "workflow_id": "author" | "plan_execute" | ..., "lm_config": { model, api_key, api_base, temperature, max_tokens, max_tool_calls }, "enabled_servers": ["caldera_core", ...], - "enabled_capabilities": ["rag", ...], - "capability_settings": { "rag": { "rag_files": [...], "topk": 5, "embed_model": "..." } }, + "enabled_capabilities": ["cti", ...], + "capability_settings": { "cti": { "cti_files": [...] } }, "workflow_context": { ...workflow-specific UI/runtime options... } } Legacy payload (still accepted, mapped by mcp_svc): { "text": "...", "type": "factory"|"planner"|"rag_factory"|"rag_planner", - "config": { ...lm + rag_files/rag_topk/rag_embed_model }, "enabled_servers": [...] } + "config": { ...lm + cti_files }, "enabled_servers": [...] } """ self.log.info("[MCP] Executing request") try: @@ -287,13 +275,6 @@ async def defaults(self, request): "ssl_verify": cfg.get("ssl_verify", True), "fields_locked": cfg.get("fields_locked") or {}, } - for key, fallback in _RAG_DEFAULTS.items(): - payload[key] = cfg.get(key, fallback) - payload["rag_embed_model"] = ( - cfg.get("rag_embed_model") - or cfg.get("embed_model") - or cfg.get("model") - ) return web.json_response(payload) except Exception as e: self.log.error(f"[MCP] Error fetching defaults: {e}") @@ -407,71 +388,6 @@ async def status(self, request): self.log.debug(f"[MCP] Status for run {run_id} served from cache") return web.json_response({"run_id": run_id, **snapshot}) - async def upload_rag(self, request): - try: - reader = await request.multipart() - part = await reader.next() - if not part or part.name != "file": - return web.json_response({"error": 'Missing "file" field in form-data'}, status=400) - - raw_filename = part.filename or "rag.json" - filename = os.path.basename(raw_filename) - - if not filename.lower().endswith(".json"): - return web.json_response({"error": "Only .json files are allowed"}, status=400) - - base_dir = (Path(__file__).resolve().parent.parent / "data") - base_dir.mkdir(parents=True, exist_ok=True) - - target_path = base_dir / filename - if target_path.exists(): - ts = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S") - stem = Path(filename).stem - suffix = Path(filename).suffix - filename = f"{stem}_{ts}{suffix}" - target_path = base_dir / filename - - file_bytes = await part.read() - try: - _ = json.loads(file_bytes.decode("utf-8")) - except Exception as e: - return web.json_response({"error": f"Invalid JSON: {e}"}, status=400) - - with open(target_path, "wb") as f: - f.write(file_bytes) - - stat = target_path.stat() - self.log.info(f"[MCP] Uploaded RAG file to {target_path} ({stat.st_size} bytes)") - return web.json_response({ - "message": "RAG file uploaded", - "filename": filename, - "size": stat.st_size, - "path": str(target_path) - }) - except Exception as e: - self.log.error(f"[MCP] Error uploading RAG file: {e}") - return web.json_response({"error": str(e)}, status=500) - - async def list_rag(self, request): - try: - base_dir = (Path(__file__).resolve().parent.parent / "data") - files = [] - if base_dir.exists(): - for p in base_dir.glob("*.json"): - try: - stat = p.stat() - files.append({ - "filename": p.name, - "size": stat.st_size, - "modified": datetime.fromtimestamp(stat.st_mtime, timezone.utc).isoformat().replace("+00:00", "Z") - }) - except Exception: - continue - return web.json_response({"files": sorted(files, key=lambda x: x["filename"].lower())}) - except Exception as e: - self.log.error(f"[MCP] Error listing RAG files: {e}") - return web.json_response({"error": str(e)}, status=500) - async def list_runs(self, request): """List all MLflow runs with basic information.""" try: @@ -750,6 +666,7 @@ async def list_stix_cti(self, request): model = None provider = None extractor = None + objects = None try: with p.open("r", encoding="utf-8") as f: @@ -757,6 +674,12 @@ async def list_stix_cti(self, request): model = bundle.get("x_cti_model") provider = bundle.get("x_cti_provider") extractor = (bundle.get("x_cti_config") or {}).get("extractor") + # What the cti capability will actually put in the + # prompt, so the picker can say so before a run. + objects = sum( + 1 for o in (bundle.get("objects") or []) + if o.get("type") in ("attack-pattern", "threat-actor") + ) except Exception: # Do NOT fail listing if one file is malformed pass @@ -769,6 +692,7 @@ async def list_stix_cti(self, request): "model": model, "provider": provider, "extractor": extractor, + "objects": objects, }) self.log.info(f"[MCP] listing stix cti files: {files}") diff --git a/app/mcp_server.py b/app/mcp_server.py index c26acf1..a98f3cb 100644 --- a/app/mcp_server.py +++ b/app/mcp_server.py @@ -431,7 +431,12 @@ def create_windows_ability( command_description: Detailed description of the command to use in the windows ability tactic: MITRE ATT&CK tactic (e.g., 'privilege-escalation', 'discovery') technique_name: Name of the MITRE ATT&CK technique - technique_id: Optional MITRE ATT&CK technique ID (e.g., 'T1548.002') + technique_id: MITRE ATT&CK technique id (e.g., 'T1548.002'). Optional, + but CTI coverage is matched on it: an ability with no id is + invisible to cti_pipeline_build_adversary, so a technique you + just wrote an ability for still reports as unmatched. Supply it + whenever the technique is known, which it is when the id came + from cti_context. payloads: Optional list of payload files needed Returns: @@ -439,10 +444,11 @@ def create_windows_ability( """ ability_id = str(uuid.uuid4()) created_command = create_command(command_description, "windows") - # Create the executor object with default values for optional fields + # Caldera keys an executor as (name, platform): name is the runner, platform + # is the OS. Transposed, no agent ever matches and the ability cannot run. executor = { - "name": "windows", - "platform": "psh", + "name": "psh", + "platform": "windows", "command": created_command, "code": None, "language": None, @@ -497,7 +503,12 @@ def create_linux_ability( command_description: Detailed description of the command to use in the linux ability tactic: MITRE ATT&CK tactic (e.g., 'privilege-escalation', 'discovery') technique_name: Name of the MITRE ATT&CK technique - technique_id: Optional MITRE ATT&CK technique ID (e.g., 'T1548.002') + technique_id: MITRE ATT&CK technique id (e.g., 'T1548.002'). Optional, + but CTI coverage is matched on it: an ability with no id is + invisible to cti_pipeline_build_adversary, so a technique you + just wrote an ability for still reports as unmatched. Supply it + whenever the technique is known, which it is when the id came + from cti_context. payloads: Optional list of payload files needed Returns: @@ -506,10 +517,11 @@ def create_linux_ability( ability_id = str(uuid.uuid4()) created_command = create_command(command_description, "linux") - # Create the executor object with default values for optional fields + # Caldera keys an executor as (name, platform): name is the runner, platform + # is the OS. Transposed, no agent ever matches and the ability cannot run. executor = { - "name": "linux", - "platform": "sh", + "name": "sh", + "platform": "linux", "command": created_command, "code": None, "language": None, diff --git a/app/mcp_svc.py b/app/mcp_svc.py index 64f25a7..a80bcb5 100644 --- a/app/mcp_svc.py +++ b/app/mcp_svc.py @@ -21,8 +21,8 @@ ) -# Where chat session history lives on disk. Co-located with RAG uploads -# under plugins/mcp/data/. Single JSON file rather than one-per-session +# Where chat session history lives on disk, under plugins/mcp/data/. +# Single JSON file rather than one-per-session # because the load is once at boot and the writes happen at the end of # successful runs, where one extra ~tens-of-KB write is invisible against # the LLM round-trip we just finished. @@ -56,8 +56,8 @@ def _scrub_secrets(text: str) -> str: _LEGACY_TYPE_MAP = { "factory": ("author", []), "planner": ("plan_execute", []), - "rag_factory": ("author", ["rag"]), - "rag_planner": ("plan_execute", ["rag"]), + "rag_factory": ("author", ["cti"]), + "rag_planner": ("plan_execute", ["cti"]), } # Bound on the in-memory live-run cache. Older snapshots fall off LRU-style @@ -91,7 +91,6 @@ def __init__(self, services, server_registry=None, workflow_registry=None, capab self.auth_svc = services.get("auth_svc") self.log = logging.getLogger("plugins.mcp") - self.rag_service = None self.server_registry = server_registry or {} self.workflow_registry = workflow_registry or {} self.capability_registry = capability_registry or {} @@ -409,7 +408,7 @@ async def execute(self, focus: str = None, prompt: str = "", model_config: dict _LEGACY_TYPE_MAP and can be removed once the UI ships exclusively on the new shape. """ - legacy_rag_settings = None + legacy_cti_settings = None if workflow_id is None and focus is not None: mapped = _LEGACY_TYPE_MAP.get(focus) if mapped is None: @@ -420,14 +419,10 @@ async def execute(self, focus: str = None, prompt: str = "", model_config: dict if c not in enabled_capabilities: enabled_capabilities.append(c) mc = model_config or {} - if mc.get("rag_files") and "rag" not in enabled_capabilities: - enabled_capabilities.append("rag") - if "rag" in enabled_capabilities: - legacy_rag_settings = { - "rag_files": mc.get("rag_files") or [], - "topk": mc.get("rag_topk"), - "embed_model": mc.get("rag_embed_model"), - } + if mc.get("cti_files") and "cti" not in enabled_capabilities: + enabled_capabilities.append("cti") + if "cti" in enabled_capabilities: + legacy_cti_settings = {"cti_files": mc.get("cti_files") or []} if lm_config is None: lm_config = mc @@ -456,25 +451,8 @@ async def execute(self, focus: str = None, prompt: str = "", model_config: dict lm_obj = self._create_dspy_client(resolved_lm) cap_settings = dict(capability_settings or {}) - if legacy_rag_settings is not None and "rag" not in cap_settings: - cap_settings["rag"] = legacy_rag_settings - if "rag" in scoped_capabilities: - rag_cfg = dict(cap_settings.get("rag") or {}) - if not rag_cfg.get("api_key"): - rag_cfg["api_key"] = resolved_lm.get("api_key", "") - if not rag_cfg.get("api_base"): - rag_cfg["api_base"] = resolved_lm.get("api_base", "") - if not rag_cfg.get("embed_model"): - # Last rung is the chat model: one endpoint, one model, unless - # a deployment explicitly names a separate embedding model. - rag_cfg["embed_model"] = ( - resolved_lm.get("embed_model") - or resolved_lm.get("rag_embed_model") - or resolved_lm.get("model") - ) - if rag_cfg.get("ssl_verify") is None: - rag_cfg["ssl_verify"] = resolved_lm.get("ssl_verify") - cap_settings["rag"] = rag_cfg + if legacy_cti_settings is not None and "cti" not in cap_settings: + cap_settings["cti"] = legacy_cti_settings # Mint the run record without activating it. Every write against it # goes through a RunTracker bound to this id, so concurrent requests @@ -621,8 +599,14 @@ async def _run_execution(self, workflow, prompt, run_id, lm_obj, ",".join(sorted(contrib.keys())), ) except Exception as e: + # The operator enabled this capability deliberately. A run + # that continues without it finishes green and is + # indistinguishable from a grounded one, so it has to stop. tracker.set_tag(f"capability_{cap_id}_stage", "error") - self.log.warning(f"[MCP] Capability '{cap_id}' enrich failed: {e}") + self.log.error(f"[MCP] Capability '{cap_id}' enrich failed: {e}") + raise RuntimeError( + f"{cap_id} was enabled for this run but failed: {e}" + ) from e if workflow.run is None: raise RuntimeError(f"Workflow {workflow.id} has no run() function") @@ -639,6 +623,11 @@ async def _run_execution(self, workflow, prompt, run_id, lm_obj, run_kwargs["chat_history"] = chat_history if isinstance(workflow_context, dict) and workflow_context: run_kwargs["workflow_context"] = workflow_context + # getattr, not attribute access: a workflow contributed by another + # plugin may predate this field, and so do the test fakes. + denied_tools = getattr(workflow, "denied_tools", None) + if denied_tools: + run_kwargs["denied_tools"] = list(denied_tools) result = await workflow.run( prompt, lm_obj, diff --git a/app/utilities/cti_stix_builders.py b/app/utilities/cti_stix_builders.py index 5b200a9..6e9b557 100644 --- a/app/utilities/cti_stix_builders.py +++ b/app/utilities/cti_stix_builders.py @@ -133,7 +133,7 @@ def make_attack_pattern(ttp_text, taxonomy: dict): - "Exfiltration Over Alternative Protocol (T1048)" - "T1070.004" - plain English name - - dict from IR with keys: id, name, confidence, evidence + - dict from IR with keys: id, name, x_cti_confidence, x_cti_evidence """ confidence = None @@ -143,8 +143,10 @@ def make_attack_pattern(ttp_text, taxonomy: dict): # Normalize IR dict → string (WITHOUT losing metadata) # ---------------------------- if isinstance(ttp_text, dict): - confidence = ttp_text.get("confidence") - evidence = ttp_text.get("evidence") + # cti_technique_grounding writes the x_cti_ prefix. The bare names are + # still accepted so a hand-written IR keeps its provenance. + confidence = ttp_text.get("x_cti_confidence", ttp_text.get("confidence")) + evidence = ttp_text.get("x_cti_evidence", ttp_text.get("evidence")) if "id" in ttp_text: t = ttp_text["id"].strip() diff --git a/app/utilities/llm_client.py b/app/utilities/llm_client.py index e9bb830..f4a33e7 100644 --- a/app/utilities/llm_client.py +++ b/app/utilities/llm_client.py @@ -192,8 +192,7 @@ def _aiohttp_ssl_arg(ssl_verify): LLM_OVERRIDABLE = frozenset({ "provider", "model", "api_base", "api_base_env", "api_key_env", "ssl_verify", "offline", "temperature", "top_p", "max_tokens", - "max_tool_calls", "timeout", "stream", "embed_model", - "rag_topk", "rag_embed_model", + "max_tool_calls", "timeout", "stream", }) WORKLOAD_OVERRIDABLE = frozenset({ @@ -203,9 +202,10 @@ def _aiohttp_ssl_arg(ssl_verify): "stream", # Per-workload run mode: extraction can go offline without silencing chat. "offline", - # A different embedding model is a separate axis from the chat endpoint, - # so it cannot be confused with one. - "embed_model", + # How much attached intel the cti capability may put in a prompt. A + # deployment on a short-window model has to be able to lower it, and + # set_config is the only writer, so it has to be listed here. + "context_budget_tokens", }) @@ -271,8 +271,6 @@ def get_llm_provenance(profile: str = "llm", *, runtime: bool = False) -> dict: "max_tokens": llm.get("max_tokens"), "timeout": llm.get("timeout", 60), "ssl_verify": llm.get("ssl_verify", True), - # Optional: allow config to specify embedding model explicitly - "embed_model": llm.get("embed_model") or llm.get("model"), } if not runtime: diff --git a/app/workflows/author.py b/app/workflows/author.py index e6d7894..35761a5 100644 --- a/app/workflows/author.py +++ b/app/workflows/author.py @@ -25,6 +25,7 @@ dspy_lm_kwargs_from_settings, ) from plugins.mcp.app.dspy_runner import safe_react_acall +from plugins.mcp.app.workflows.prompts.common import CTI_CONTEXT_DESC def get_env(lm_settings=None): @@ -120,7 +121,7 @@ class DSPyCalderaFactoryClient(dspy.Signature): adversary_emulation_task: str = dspy.InputField() process_result: str = dspy.OutputField(desc=_AUTHOR_OUTPUT_DESC) -class DSPyCalderaFactoryClientWithRAG(dspy.Signature): +class DSPyCalderaFactoryClientWithCTI(dspy.Signature): """You are an ability factory for the Caldera adversary emulation platform, enhanced with Cyber Threat Intelligence (CTI) data. You have access to MCP tool servers that wrap Caldera's core API and any installed plugins. Your @@ -137,9 +138,7 @@ class DSPyCalderaFactoryClientWithRAG(dspy.Signature): """ adversary_emulation_task: str = dspy.InputField() - cti_context: str = dspy.InputField( - desc="Relevant CTI (Cyber Threat Intelligence) information including attack patterns, techniques, and threat actor behaviors" - ) + cti_context: str = dspy.InputField(desc=CTI_CONTEXT_DESC) process_result: str = dspy.OutputField( desc=( _AUTHOR_OUTPUT_DESC @@ -150,38 +149,17 @@ class DSPyCalderaFactoryClientWithRAG(dspy.Signature): ) -def format_rag_context(rag_context): - """Format RAG context into a string for the DSPy signature.""" - if not rag_context: - return "No CTI context available." - - formatted_parts = [] - - # Add search results summary - if "search_results" in rag_context: - formatted_parts.append("Relevant CTI findings:") - for i, result in enumerate(rag_context["search_results"][:3], 1): - formatted_parts.append(f"{i}. {result}") - - # Add detailed context - if "detailed_context" in rag_context: - formatted_parts.append("\nDetailed CTI Information:") - for ctx in rag_context["detailed_context"]: - formatted_parts.append(f"\n{ctx['name']}:") - formatted_parts.append(f"{ctx['description']}") - - return "\n".join(formatted_parts) - -async def run(adversary_emulation_task: str, lm_obj=None, rag_context=None, run_id=None, enabled_servers=None, server_registry=None, cti_context: str = "", **_extra_capability_context): - # cti_context (formatted string) is the new orchestrator's path. rag_context - # (raw dict) is the legacy mcp_svc shim. Either may be supplied; cti_context - # wins. **_extra_capability_context absorbs unknown kwargs from future - # capabilities so adding one doesn't break this signature. +async def run(adversary_emulation_task: str, lm_obj=None, run_id=None, enabled_servers=None, server_registry=None, cti_context: str = "", workflow_context: dict | None = None, denied_tools=None, + **_extra_capability_context): + # cti_context is the attached intel, already rendered by the cti + # capability. **_extra_capability_context absorbs unknown kwargs from + # future capabilities so adding one doesn't break this signature. # # lm_obj is the dict mcp_svc produces from resolve_llm_config (yaml + .env # + UI overrides, with fields_locked enforced). Workflows trust it; they # do not re-merge yaml here. Tests that call run() directly without # lm_obj fall back to the same yaml-resolved defaults. + denied = set(denied_tools or ()) _ensure_mlflow() lm_settings = dict(lm_obj) if lm_obj else llm_defaults() max_tool_calls = lm_settings.get("max_tool_calls") or 5 @@ -259,6 +237,8 @@ async def run(adversary_emulation_task: str, lm_obj=None, rag_context=None, run_ for server_name, session in zip(enabled_servers, sessions): tool_list = (await session.list_tools()).tools for tool in tool_list: + if tool.name in denied: + continue if tool.name in seen: raise ValueError( f"Tool name collision: '{tool.name}' defined by both " @@ -275,22 +255,16 @@ async def run(adversary_emulation_task: str, lm_obj=None, rag_context=None, run_ tracker.set_tag("stage", "creating DSPy ReAct instance") # Resolve CTI context: prefer the orchestrator-supplied string, - # fall back to formatting the legacy structured dict. resolved_cti = cti_context - if not resolved_cti and rag_context: - resolved_cti = format_rag_context(rag_context) if resolved_cti: - signature = DSPyCalderaFactoryClientWithRAG + signature = DSPyCalderaFactoryClientWithCTI tracker.log_param("cti_context_preview", resolved_cti[:1000]) tracker.set_tag("cti_context_length", len(resolved_cti)) - if rag_context: - tracker.set_tag("cti_search_results_count", len(rag_context.get("search_results", []))) - tracker.set_tag("cti_detailed_context_count", len(rag_context.get("detailed_context", []))) print(f"[MCP] Passing CTI context to LLM ({len(resolved_cti)} chars)") react = dspy.ReAct(signature, tools=dspy_tools, max_iters=max_tool_calls) - tracker.set_tag("stage", "executing DSPy ReAct with RAG") + tracker.set_tag("stage", "executing DSPy ReAct with CTI") result = await safe_react_acall( react, adversary_emulation_task=adversary_emulation_task, @@ -368,8 +342,18 @@ async def run(adversary_emulation_task: str, lm_obj=None, rag_context=None, run_ ), signature=DSPyCalderaFactoryClient, required_servers=["caldera_core"], - optional_servers=[], - accepted_capabilities=["rag"], + optional_servers=["cti_pipeline"], + accepted_capabilities=["cti"], + # The signature tells the model not to run operations. caldera_core is + # required and cti_pipeline is now offered, and between them they + # expose all three of these, so without the scope that instruction is + # the only thing standing between an authoring run and a live + # operation. Authoring an adversary stays allowed; running one does not. + denied_tools=[ + "core_create_operation", + "core_add_link_to_operation", + "cti_pipeline_run_operation", + ], mlflow_experiment=_MLFLOW_EXPERIMENT, ui_component="author.vue", example_prompts=[ diff --git a/app/workflows/base.py b/app/workflows/base.py index 019bb68..d54843b 100644 --- a/app/workflows/base.py +++ b/app/workflows/base.py @@ -51,6 +51,13 @@ class Workflow: # capability's enrich() when its id appears here AND the user enabled it. accepted_capabilities: list[str] = field(default_factory=list) + # Tool names this workflow must never be given, even when the server that + # exposes them is enabled. Server scoping is all-or-nothing, so a workflow + # that needs most of a server but must not reach two of its tools has no + # other way to say so. The agent cannot call a tool that is absent from + # its list, which is the boundary a persona docstring can only request. + denied_tools: list[str] = field(default_factory=list) + # MLflow experiment this workflow's runs belong to. mcp_svc mints the run # here, and the boot-time orphan sweep only reconciles runs in the # experiments the registered workflows declare. diff --git a/app/workflows/plan_execute.py b/app/workflows/plan_execute.py index f9c119d..3274b77 100644 --- a/app/workflows/plan_execute.py +++ b/app/workflows/plan_execute.py @@ -31,7 +31,6 @@ CHAT_HISTORY_DESC, CTI_CONTEXT_DESC, PLAN_EXECUTE_OUTPUT_DESC, - format_rag_context, ) from plugins.mcp.app.workflows.prompts.plan_execute import ( PLAN_EXECUTE_AGENT_DOC, @@ -134,7 +133,7 @@ class DSPyCalderaPlannerClient(dspy.Signature): chat_history: str = dspy.InputField(default="", desc=CHAT_HISTORY_DESC) process_result: str = dspy.OutputField(desc=PLAN_EXECUTE_OUTPUT_DESC) -class DSPyCalderaPlannerClientWithRAG(dspy.Signature): +class DSPyCalderaPlannerClientWithCTI(dspy.Signature): """Plan and execute a CTI-grounded CALDERA adversary-emulation request.""" adversary_emulation_task: str = dspy.InputField() operation_context: str = dspy.InputField( @@ -153,13 +152,13 @@ class DSPyCalderaPlannerClientWithRAG(dspy.Signature): DSPyCalderaPlannerClient.__doc__ = PLAN_EXECUTE_AGENT_DOC -DSPyCalderaPlannerClientWithRAG.__doc__ = PLAN_EXECUTE_AGENT_WITH_CTI_DOC +DSPyCalderaPlannerClientWithCTI.__doc__ = PLAN_EXECUTE_AGENT_WITH_CTI_DOC -async def run(adversary_emulation_task: str, lm_obj=None, rag_context=None, +async def run(adversary_emulation_task: str, lm_obj=None, run_id=None, enabled_servers=None, server_registry=None, cti_context: str = "", chat_history: str = "", - workflow_context: dict | None = None, + workflow_context: dict | None = None, denied_tools=None, **_extra_capability_context): """ lm_obj can be: @@ -170,6 +169,7 @@ async def run(adversary_emulation_task: str, lm_obj=None, rag_context=None, - None, to fall back to llm_defaults() from the shared config module (used by tests / direct invocation) """ + denied = set(denied_tools or ()) _ensure_mlflow() if isinstance(lm_obj, dspy.LM): lm_instance = lm_obj @@ -265,6 +265,8 @@ async def run(adversary_emulation_task: str, lm_obj=None, rag_context=None, for server_name, session in zip(enabled_servers, sessions): tool_list = (await session.list_tools()).tools for tool in tool_list: + if tool.name in denied: + continue if tool.name in seen: raise ValueError( f"Tool name collision: '{tool.name}' defined by both " @@ -281,24 +283,19 @@ async def run(adversary_emulation_task: str, lm_obj=None, rag_context=None, # Resolve CTI context: prefer the orchestrator-supplied string, # fall back to formatting the legacy structured dict. resolved_cti = cti_context - if not resolved_cti and rag_context: - resolved_cti = format_rag_context(rag_context) operation_context = format_plan_execute_context(workflow_context) if operation_context: tracker.log_param("operation_context_preview", operation_context[:1000]) tracker.set_tag("operation_context_length", len(operation_context)) if resolved_cti: - signature = DSPyCalderaPlannerClientWithRAG + signature = DSPyCalderaPlannerClientWithCTI tracker.log_param("cti_context_preview", resolved_cti[:1000]) tracker.set_tag("cti_context_length", len(resolved_cti)) - if rag_context: - tracker.set_tag("cti_search_results_count", len(rag_context.get("search_results", []))) - tracker.set_tag("cti_detailed_context_count", len(rag_context.get("detailed_context", []))) print(f"[MCP] Passing CTI context to LLM ({len(resolved_cti)} chars)") react = dspy.ReAct(signature, tools=dspy_tools, max_iters=max_tool_calls) - tracker.set_tag("stage", "executing DSPy ReAct with RAG") + tracker.set_tag("stage", "executing DSPy ReAct with CTI") result = await safe_react_acall( react, adversary_emulation_task=adversary_emulation_task, @@ -396,7 +393,7 @@ async def run(adversary_emulation_task: str, lm_obj=None, rag_context=None, signature=DSPyCalderaPlannerClient, required_servers=["caldera_core"], optional_servers=["cti_pipeline"], - accepted_capabilities=["rag"], + accepted_capabilities=["cti"], mlflow_experiment=_MLFLOW_EXPERIMENT, ui_component="plan_execute.vue", example_prompts=PLAN_EXECUTE_EXAMPLES, diff --git a/app/workflows/prompts/common.py b/app/workflows/prompts/common.py index c419cbc..5e62af0 100644 --- a/app/workflows/prompts/common.py +++ b/app/workflows/prompts/common.py @@ -2,7 +2,6 @@ from __future__ import annotations -from typing import Any CHAT_HISTORY_DESC = ( @@ -16,8 +15,9 @@ CTI_CONTEXT_DESC = ( - "ATT&CK technique descriptions retrieved from the selected STIX " - "bundles, and the named threat actor when the report identifies one. " + "The ATT&CK techniques carried by the attached intel, with the report's " + "own wording where the pipeline captured it, and the named threat actor " + "when the report identifies one. " "Technique-level only: it never contains hosts, accounts, domains, " "software or infrastructure. Use this as grounding, not as permission " "to invent missing facts." @@ -37,21 +37,3 @@ ) -def format_rag_context(rag_context: dict[str, Any] | None) -> str: - """Format RAG capability output for workflow signatures.""" - if not rag_context: - return "No CTI context available." - - formatted_parts: list[str] = [] - if "search_results" in rag_context: - formatted_parts.append("Relevant CTI findings:") - for i, result in enumerate(rag_context["search_results"][:3], 1): - formatted_parts.append(f"{i}. {result}") - - if "detailed_context" in rag_context: - formatted_parts.append("\nDetailed CTI Information:") - for ctx in rag_context["detailed_context"]: - formatted_parts.append(f"\n{ctx['name']}:") - formatted_parts.append(f"{ctx['description']}") - - return "\n".join(formatted_parts) diff --git a/app/workflows/prompts/plan_execute.py b/app/workflows/prompts/plan_execute.py index c94a6ff..404e3b5 100644 --- a/app/workflows/prompts/plan_execute.py +++ b/app/workflows/prompts/plan_execute.py @@ -143,11 +143,11 @@ def format_plan_execute_context(context: dict[str, Any] | None) -> str: return "" parts: list[str] = [] - selected = _as_text_list(context.get("selected_stix_files")) - if selected: - parts.append("Selected STIX bundles:\n" + "\n".join(f"- {x}" for x in selected)) - - if context.get("cti_uses_rag"): - parts.append("CTI selection mode: selected STIX bundles are also enabled as RAG context.") + # Names only. The techniques these bundles carry arrive separately as + # cti_context; this is so the agent can name its source and pass a real + # filename to a CTI tool instead of guessing one. + attached = _as_text_list(context.get("attached_intel")) + if attached: + parts.append("Attached CTI bundles:\n" + "\n".join(f"- {x}" for x in attached)) return "\n\n".join(parts) diff --git a/conf/default.yml b/conf/default.yml index 0e0f520..e219cea 100644 --- a/conf/default.yml +++ b/conf/default.yml @@ -41,6 +41,15 @@ cti: top_p: 1.0 timeout: 120 offline: true + # Ceiling on the intel handed to a workflow as cti_context. A one-report + # bundle is a few thousand tokens and a large one around 13k, so this holds + # roughly two without trimming. + # + # It is a ceiling, not a target. Answer quality against context length is an + # inverted U and published long-context evaluations put the turn well above + # this figure, so raise it if operators routinely attach several bundles and + # lower it for a model with a short usable window. + context_budget_tokens: 24000 # Caldera's REST API, used by the caldera_core subprocess. Both values # auto-resolve from caldera's own config, so a stock deployment needs no setup; diff --git a/gui/views/chat/ChatRagPanel.vue b/gui/views/chat/ChatRagPanel.vue deleted file mode 100644 index fd08f24..0000000 --- a/gui/views/chat/ChatRagPanel.vue +++ /dev/null @@ -1,389 +0,0 @@ - - - - - - diff --git a/gui/views/chat/ChatSidebar.vue b/gui/views/chat/ChatSidebar.vue index 670f46f..31b285e 100644 --- a/gui/views/chat/ChatSidebar.vue +++ b/gui/views/chat/ChatSidebar.vue @@ -1,11 +1,10 @@