diff --git a/.github/workflows/publish-canary.yml b/.github/workflows/publish-canary.yml index 97d844b6a..8c527fb76 100644 --- a/.github/workflows/publish-canary.yml +++ b/.github/workflows/publish-canary.yml @@ -13,10 +13,6 @@ on: branches: - canary -concurrency: - group: publish-canary-${{ github.ref }} - cancel-in-progress: true - permissions: contents: read diff --git a/pyproject.toml b/pyproject.toml index 1fc1f04e1..b1c3f0c21 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -92,7 +92,11 @@ gcs = [ ] ai = [ 'litellm < 2', - 'safecmd' + 'safecmd', + # safecmd shells out to the `shfmt` binary (via shutil.which); pin shfmt-py + # explicitly so `secator install addons ai` always ships the binary, not just + # the safecmd Python package. Without it the guardrail shell parser degrades. + 'shfmt-py' ] [project.scripts] diff --git a/secator/ai/actions.py b/secator/ai/actions.py index 63c010d2c..f97889782 100644 --- a/secator/ai/actions.py +++ b/secator/ai/actions.py @@ -1,7 +1,6 @@ """Action handlers for AI task.""" import json -import os -import subprocess +import threading import uuid from concurrent.futures import ThreadPoolExecutor, as_completed from dataclasses import dataclass, field, fields @@ -12,6 +11,28 @@ from secator.output_types import Ai, Error, Info, Warning, OutputType, FINDING_TYPES from secator.template import TemplateLoader from secator.utils import format_token_count +from secator.ai.utils import ( + _sanitized_env, _build_action_display, _is_approved, _truncate, _format_action_error, + _is_heavy_runner, _sanitize_child_opts, build_subagent_prompt, _union_live_results, + _coerce_finding_fields, _get_action_label, _decrypt_dict, +) + + +# Bound recursive AI-subagent fan-out so injected output can't drive an +# exponential subagent/token blow-up. Depth caps recursion (child inherits +1 via +# context); breadth caps how many subagents one parent turn may spawn. +_MAX_SUBAGENT_DEPTH = 3 +_MAX_SUBAGENTS_PER_TURN = 5 +_SUBAGENT_TURN_LOCK = threading.Lock() + +# Cap shell stdout before it enters AI history so a huge command can't blow up +# the next prompt's token budget; head+tail keeps both the start and the result. +_MAX_SHELL_OUTPUT_CHARS = 4000 + +# Cap on ad-hoc AI shell commands (dispatched as the `command` task). Applied as an +# instance attribute post-construction (see _handle_shell) since max_timeout is not a +# run_opts-settable field. +_SHELL_TIMEOUT = 60 @dataclass @@ -27,6 +48,8 @@ class ActionContext: """ targets: List[str] model: str + api_key: str = "" + api_base: str = "" encryptor: Any = None dry_run: bool = False verbose: bool = False @@ -34,6 +57,7 @@ class ActionContext: scope: str = "workspace" results: Optional[List[Dict]] = None max_workers: int = 3 + in_batch: bool = False # set on the per-batch ctx so the per-turn fan-out cap applies subagent: bool = False silent: bool = False sync: bool = True @@ -55,40 +79,82 @@ def get_query_engine(self): return self._query_engine -SENSITIVE_ENV_PREFIXES = ( - "SECATOR_", - "ANTHROPIC_", "OPENAI_", "GOOGLE_", "AZURE_", "AWS_", "GCP_", - "GITHUB_TOKEN", "GITLAB_TOKEN", "SLACK_TOKEN", "DISCORD_TOKEN", - "SECRET_", "TOKEN_", "API_KEY", "PRIVATE_KEY", -) +def _build_hooks_from_context(context: Dict) -> Dict: + """Build the runner hooks dict from ``context['drivers']``. + + Sub-runners dispatched by the ai task are constructed in-process and run + synchronously, so the framework's pickle path (``__setstate__``, which + re-registers driver hooks from ``context['drivers']``) never runs for them. + Without this, a sub-runner inherits the ai task's ``workspace_id`` / + ``drivers`` in its context but registers *no* driver hooks — so its + ``mongodb``/``api`` ``update_runner``/``update_finding`` hooks never fire and + its runner doc + findings are never persisted to the workspace. The result: + sub-runs are absent from the workspace History. + This mirrors the normal CLI entrypoint (``cli_helper._run``): import each + driver's ``secator.hooks..HOOKS`` and ``deep_merge_dicts`` them into a + single class-keyed dict (keyed by ``Scan``/``Workflow``/``Task``). The dict is + returned raw (not flattened) because ``Task``/``Workflow`` forward + ``self._hooks.get(Task, {})`` down to their command/task signatures. + + Args: + context: Runner context dict (expects ``drivers`` list). + + Returns: + dict: Merged hooks dict suitable for ``runner_cls(..., hooks=hooks)``. + """ + from secator.loader import discover_external_drivers, get_available_drivers, order_drivers + from secator.utils import import_dynamic, deep_merge_dicts + + drivers = list(context.get('drivers', [])) + if not drivers: + return {} + discover_external_drivers() + # Order by canonical priority so authoritative backends (e.g. mongodb) register + # their hooks before relay drivers (e.g. api) — same ordering as __setstate__. + drivers = order_drivers(drivers) + supported = set(get_available_drivers()) + hooks_list = [] + for driver in drivers: + if driver not in supported: + continue + driver_hooks = import_dynamic(f'secator.hooks.{driver}', 'HOOKS') + if driver_hooks: + hooks_list.append(driver_hooks) + if not hooks_list: + return {} + return deep_merge_dicts(*hooks_list) -def _sanitized_env() -> dict: - """Return a copy of os.environ with sensitive variables removed.""" - return {k: v for k, v in os.environ.items() - if not any(k.startswith(p) for p in SENSITIVE_ENV_PREFIXES) - and "KEY" not in k and "SECRET" not in k and "TOKEN" not in k and "PASSWORD" not in k} +def _build_child_hooks_or_denial(context: Dict) -> Tuple[Dict, Optional["Warning"]]: + """Rebuild the child's persistence hooks, refusing a persistence-less child. -def _build_action_display(action: Dict) -> str: - """Build a display string for the action being checked. + ``context`` carries the parent's ``drivers`` (copied via ``_get_result_context``), + so an empty/failed rebuild while the parent HAS drivers means the child would run + to completion and silently persist nothing (lost findings/docs). In that case + return a denial ``Warning`` (same shape other denials use) so the caller yields it + and skips the spawn. When the parent itself has no drivers (pure local/no-persistence + run) an empty-hooks child is expected and allowed. - Returns a concise description of the command/task/workflow for prompt context. + Returns ``(hooks, denial)``; if ``denial`` is non-None the caller must not spawn. """ - action_type = action.get("action", "") - if action_type == "shell": - return action.get("command", "") - elif action_type in ("task", "workflow"): - name = action.get("name", "") - targets = action.get("targets", []) - opts = action.get("opts", {}) - parts = [f"{action_type}: {name}"] - if targets: - parts.append(f"targets={targets}") - if opts: - parts.append(f"opts={opts}") - return " ".join(parts) - return "" + parent_has_drivers = bool(context.get('drivers')) + try: + hooks = _build_hooks_from_context(context) + except Exception as e: # narrow to the rebuild — surface, don't degrade to hooks={} + if parent_has_drivers: + return {}, Warning( + message=f"Subagent spawn denied: persistence hook rebuild failed — {type(e).__name__}: {e}", + _context=context, + ) + return {}, None + if parent_has_drivers and not hooks: + return {}, Warning( + message="Subagent spawn denied: parent has persistence drivers but child hook rebuild " + "was empty (would silently drop findings/docs)", # noqa: E131 + _context=context, + ) + return hooks, None def check_guardrails_sync(action: Dict, ctx: ActionContext) -> Tuple[Optional[str], List]: @@ -106,6 +172,39 @@ def check_guardrails_sync(action: Dict, ctx: ActionContext) -> Tuple[Optional[st return e.value, items +def _ask_and_check(ctx: ActionContext, is_remote: bool, question: str, permission_type: str, + value: str, deny_message: str, command: Optional[str] = None, + reason: Optional[str] = None): + """Ask the user/backend to approve one "ask" guardrail layer (shell/target/path). + + Builds the common ask_kwargs, emits the remote pending-prompt (if any), then + calls ``ctx.backend.ask_user()`` and checks approval. This is a generator so a + remote backend's pending prompt can be yielded up through the caller's + ``yield from``. Returns ``None`` if approved, else ``deny_message``. + """ + ask_kwargs = dict( + question=question, + choices=["allow", "allow_all", "deny"], + session_id=ctx.session_id, + prompt_type="permission", + permission_type=permission_type, + value=value, + engine=ctx.permission_engine, + # unique id per prompt so its remote poll matches only its own answer + prompt_uuid=str(uuid.uuid4()), + ) + if command is not None: + ask_kwargs["command"] = command + if reason is not None: + ask_kwargs["reason"] = reason + if is_remote: + yield ctx.backend.build_pending_prompt(**ask_kwargs) + response = ctx.backend.ask_user(**ask_kwargs) if ctx.backend else None + if not _is_approved(response): + return deny_message + return None + + def check_guardrails(action: Dict, ctx: ActionContext): """Check action against guardrails before dispatching. @@ -146,11 +245,8 @@ def check_guardrails(action: Dict, ctx: ActionContext): is_remote = isinstance(ctx.backend, RemoteBackend) - # Prompt loop: each check_action returns the first "ask" it encounters - # (shell, then targets, then paths). We prompt for that layer, then re-check - # to surface the next layer, until everything is resolved. - # All prompting goes through ctx.backend.ask_user() — the backend handles - # the UX differences (CLI menu, DB polling, or auto-deny). + # Prompt loop: check_action returns the first unresolved "ask" layer (shell, then + # targets, then paths); prompt via ctx.backend.ask_user() and re-check until resolved. max_rounds = 5 rounds = 0 while result.decision == "ask" and rounds < max_rounds: @@ -160,21 +256,16 @@ def check_guardrails(action: Dict, ctx: ActionContext): # Handle shell command prompts (unknown commands or parse failures) if result.shell_command: parse_failed = "Could not parse" in (result.reason or "") - ask_kwargs = dict( + denial = yield from _ask_and_check( + ctx, is_remote, question=result.reason or "Shell command requires approval", - choices=["allow", "allow_all", "deny"], - session_id=ctx.session_id, - prompt_type="permission", permission_type="shell", value=result.shell_command, + deny_message="Action denied: shell command not approved", reason=result.reason, - engine=ctx.permission_engine, ) - if is_remote: - yield ctx.backend.build_pending_prompt(**ask_kwargs) - response = ctx.backend.ask_user(**ask_kwargs) if ctx.backend else None - if response is None or response.get("answer") == "deny": - return "Action denied: shell command not approved" + if denial: + return denial if parse_failed: return None @@ -183,21 +274,16 @@ def check_guardrails(action: Dict, ctx: ActionContext): recheck = ctx.permission_engine._check_value("target", target) if recheck.decision == "allow": continue - ask_kwargs = dict( + denial = yield from _ask_and_check( + ctx, is_remote, question=f"Target {target} requires approval", - choices=["allow", "allow_all", "deny"], - session_id=ctx.session_id, - prompt_type="permission", permission_type="target", value=target, + deny_message=f"Action denied: target {target} not approved", command=cmd_display, - engine=ctx.permission_engine, ) - if is_remote: - yield ctx.backend.build_pending_prompt(**ask_kwargs) - response = ctx.backend.ask_user(**ask_kwargs) if ctx.backend else None - if response is None or response.get("answer") == "deny": - return f"Action denied: target {target} not approved" + if denial: + return denial # Handle path prompts if result.paths: @@ -205,27 +291,26 @@ def check_guardrails(action: Dict, ctx: ActionContext): path_access_map = {p: a for p, a in detect_paths_with_access(cmd)} for path in result.paths: access_type = path_access_map.get(path, "read") - ask_kwargs = dict( + denial = yield from _ask_and_check( + ctx, is_remote, question=f"{access_type.capitalize()} access to {path} requires approval", - choices=["allow", "allow_all", "deny"], - session_id=ctx.session_id, - prompt_type="permission", permission_type=access_type, value=path, + deny_message=f"Action denied: {access_type} access to {path} not approved", command=cmd_display, - engine=ctx.permission_engine, ) - if is_remote: - yield ctx.backend.build_pending_prompt(**ask_kwargs) - response = ctx.backend.ask_user(**ask_kwargs) if ctx.backend else None - if response is None or response.get("answer") == "deny": - return f"Action denied: {access_type} access to {path} not approved" + if denial: + return denial # Re-check to see if more layers need prompting result = ctx.permission_engine.check_action(action) if result.decision == "deny": return f"Action denied after prompt: {result.reason}" + # fail closed: prompts exhausted with the decision still unresolved -> block + if result.decision == "ask": + return f"Action denied: guardrail check unresolved after {max_rounds} prompts" + return None @@ -259,6 +344,118 @@ def dispatch_action(action: Dict, ctx: ActionContext) -> Generator: yield Warning(message=f"Unknown action: {action_type}", _context=context) +def safe_dispatch_action(action: Dict, ctx: ActionContext) -> Generator: + """Dispatch a single action, converting any raised ``Exception`` into an + ``Error`` output item instead of letting it abort the AI loop. + + A Python error during a handler (e.g. ``TypeError: 'str' object is not a + mapping`` from a malformed LLM action/opts) must NOT kill the main loop. We + wrap the per-action generator so the failure becomes an ``Error`` carrying + the action's ``tool_call_id``/``tool_call_name`` in ``_context`` — that lets + the caller group it into a tool result and feed the error back to the LLM so + it can correct itself on the next turn. + + Only ``Exception`` is caught: ``KeyboardInterrupt`` / ``SystemExit`` / + ``GeneratorExit`` (all ``BaseException`` subclasses) propagate so legitimate + control-flow and generator close are never swallowed. + """ + import traceback as _traceback + try: + yield from dispatch_action(action, ctx) + except Exception as e: # noqa: BLE001 - per-action resilience: feed error back to LLM, never abort the loop + context = _get_result_context(action, ctx) + yield Error( + message=_format_action_error(e), + traceback=_traceback.format_exc(), + _context=context, + ) + + +def _guard_subagent_fanout(ctx: "ActionContext", context: Dict) -> Optional["Warning"]: + """Cap AI-subagent recursion depth + per-turn fan-out. + + Returns a denial ``Warning`` if a cap is hit (caller yields it and skips the + spawn); otherwise stamps the child's depth (+1) into ``context`` and bumps the + per-turn counter. Breadth is only counted within a batch (one LLM turn); a + lone spawn is inherently breadth-1. + """ + depth = int(ctx.context.get("ai_subagent_depth", 0) or 0) + if depth >= _MAX_SUBAGENT_DEPTH: + return Warning( + message=f"Subagent spawn denied: recursion depth cap ({_MAX_SUBAGENT_DEPTH}) reached", + _context=context, + ) + if ctx.in_batch: # per-turn breadth only bites within a batch + with _SUBAGENT_TURN_LOCK: + turn = int(ctx.context.get("ai_subagent_turn_count", 0) or 0) + over_breadth = turn >= _MAX_SUBAGENTS_PER_TURN + if not over_breadth: + ctx.context["ai_subagent_turn_count"] = turn + 1 + if over_breadth: + return Warning( + message=f"Subagent spawn denied: per-turn fan-out cap ({_MAX_SUBAGENTS_PER_TURN}) reached", + _context=context, + ) + context["ai_subagent_depth"] = depth + 1 # child inherits depth+1 + return None + + +def _gather_subagent_evidence(ctx: "ActionContext", targets: list, limit: int = 40) -> str: + """Auto-assemble prior findings for the subagent's targets so it doesn't redo work. + + Queries the workspace (the single source of truth — incl. this run's live findings) + for findings whose host/ip/url match any target, capped at `limit`. Best-effort: + any failure returns "" (evidence is a nicety, never a blocker). + """ + targets = [t for t in (targets or []) if t] + if not targets: + return "" + query = {"$or": [{"host": {"$in": targets}}, {"ip": {"$in": targets}}, {"url": {"$in": targets}}]} + try: + results = ctx.get_query_engine().search(query, limit=limit) or [] + except Exception: # noqa: BLE001 - evidence is best-effort; never break the spawn + return "" + lines = [] + for r in results[:limit]: + d = r.toDict() if hasattr(r, "toDict") else r + t = d.get("_type", "finding") + key = d.get("url") or d.get("matched_at") or f"{d.get('ip', '') or d.get('host', '')}" + extra = f":{d.get('port')}" if d.get("port") else "" + name = f" {d.get('name')}" if d.get("name") else "" + lines.append(f"- {t} {key}{extra}{name}".rstrip()) + return "\n".join(lines) + + +def _child_run_opts(ctx: ActionContext) -> Dict: + """Common run_opts shared by every child runner (task/workflow/shell command).""" + return { + "print_item": not ctx.silent, + "print_line": ctx.verbose and not ctx.silent, + "print_progress": False, + "print_reports_message": False, + "enable_reports": True, + "exporters": [], + "sync": ctx.sync, + } + + +def _child_preamble(ctx: ActionContext, context: Dict) -> Tuple[Dict, Optional["Warning"]]: + """Shared child-runner prelude: stamp task_chunk_id + subagent flag, then rebuild + persistence hooks (or return a denial). + + Propagates driver hooks (mongodb/api): a sync sub-runner skips the pickle path + that normally re-registers them, so without this its results never persist. + Don't silently spawn a persistence-less child when the parent has drivers. + + Returns ``(hooks, denial)``; if ``denial`` is non-None the caller must yield it + and skip the spawn. + """ + context["task_chunk_id"] = str(uuid.uuid4()) + if ctx.subagent: + context["subagent"] = ctx.context.get("subagent", True) + return _build_child_hooks_or_denial(context) + + def _run_runner(action: Dict, ctx: ActionContext, runner_type: str) -> Generator: """Execute a secator task or workflow. @@ -269,13 +466,34 @@ def _run_runner(action: Dict, ctx: ActionContext, runner_type: str) -> Generator """ name = action.get("name", "") targets = action.get("targets", ctx.targets) - opts = action.get("opts", {}) + # drop LLM-set control keys (notably `dangerous`) before they reach the child + opts = _sanitize_child_opts(action.get("opts", {})) context = _get_result_context(action, ctx) # Force subagent flags when spawning an AI task from a parent AI task if runner_type == "task" and name.lower() == "ai": + # Bound recursive fan-out before constructing/running the child + denial = _guard_subagent_fanout(ctx, context) + if denial is not None: + yield denial + return opts["subagent"] = True opts["interactive"] = False + # Inherit the parent's resolved LLM config (else it falls back to the default + # model/provider with no key set -> AuthenticationError). setdefault so an + # explicit LLM-supplied model/key still wins. + opts.setdefault("model", ctx.model) + if ctx.api_key: + opts.setdefault("api_key", ctx.api_key) + if ctx.api_base: + opts.setdefault("api_base", ctx.api_base) + # 1.b/1.c: structure the subagent's prompt and inject prior findings for its + # scope so it doesn't re-run work already done. + _objective = opts.get("prompt", "") + opts["prompt"] = build_subagent_prompt(_objective, targets, _gather_subagent_evidence(ctx, targets)) + + # defense in depth: a spawned runner is never dangerous (CLI --dangerous unaffected) + opts["dangerous"] = False if runner_type == "task": tpl = TemplateLoader(input={'type': 'task', 'name': name}) @@ -292,19 +510,10 @@ def _run_runner(action: Dict, ctx: ActionContext, runner_type: str) -> Generator yield Info(message=f"[DRY RUN] Would run {runner_type}: {name} on {targets}", _context=context) return - if not ctx.silent: - yield Ai(content=name, ai_type=runner_type, extra_data={"targets": targets, "opts": opts}, _context=context) - run_opts = { - "print_item": not ctx.silent, - "print_line": ctx.verbose and not ctx.silent, + **_child_run_opts(ctx), "print_cmd": not ctx.silent and not ctx.subagent, "print_cmd_icon": "└", - "print_progress": False, - "print_reports_message": False, - "enable_reports": True, - "exporters": [], - "sync": ctx.sync, "tty": not ctx.subagent and ctx.sync, **opts, } @@ -312,14 +521,40 @@ def _run_runner(action: Dict, ctx: ActionContext, runner_type: str) -> Generator run_opts["print_start"] = not ctx.silent and not ctx.subagent run_opts["print_end"] = not ctx.silent and not ctx.subagent - context["task_chunk_id"] = str(uuid.uuid4()) - if ctx.subagent: - context["subagent"] = ctx.context.get("subagent", True) + # A heavy sub-task (e.g. nuclei) must not run sync in the ai task's small worker + # pool (OOM risk) — dispatch it async to its own profile's queue when in a worker. + if run_opts.get("sync") and _is_heavy_runner(runner_type, name, opts): + from secator.celery import IN_WORKER + if IN_WORKER: + run_opts["sync"] = False + run_opts["tty"] = False + + hooks, denial = _child_preamble(ctx, context) + if denial is not None: + yield denial + return try: - runner = runner_cls(tpl, targets, run_opts=run_opts, context=context) + runner = runner_cls(tpl, targets, run_opts=run_opts, hooks=hooks, context=context) except TaskNotFoundError as e: yield Error(message=str(e), _context=context) return + + # Emit the action Ai item now the runner exists (on_init stamped the runner id) so + # the UI can render a RunnerCard; always emitted, even when silent. Prefer context + # `{type}_id` (the persisted doc's `_id`) over `runner.id` (internal, doesn't match). + runner_id = context.get(f"{runner_type}_id", "") or runner.id + yield Ai( + content=name, + ai_type=runner_type, + extra_data={ + "targets": targets, + "opts": opts, + "runner_id": runner_id, + "runner_type": runner_type, + }, + _context=context, + ) + yield from runner # Auto-allow reading from the spawned runner's reports folder @@ -329,15 +564,28 @@ def _run_runner(action: Dict, ctx: ActionContext, runner_type: str) -> Generator def _get_result_context(action, ctx): - """Get result context from action""" - ctx = ctx.context.copy() + """Build the CHILD runner's context. + + Stamps the conversation ``session_id`` (parenting link — see the runner-parenting + design) and marks the child ``has_parent``. Critically, it STRIPS the parent's + runner-identity keys (`task_id`/`workflow_id`/`scan_id`): a child that inherited + them would make `update_runner`/`runner_id` target the PARENT's doc instead of + minting its own. The child keeps drivers/workspace so it persists into the same + workspace, linked to the conversation by ``session_id``. + """ + new_ctx = ctx.context.copy() + for identity_key in ("task_id", "workflow_id", "scan_id", "task_chunk_id"): + new_ctx.pop(identity_key, None) + if ctx.session_id and not new_ctx.get("session_id"): + new_ctx["session_id"] = ctx.session_id + new_ctx["has_parent"] = True action_context = {} tool_call_id = action.get("tool_call_id") tool_call_name = action.get("tool_call_name") if tool_call_id: action_context["tool_call_id"] = tool_call_id action_context["tool_call_name"] = tool_call_name - return {**ctx, **action_context} + return {**new_ctx, **action_context} def _handle_task(action: Dict, ctx: ActionContext) -> Generator: @@ -351,7 +599,14 @@ def _handle_workflow(action: Dict, ctx: ActionContext) -> Generator: def _handle_shell(action: Dict, ctx: ActionContext) -> Generator: - """Execute a shell command. + """Execute a shell command as a `command` task runner. + + Dispatches the built-in `command` task (a Command subclass that runs an arbitrary + shell command line verbatim) through the normal runner lifecycle, instead of a raw + `subprocess.run`. This makes the shell invocation persist as a runner doc (via the + driver hooks rebuilt from `context['drivers']`) and appear in history, parented + under the conversation via `context['session_id']` — exactly like `_run_runner` + does for AI-spawned tasks/workflows. Args: action: Action dict with command @@ -367,18 +622,58 @@ def _handle_shell(action: Dict, ctx: ActionContext) -> Generator: yield Info(message=f"[DRY RUN] Would run: {command}", _context=context) return - yield Ai(content=command, ai_type="shell", _context=context) - try: - result = subprocess.run( - command, - shell=True, - capture_output=True, - text=True, - timeout=60, - env=_sanitized_env() + # Don't silently run a persistence-less child when the parent has drivers + # (same guard _run_runner uses for spawned tasks/workflows). + hooks, denial = _child_preamble(ctx, context) + if denial is not None: + yield denial + return + + # hooks is CLASS-keyed ({Task: {...}}); we bypass the Task wrapper with a direct + # `command(...)` instantiation, so extract hooks[Task] ourselves — else register_hooks + # finds no match and the runner doc is silently never persisted (no error, no doc). + hooks = hooks.get(Task, {}) + + # Mirrors _run_runner's wiring: quiet, reports enabled, never dangerous (defense + # in depth). `env` is the sanitized process env so `env`/`printenv` can't leak secrets. + run_opts = { + **_child_run_opts(ctx), + "print_cmd": False, + "dangerous": False, + "env": _sanitized_env(), + } + + # Instantiate `command` directly (bypasses the Task wrapper, which discards `.output`) + # so stdout survives while persist hooks still fire. Spread **run_opts, not `run_opts=` + # (would nest and drop `env`); import locally to avoid a circular import. + from secator.tasks.command import command as CommandTask + runner = CommandTask([command], hooks=hooks, context=context, **run_opts) + + # 60s cap on ad-hoc AI shell commands. max_timeout is NOT run_opts-settable + # (Command.__init__ resolves it from CONFIG.tasks.overrides); setting the + # instance attribute here is honored by get_max_timeout(). + runner.max_timeout = _SHELL_TIMEOUT + + # Emit the command Ai now that the runner exists: its on_init hook has + # stamped the runner id into context, so the UI can link this item to the + # persisted runner doc (mirrors _run_runner:688-699). + yield Ai( + content=command, + ai_type="shell", + extra_data={ + "runner_id": context.get("task_id", "") or runner.id, + "runner_type": "task", + }, + _context=context, ) - output = result.stdout or result.stderr or "(no output)" + + # Run to completion in-process (fires persist hooks like a normal task/workflow). + # Do NOT `yield from runner` — raw stdout lines aren't separate transcript items; + # the single shell_output below is the contract. + runner.run() + + output = _truncate(runner.output or "(no output)", _MAX_SHELL_OUTPUT_CHARS) # cap so it can't blow up history yield Ai(content=output, ai_type="shell_output", _context=context) except Exception as e: @@ -394,20 +689,55 @@ def _handle_query(action: Dict, ctx: ActionContext) -> Generator: """ context = _get_result_context(action, ctx) query_filter = action.get("query", {}) + # The schema declares `limit` an integer, but some models send it as a string + # ("10"); a str limit reaches the backend and raises `'>=' not supported between + # int and str`. Coerce to int (bad/None values fall back to the default). limit = action.get("limit", 100) + try: + limit = int(limit) + except (TypeError, ValueError): + limit = 100 + + # Some providers serialize `query` as a JSON string despite the object schema + # (known tool-calling quirk); coerce it back, else fail with a clear LLM error. + if isinstance(query_filter, str): + try: + query_filter = json.loads(query_filter) + except (json.JSONDecodeError, TypeError): + yield Error( + message='query must be a JSON object (e.g. {"_type": "vulnerability"}); ' + f'got an unparseable string: {query_filter[:120]!r}', + _context=context, + ) + return + if not isinstance(query_filter, dict): + yield Error( + message=f'query must be a JSON object; got {type(query_filter).__name__}.', + _context=context, + ) + return # Decrypt query values if ctx.encryptor: query_filter = _decrypt_dict(query_filter, ctx.encryptor) - if ctx.scope != "current" and not ctx.context.get("workspace_id"): + engine = ctx.get_query_engine() + is_local = getattr(engine.backend, "name", "") == "json" + + # A non-local backend (mongodb/api) needs a workspace to query. The local (json) + # driver can always answer from this run's in-memory findings (unioned below), so + # it is exempt from the workspace_id requirement. + if not is_local and ctx.scope != "current" and not ctx.context.get("workspace_id"): yield Warning(message="No workspace available for query", _context=context) return try: query_str = json.dumps(query_filter, separators=(',', ':')) - engine = ctx.get_query_engine() results = engine.search(query_filter, limit=limit) + # Local driver only writes to disk at end-of-run, so union in-memory live results + # to make query_workspace the source of truth (mongodb/api persist live already). + if is_local and ctx.scope != "current": + results = _union_live_results(results, ctx.results or [], query_filter, limit) yield Ai( content=query_str, ai_type="query", @@ -444,7 +774,10 @@ def _handle_follow_up(action: Dict, ctx: ActionContext) -> Generator: context = _get_result_context(action, ctx) reason = action.get("reason", "completed") choices = action.get("choices", []) - yield Ai(content=reason, ai_type="follow_up", extra_data={"choices": choices}, _context=context) + # Store choices on the top-level `choices` field (what the web UI reads) AND in + # extra_data (back-compat). Without the top-level field, the persisted follow-up + # doc has `choices: []` and the UI renders no choice buttons. + yield Ai(content=reason, ai_type="follow_up", choices=choices, extra_data={"choices": choices}, _context=context) def _handle_stop(action: Dict, ctx: ActionContext) -> Generator: @@ -507,6 +840,10 @@ def _handle_add_finding(action: Dict, ctx: ActionContext) -> Generator: extra.update(unknown) finding_data['extra_data'] = extra + # Coerce AI-provided scalars to declared field types (LLMs send wrong-typed + # scalars, e.g. a bool field as the string "true") before validating. + finding_data = _coerce_finding_fields(cls, finding_data) + # Validate field types before instantiation errors = cls.validate_fields(finding_data) if errors: @@ -519,6 +856,9 @@ def _handle_add_finding(action: Dict, ctx: ActionContext) -> Generator: yield Ai( content=f'{str(finding)}', ai_type="add_finding", + # Carry the created finding so the web UI can render its FindingCard + # (VulnerabilityCard/SubdomainCard/…) — it routes on `_type`. + extra_data={"finding": finding.toDict()}, _context=context ) yield finding @@ -526,24 +866,6 @@ def _handle_add_finding(action: Dict, ctx: ActionContext) -> Generator: yield Error(message=f"Failed to create {finding_type}: {e}\nExpected schema:\n{cls.schema()}", _context=context) -def _get_action_label(action: Dict) -> str: - """Get a display label for an action.""" - act_type = action.get("action", "unknown") - if act_type in ("task", "workflow"): - name = action.get("name", "?") - opts = action.get("opts", {}) - session_name = opts.get("session_name", "") - if session_name: - return session_name - targets = action.get("targets", []) - target_str = targets[0] if len(targets) == 1 else f"{len(targets)} targets" - return f"{name} on {target_str}" - elif act_type == "shell": - cmd = action.get("command", "")[:40] - return f"shell: {cmd}" - return act_type - - def _run_batch(actions: List[Dict], ctx: ActionContext) -> Generator: """Execute multiple actions in parallel with Rich progress display. @@ -569,8 +891,11 @@ def _run_batch(actions: List[Dict], ctx: ActionContext) -> Generator: max_workers = ctx.max_workers or 3 + # Fresh per-turn subagent fan-out budget for this batch (one LLM turn) + ctx.context["ai_subagent_turn_count"] = 0 + # Silence console output for parallel tasks to avoid interleaved printing - batch_ctx = replace(ctx, silent=True) + batch_ctx = replace(ctx, silent=True, in_batch=True) # Skip Rich progress panel when we are a subagent, or when the batch # contains an AI subagent task (its output conflicts with the Live display) @@ -593,8 +918,10 @@ def _run_batch(actions: List[Dict], ctx: ActionContext) -> Generator: progress_ids = {} def run_single(act: Dict, idx: int) -> Dict: + # safe_dispatch_action so one action raising doesn't abort the batch — the + # error becomes an Error item (tagged with tool_call_id) fed back to the LLM. results = [] - for item in dispatch_action(act, batch_ctx): + for item in safe_dispatch_action(act, batch_ctx): if isinstance(item, Ai) and item.ai_type == "token_usage": if progress: extra = item.extra_data or {} @@ -667,31 +994,3 @@ def get_renderables(self): for idx, result in sorted(all_results, key=lambda x: x[0]): for item in result["results"]: yield item - - -def _decrypt_dict(d: Dict, encryptor: Any) -> Dict: - """Recursively decrypt all string values in a dict. - - Args: - d: Dictionary to decrypt - encryptor: SensitiveDataEncryptor instance - - Returns: - Decrypted dictionary - """ - result = {} - for k, v in d.items(): - if isinstance(v, str): - result[k] = encryptor.decrypt(v) - elif isinstance(v, dict): - result[k] = _decrypt_dict(v, encryptor) - elif isinstance(v, list): - result[k] = [ - encryptor.decrypt(i) if isinstance(i, str) - else _decrypt_dict(i, encryptor) if isinstance(i, dict) - else i - for i in v - ] - else: - result[k] = v - return result diff --git a/secator/ai/encryption.py b/secator/ai/encryption.py index 4aa3ed3f0..ff897d197 100644 --- a/secator/ai/encryption.py +++ b/secator/ai/encryption.py @@ -129,14 +129,10 @@ def decrypt(self, text: str) -> str: """Restore original sensitive values from placeholders.""" result = text - # Full placeholders [TYPE:hash] + # Full placeholders [TYPE:hash] and their bracket-stripped form TYPE:hash for placeholder, original in self.pii_map.items(): result = result.replace(placeholder, original) - - # Without brackets TYPE:hash - for placeholder, original in self.pii_map.items(): - no_brackets = placeholder[1:-1] - result = result.replace(no_brackets, original) + result = result.replace(placeholder[1:-1], original) # Bare hashes for hash_value, original in self.hash_map.items(): diff --git a/secator/ai/guardrails.py b/secator/ai/guardrails.py index 37e21029a..375ed4263 100644 --- a/secator/ai/guardrails.py +++ b/secator/ai/guardrails.py @@ -1,10 +1,11 @@ """Permission engine for AI guardrails.""" import fnmatch +import ipaddress import re import socket from dataclasses import dataclass, field from functools import lru_cache -from typing import Dict, List, Tuple +from typing import Dict, List, Optional, Tuple, Union from secator.ai.encryption import PII_PATTERNS @@ -26,6 +27,60 @@ # Execute-type commands EXECUTE_COMMANDS = frozenset({"python", "python3", "bash", "sh", "node", "ruby", "perl", "gcc", "g++", "make", "go"}) +# Download tools that write to a file via an OUTPUT FLAG — the flag's destination +# is a WRITE, not a read (else `deny write(/etc/*)` never fires). Focused set; residual +# write-vs-read gaps (dd of=, tar -f, cp/install dest, >() ) are tracked separately. +OUTPUT_FLAG_COMMANDS = { + "curl": frozenset({"-o", "--output"}), + "wget": frozenset({"-O", "--output-document"}), +} + +# Exec-wrappers run a different inner command (`timeout 60 rm -rf /`) — peel the +# wrapper and check the INNER command, not the allow-listed wrapper name. +EXEC_WRAPPERS = frozenset({ + "timeout", "xargs", "env", "nice", "ionice", "nohup", "stdbuf", + "setsid", "sudo", "doas", "watch", "time", "chroot", "unbuffer", + # Laundering-vector wrappers + "flock", "runuser", "su", "script", "proxychains", "proxychains4", + "firejail", "torsocks", "torify", "unshare", "catchsegv", "chrt", "taskset", +}) + +# Per-wrapper arg grammar so the REAL command is located, not a lockfile/config/user. +# (opts_taking_a_value, positional_args_before_cmd, cmd_string_opts) — cmd_string_opts values +# (e.g. `-c 'curl evil'`) are re-parsed and peeled so the payload is checked, not skipped. +_EMPTY = frozenset() +_WRAPPER_ARG_GRAMMAR = { + "flock": (frozenset({"-w", "--timeout", "-E", "--conflict-exit-code"}), 1, frozenset({"-c", "--command"})), + "runuser": (frozenset({"-u", "--user", "-g", "--group", "-G", "--supp-group", "-s", "--shell"}), 0, frozenset({"-c", "--command"})), # noqa: E501 + "su": (frozenset({"-s", "--shell", "-g", "--group", "-G", "--supp-group"}), 1, frozenset({"-c", "--command"})), # noqa: E501 + "script": (_EMPTY, 0, frozenset({"-c", "--command"})), + "proxychains": (frozenset({"-f"}), 0, _EMPTY), + "proxychains4": (frozenset({"-f"}), 0, _EMPTY), + "sudo": (frozenset({"-u", "--user", "-g", "--group", "-U", "-C", "-p", "-r", "-t", "-T"}), 0, _EMPTY), +} + + +def _exec_wrappers() -> frozenset: + """Built-in wrappers plus any ops-configured extras. Config EXTENDS the security baseline.""" + try: + from secator.config import CONFIG + extra = getattr(CONFIG.addons.ai, "exec_wrappers", None) or [] + extra = {str(w).strip() for w in extra if str(w).strip()} + if extra: + return EXEC_WRAPPERS | extra + except Exception: + pass + return EXEC_WRAPPERS + + +def _split_cmd_string(s: str) -> List[str]: + """Best-effort tokenize a `-c ''` payload so the nested command can be re-checked.""" + import shlex + try: + return shlex.split(s) + except ValueError: + return s.split() + def parse_rule(rule: str) -> Tuple[str, List[str]]: """Parse a rule string like 'target(10.0.0.1,example.com)' into (type, patterns). @@ -44,6 +99,59 @@ def parse_rule(rule: str) -> Tuple[str, List[str]]: return rule_type, values +_IP_INT_RE = re.compile(r'0[xX][0-9a-fA-F]+|0[oO][0-7]+|\d+') +_DOTTED_ODD_RE = re.compile(r'(?:0[xX][0-9a-fA-F]+|0[0-7]+|\d+)(?:\.(?:0[xX][0-9a-fA-F]+|0[0-7]+|\d+)){3}') +IPAddress = Union[ipaddress.IPv4Address, ipaddress.IPv6Address] + + +def _normalize_ip(candidate: str) -> Optional[IPAddress]: + """Normalize encoded IPs (decimal/hex/octal int, dotted-hex/octal, IPv6-mapped) to an ip_address. + + Returns None if the candidate is not an IP (e.g. a hostname) so callers fall back to literal matching. + Hostnames are NOT resolved here (DNS rebinding is a documented residual). + """ + s = candidate.strip() + if not s: + return None + if s.startswith('[') and s.endswith(']'): # [::1] / [::ffff:1.2.3.4] + s = s[1:-1] + ip = None + # Plain dotted-quad / standard IPv6 first (leaves normal targets untouched) + try: + ip = ipaddress.ip_address(s) + except ValueError: + # Integer form: decimal (2852039166), hex (0xA9FEA9FE), octal (0o...) + if _IP_INT_RE.fullmatch(s): + try: + ip = ipaddress.ip_address(int(s, 0) if s[:2].lower() in ('0x', '0o') else int(s)) + except (ValueError, ipaddress.AddressValueError): + return None + # Dotted octets with hex/octal parts (0xA9.0xFE.0xA9.0xFE, 0251.0376.0251.0376) + elif _DOTTED_ODD_RE.fullmatch(s): + try: + octets = [int(p, 0) if p[:2].lower() == '0x' else int(p, 8) if p.startswith('0') and len(p) > 1 else int(p) + for p in s.split('.')] + if all(0 <= o <= 255 for o in octets): + ip = ipaddress.ip_address('.'.join(str(o) for o in octets)) + except (ValueError, ipaddress.AddressValueError): + return None + if ip is None: + return None + # Collapse IPv6-mapped/compatible IPv4 (::ffff:169.254.169.254) down to the v4 address + if isinstance(ip, ipaddress.IPv6Address) and ip.ipv4_mapped is not None: + ip = ip.ipv4_mapped + return ip + + +def _ip_in_pattern(ip: IPAddress, pattern: str) -> Optional[bool]: + """True/False if `pattern` is an IP/CIDR literal, else None (pattern isn't an address rule).""" + try: + net = ipaddress.ip_network(pattern, strict=False) + except ValueError: + return None + return ip.version == net.version and ip in net + + def match_rule(value: str, patterns: List[str]) -> bool: """Check if a value matches any of the given patterns. @@ -53,6 +161,7 @@ def match_rule(value: str, patterns: List[str]) -> bool: - Glob patterns (fnmatch) - {port} variable (matches :\\d+) - Basename matching for path-like values (e.g. '.env' matches '/home/user/.env') + - IP/CIDR patterns are matched by normalized address (encoded IPs are canonicalized first) Args: value: The value to check @@ -61,9 +170,21 @@ def match_rule(value: str, patterns: List[str]) -> bool: Returns: True if value matches any pattern """ + # Normalize encoded IPs before deny/allow match so alternate encodings can't evade IP rules + norm_ip = _normalize_ip(value) + canon = str(norm_ip) if norm_ip is not None else None for pattern in patterns: if pattern == "*": return True + if norm_ip is not None: + in_pat = _ip_in_pattern(norm_ip, pattern) + if in_pat is not None: + if in_pat: + return True + continue # IP/CIDR pattern that doesn't contain this address — no string fallback + # Non-address pattern (glob/{port}): also test the canonical dotted form + if canon != value and fnmatch.fnmatch(canon, pattern): + return True if "{port}" in pattern: regex_pattern = re.escape(pattern).replace(r"\{port\}", r"\d+") if re.fullmatch(regex_pattern, value): @@ -108,10 +229,8 @@ def _is_file_path(value: str) -> bool: def _is_network_target(value: str) -> bool: - """Check if a value looks like a valid network target (IP, hostname, URL, CIDR). - - Filters out descriptive strings that aren't actual targets. - """ + """Check if a value looks like a valid network target (IP, hostname, URL, CIDR), + filtering out descriptive strings that aren't actual targets.""" if ' ' in value.strip(): return False if value.startswith(('http://', 'https://')): @@ -226,37 +345,131 @@ def _check_arg(arg: str): return targets -def _extract_cmd_names(command: str) -> List[str]: - """Extract command names from a shell command using safecmd's bash parser. +_SHELL_PARSER_WARNED = False + + +def _warn_shell_parser_unavailable(reason: str) -> None: + """Warn ONCE that the shfmt-based shell parser is unavailable, then let the + caller fall back to the non-shfmt path (whole-command approval). + + This is deliberately a Warning, not an Error, and it does NOT claim the ai + addon is missing: ``litellm`` (the ai addon) can be installed while the shell + parser — ``safecmd`` + the ``shfmt`` binary it shells out to — is not. Without + it the guardrail can't split a command into sub-commands, so + ``_check_action_type`` falls back to asking the user to approve the whole + command (safe, just coarser). Warn once so a long agent run isn't spammed on + every shell command. + """ + global _SHELL_PARSER_WARNED + if _SHELL_PARSER_WARNED: + return + _SHELL_PARSER_WARNED = True + from secator.rich import console + from secator.output_types import Warning + console.print(Warning( + message=f'{reason}: shell commands cannot be sub-parsed for guardrails — ' + 'falling back to whole-command approval. Run "secator install addons ai" ' + 'to enable precise per-subcommand parsing.' + )) + + +def _parse_subcommands(command: str) -> List[List[str]]: + """Parse a shell command into sub-command token lists via safecmd's parser. Uses shfmt (via safecmd) to properly parse pipes, &&, ||, ;, subshells, - and command substitutions. Returns empty list if parsing fails (caller + and command substitutions. Returns an empty list if parsing fails (caller should prompt the user to approve the whole command). Args: command: Full shell command string Returns: - List of command name strings (first token of each sub-command), - or empty list if parsing fails. + List of token lists, one per sub-command, or [] if parsing fails. """ - import re try: from safecmd.bashxtract import extract_commands except ImportError: - from secator.rich import console - console.print('[bold red][ERR][/] Missing ai addon: please run "secator install addons ai".') + # NOT a missing *ai* addon (litellm can be present without the shell parser). + _warn_shell_parser_unavailable('Missing safecmd shell parser') return [] try: # Normalize LLM-generated multiline commands: join lines where a pipe/operator # starts the next line (e.g. "cmd1\n| cmd2" -> "cmd1 | cmd2") command = re.sub(r'\s*\n\s*(\||\&\&|\|\|)', r' \1', command) - cmds, ops, redirects = extract_commands(command) - return [c[0] for c in cmds if c] + cmds, _, _ = extract_commands(command) + return [c for c in cmds if c] + except FileNotFoundError: + # safecmd is installed but the `shfmt` binary it shells out to isn't on PATH. + _warn_shell_parser_unavailable('Missing shfmt binary') + return [] except Exception: return [] +def _extract_cmd_names(command: str) -> List[str]: + """Extract command names (first token of each sub-command); [] on parse failure.""" + return [c[0] for c in _parse_subcommands(command)] + + +def _is_wrapper_operand(token: str) -> bool: + """Heuristic: is this token a wrapper operand (numeric duration / KEY=VALUE), not the inner cmd?""" + if re.fullmatch(r'\d+(?:\.\d+)?[smhd]?', token): + return True + if re.fullmatch(r'[A-Za-z_][A-Za-z0-9_]*=.*', token): + return True + return False + + +def _peel_wrapper(args: List[str]) -> List[str]: + """Strip leading exec-wrapper binaries to reach the inner command's tokens. + + Bare `env`/`sudo` (no inner command) is returned as-is so it's still checked by name. + """ + wrappers = _exec_wrappers() + tokens = args + for _ in range(len(args)): # bounded peels (guards against pathological nesting) + if not tokens: + return tokens + name = tokens[0].rsplit('/', 1)[-1] + if name not in wrappers: + return tokens + rest = tokens[1:] + # Peel proxychains/firejail/flock/runuser/... past their OWN args (value-opts, + # positional lockfile/config, `-c ''`) so the leaf payload is what gets classified. + opts_with_val, n_pos, cmd_opts = _WRAPPER_ARG_GRAMMAR.get(name, (_EMPTY, 0, _EMPTY)) + i = 0 + pos_seen = 0 + while i < len(rest): + tok = rest[i] + if tok == '--': # end-of-options: the inner command starts next + i += 1 + break + if tok in cmd_opts and i + 1 < len(rest): # `-c ''` — re-parse & peel the nested payload + nested = _split_cmd_string(rest[i + 1]) + return _peel_wrapper(nested) if nested else tokens + if tok in opts_with_val and i + 1 < len(rest): # option that consumes its value + i += 2 + continue + if tok.startswith('-') or _is_wrapper_operand(tok): + i += 1 + continue + if pos_seen < n_pos: # wrapper's own positional (flock lockfile / su user) + pos_seen += 1 + i += 1 + continue + break + if i >= len(rest): + return tokens # wrapper with no inner command — check it by name + tokens = rest[i:] + return tokens + + +def _match_command_glob(command: str, pattern: str) -> bool: + """Anchored glob match where '*' does NOT cross '/' (so `rm -rf /*` spares `rm -rf /tmp/x`).""" + regex = ''.join('[^/]*' if ch == '*' else re.escape(ch) for ch in pattern) + return re.fullmatch(regex, command) is not None + + def _resolve_path(path: str, cwd: str = "") -> str: """Resolve a path to absolute for consistent rule matching. @@ -359,11 +572,31 @@ def _extract_docker_volumes(args: List[str]): cmd_class = classify_command(cmd_name) base_access = "write" if cmd_class == "write" else "read" - for arg in args[1:]: - if arg.startswith('-'): - continue - if _is_file_path(arg): + # Output-flag destinations are writes (curl -o/wget -O), not reads. + write_flags = OUTPUT_FLAG_COMMANDS.get(cmd_name.rsplit('/', 1)[-1], frozenset()) + + sub_args = args[1:] + i = 0 + while i < len(sub_args): + arg = sub_args[i] + if write_flags: + dest = None + if arg in write_flags and i + 1 < len(sub_args): # -o FILE / --output FILE + dest, i = sub_args[i + 1], i + 1 + elif '=' in arg and arg.split('=', 1)[0] in write_flags: # --output=FILE + dest = arg.split('=', 1)[1] + else: # -oFILE (short attached form) + for f in write_flags: + if len(f) == 2 and arg.startswith(f) and len(arg) > 2: + dest = arg[2:] + break + if dest and dest != '-': # '-' is stdout, not a file + _add_path(dest, "write") + i += 1 + continue + if not arg.startswith('-') and _is_file_path(arg): _add_path(arg, base_access) + i += 1 return paths @@ -447,27 +680,22 @@ def build_target_choices(target: str) -> List[Dict]: { "label": f"Allow this URL only ({base_path})", "rules": [f"target({base_path}*)"], - "selected": False, }, { "label": f"Allow all URLs from {host_port}", "rules": host_rules, - "selected": False, }, { "label": f"Allow all URLs from {host} (any port)", "rules": host_rules, - "selected": False, }, { "label": "All of the above", "rules": host_rules, - "selected": False, }, { "label": "Deny (block this action)", "rules": [], - "selected": False, }, ] # Deduplicate options 2 and 3 when there's no port @@ -482,27 +710,22 @@ def build_target_choices(target: str) -> List[Dict]: { "label": f"Allow {target} only", "rules": [host_rule], - "selected": False, }, { "label": f"Allow {target} (any port)", "rules": [host_rule, port_rule], - "selected": False, }, { "label": f"Allow all URLs from {target} (any port)", "rules": [host_rule, port_rule, url_rule, f"target((http|https)://{target}/*)"], - "selected": False, }, { "label": "All of the above", "rules": [host_rule, port_rule, url_rule, f"target((http|https)://{target}/*)"], - "selected": False, }, { "label": "Deny (block this action)", "rules": [], - "selected": False, }, ] return choices @@ -518,6 +741,22 @@ class PermissionResult: shell_command: str = "" # full command when prompting for shell approval +def _is_default_deny(result: "PermissionResult") -> bool: + """True if `result` is the catch-all "no rule matched" deny, not an explicit deny rule.""" + return "No rule for" in result.reason + + +# Finding types downstream auto-trusts. tasks/ai.py _auto_approve_workspace_targets() +# searches _type:"target" findings and auto-approves them as in-scope, so an injected +# add_finding of one of these silently widens scope. +_PRIVILEGED_FINDING_TYPES = frozenset({"target"}) + + +def _is_privileged_finding_type(action: Dict) -> bool: + """True if an add_finding action would mint a downstream-trusted (scope-widening) finding.""" + return str(action.get("_type", "")).strip().lower() in _PRIVILEGED_FINDING_TYPES + + class PermissionEngine: """Evaluate AI actions against allow/deny/ask permission rules. @@ -525,18 +764,58 @@ class PermissionEngine: Two-step validation: (1) action type check, (2) target/path check. """ - def __init__(self, config: Dict, targets: List[str] = None, workspace: str = ""): + def __init__( + self, config: Dict, targets: List[str] = None, workspace: str = "", + allowed_targets: List[str] = None, denied_targets: List[str] = None + ): self.targets = targets or [] self.workspace = str(workspace) self.rules = {"allow": [], "deny": [], "ask": []} self.runtime_allow: List[Tuple[str, List[str]]] = [] + # Platform-supplied allow-list of target regexes (e.g. validated workspace mandates): + # constrains the AI to this scope. Regex full-match, falls back to literal match. + self.allowed_targets: List = self._compile_patterns(allowed_targets) + + # Platform-supplied deny-list of target regexes (mandate `deny` scope). Symmetric + # to allowed_targets but DENY WINS, mirroring the mandate scope matcher. + self.denied_targets: List = self._compile_patterns(denied_targets) + for category in ("allow", "deny", "ask"): for rule_str in config.get(category, []): resolved = self._resolve_variables(rule_str) rule_type, patterns = parse_rule(resolved) self.rules[category].append((rule_type, patterns)) + @staticmethod + def _compile_patterns(patterns: List[str]) -> List: + """Compile a list of regex patterns, falling back to a literal-escaped match on error.""" + compiled: List = [] + for pat in (patterns or []): + if not pat: + continue + try: + compiled.append(re.compile(pat)) + except re.error: + compiled.append(re.compile(re.escape(pat))) + return compiled + + @staticmethod + def _matches_any(patterns: List, value: str) -> bool: + """Check if a value matches any of the given compiled regexes (full or partial match).""" + for rx in patterns: + if rx.fullmatch(value) or rx.match(value): + return True + return False + + def _matches_allowed_targets(self, value: str) -> bool: + """Check if a target value matches any platform-supplied allowed_targets regex.""" + return self._matches_any(self.allowed_targets, value) + + def _matches_denied_targets(self, value: str) -> bool: + """Check if a target value matches any platform-supplied denied_targets regex.""" + return self._matches_any(self.denied_targets, value) + def _resolve_variables(self, rule: str) -> str: """Replace {workspace} and {targets} variables in a rule string.""" result = rule.replace("{workspace}", self.workspace) @@ -560,9 +839,10 @@ def check_action(self, action: Dict) -> PermissionResult: if result.decision in ("deny", "ask"): return result - # Step 2: Check targets (only if target rules are configured) + # Step 2: Check targets. Always enforce when targets exist — a missing + # catch-all must fall to ask (via _check_values "No rule"), never default-allow. targets_to_check = self._extract_targets(action) - if targets_to_check and self._has_rules_for("target"): + if targets_to_check: target_result = self._check_values("target", targets_to_check) if target_result.decision == "deny": return target_result @@ -577,7 +857,7 @@ def check_action(self, action: Dict) -> PermissionResult: if action_type == "shell": command = action.get("command", "") paths_with_access = detect_paths_with_access(command) - if paths_with_access and (self._has_rules_for("read") or self._has_rules_for("write")): + if paths_with_access: # Always enforce — no read/write rule must ask, not allow # Check each path with its correct access type ask_paths = [] for path, access in paths_with_access: @@ -585,7 +865,7 @@ def check_action(self, action: Dict) -> PermissionResult: if path_result.decision == "deny": # Explicit deny rule: block immediately # "No rule" default deny: prompt user instead - if "No rule for" in path_result.reason: + if _is_default_deny(path_result): ask_paths.append((path, access)) else: return PermissionResult( @@ -621,6 +901,9 @@ def check_action(self, action: Dict) -> PermissionResult: def _has_rules_for(self, rule_type: str) -> bool: """Check if any rules exist for the given rule type.""" + # allowed_targets/denied_targets force the target-check step so out-of-scope/denied targets are caught. + if rule_type == "target" and (self.allowed_targets or self.denied_targets): + return True for category in ("allow", "deny", "ask"): for rt, _ in self.rules[category]: if rt == rule_type: @@ -639,8 +922,8 @@ def _check_action_type(self, action_type: str, action: Dict) -> PermissionResult command = action.get("command", "") if not command.strip(): return PermissionResult(decision="deny", reason="Empty command") - cmd_names = _extract_cmd_names(command) - if not cmd_names: + subcommands = _parse_subcommands(command) + if not subcommands: # Parse failure — prompt user for the whole command return PermissionResult( decision="ask", @@ -649,11 +932,20 @@ def _check_action_type(self, action_type: str, action: Dict) -> PermissionResult ) most_restrictive = None unmatched = [] - for cmd_name in cmd_names: + for args in subcommands: + # peel exec-wrappers so the INNER command is checked, not the wrapper name + inner = _peel_wrapper(args) + if not inner: + continue + cmd_name = inner[0] + # multi-word denies (e.g. "rm -rf /*") match the full peeled command; names via _check_value + denied = self._match_shell_command_deny(inner) + if denied: + return PermissionResult(decision="deny", reason=f"Denied by rule: shell({denied})") result = self._check_value("shell", cmd_name) if result.decision == "deny": # Distinguish explicit deny rules from "no matching rule" default - if "No rule for" in result.reason: + if _is_default_deny(result): unmatched.append(cmd_name) else: return result # Explicit deny rule hit @@ -676,9 +968,29 @@ def _check_action_type(self, action_type: str, action: Dict) -> PermissionResult name = action.get("name", "") return self._check_value(action_type, name) elif action_type in ("query", "follow_up", "add_finding"): + # Don't let injected add_finding mint a trusted target that auto-approve later trusts + if action_type == "add_finding" and _is_privileged_finding_type(action): + ftype = str(action.get("_type", "")).strip().lower() + return PermissionResult( + decision="ask", + reason=f"add_finding of privileged type '{ftype}' requires approval", + ) return PermissionResult(decision="allow", reason=f"{action_type} is always allowed") return PermissionResult(decision="deny", reason=f"Unknown action type: {action_type}") + def _match_shell_command_deny(self, tokens: List[str]) -> str: + """Return a multi-word shell deny pattern (e.g. "rm -rf /*") hit by these tokens, else "".""" + cmd_str = ' '.join(tokens) + for rt, patterns in self.rules["deny"]: + if rt != "shell": + continue + for pattern in patterns: + if ' ' not in pattern: + continue # single-token denies are handled by name in _check_value + if _match_command_glob(cmd_str, pattern): + return pattern + return "" + def _check_value(self, rule_type: str, value: str) -> PermissionResult: """Check a single value. Order: deny > allow > ask > deny. @@ -701,6 +1013,21 @@ def _check_value(self, rule_type: str, value: str) -> PermissionResult: if match_rule(v, patterns): return PermissionResult(decision="deny", reason=f"Denied by rule: {rule_type}({v})") + # Platform-supplied denied_targets (regex) deny-list — checked before the + # allowed_targets allow-list so DENY WINS: a target matching both an allow + # and a deny mandate scope is denied (mirrors the mandate scope matcher). + if rule_type == "target" and self.denied_targets: + for v in values_to_check: + if self._matches_denied_targets(v): + return PermissionResult(decision="deny", reason=f"Denied by mandate: target({v})") + + # Platform-supplied allowed_targets (regex) allow-list — checked after deny + # (deny still wins) but before config/runtime allow rules. + if rule_type == "target" and self.allowed_targets: + for v in values_to_check: + if self._matches_allowed_targets(v): + return PermissionResult(decision="allow", reason=f"Allowed by mandate: target({v})") + for rt, patterns in self.rules["allow"]: if rt == rule_type: for v in values_to_check: @@ -727,7 +1054,7 @@ def _check_values(self, rule_type: str, values: List[str]) -> PermissionResult: result = self._check_value(rule_type, value) if result.decision == "deny": # "No rule for" default deny → ask user instead of blocking - if "No rule for" in result.reason: + if _is_default_deny(result): ask_targets.append(value) else: return result # Explicit deny rule: block @@ -771,11 +1098,8 @@ def prompt_target(self, target: str, interactive: bool = True, command: str = "" Returns: 'allow' or 'deny' """ - if not interactive: - return "deny" - choices = build_target_choices(target) - selected_indices = self._show_target_menu(target, choices, command=command) + selected_indices = self._show_target_menu(target, choices, command=command, interactive=interactive) if selected_indices is None: return "deny" @@ -807,11 +1131,6 @@ def prompt_path(self, path: str, access_type: str = "read", interactive: bool = Returns: 'allow' or 'deny' """ - if not interactive: - return "deny" - - from secator.rich import InteractiveMenu - action_label = "Read from" if access_type == "read" else "Write to" parent = '/'.join(path.split('/')[:-1]) if '/' in path else path options = [ @@ -819,16 +1138,16 @@ def prompt_path(self, path: str, access_type: str = "read", interactive: bool = {"label": f"Allow {access_type}({parent}/*)"}, {"label": "Deny (block this action)"}, ] - result = InteractiveMenu( + idx = self._show_menu( f"{action_label} {path} requires approval.", options, description=command, - ).show() + interactive=interactive, + ) - if result is None: + if idx is None: return "deny" - idx, _ = result if idx == 2: # Deny return "deny" elif idx == 0: # Exact path @@ -848,11 +1167,6 @@ def prompt_shell(self, command: str, reason: str = "", interactive: bool = True) Returns: 'allow' or 'deny' """ - if not interactive: - return "deny" - - from secator.rich import InteractiveMenu - # Extract command names; use the unmatched one(s) from reason for option 2 cmd_names = _extract_cmd_names(command) # Parse unmatched commands from reason like "No rule for command(s): ./terrapin-scanner, foo" @@ -868,48 +1182,71 @@ def prompt_shell(self, command: str, reason: str = "", interactive: bool = True) {"label": "Deny (block this action)"}, ] title = reason or "Shell command requires approval" - result = InteractiveMenu( + idx = self._show_menu( title, options, description=f"[gray42]{command}[/gray42]", - ).show() + interactive=interactive, + ) - if result is None: + if idx is None: return "deny" - idx, _ = result - if idx == 0: # Allow this specific command (one-time, no rule added) - # Add a runtime allow for each cmd name in this command - if cmd_names: - self.add_runtime_allow([f"shell({','.join(cmd_names)})"]) + if idx == 0: # Allow ONLY this invocation — no rule added, next call re-prompts return "allow" elif idx == 1: # Allow all commands with this name self.add_runtime_allow([f"shell({prompt_cmd})"]) return "allow" return "deny" - def _show_target_menu(self, target: str, choices: List[Dict], command: str = "") -> List[int]: + def _show_target_menu( + self, target: str, choices: List[Dict], command: str = "", interactive: bool = True + ) -> Optional[List[int]]: """Show interactive menu. Separated for testability. Args: target: The target being prompted about choices: List of choice dicts from build_target_choices command: The shell command triggering this prompt (for display) + interactive: If False, auto-deny without prompting Returns: List of selected indices, or None if cancelled """ - from secator.rich import InteractiveMenu - options = [{"label": choice["label"]} for choice in choices] - result = InteractiveMenu( + idx = self._show_menu( f"Target {target} is not in allowed targets. Add it?", options, description=command, - ).show() + interactive=interactive, + ) + + if idx is None: + return None + return [idx] + def _show_menu( + self, title: str, options: List[Dict], description: str = "", interactive: bool = True + ) -> Optional[int]: + """Shared interactive-menu scaffold used by prompt_path/prompt_shell/_show_target_menu. + + Args: + title: Menu title/prompt text + options: List of {"label": ...} option dicts + description: Extra context shown below the title (e.g. the shell command) + interactive: If False, auto-deny (return None) without prompting + + Returns: + The selected index, or None if not interactive or the user cancelled. + """ + if not interactive: + return None + + from secator.rich import InteractiveMenu + + result = InteractiveMenu(title, options, description=description).show() if result is None: return None idx, _ = result - return [idx] + return idx diff --git a/secator/ai/history.py b/secator/ai/history.py index 1d3efe484..36783c8ca 100644 --- a/secator/ai/history.py +++ b/secator/ai/history.py @@ -2,7 +2,6 @@ """Chat history management for AI task - litellm format.""" import json from dataclasses import dataclass, field -from datetime import datetime from pathlib import Path from typing import Dict, List, Optional, Tuple @@ -13,6 +12,32 @@ COMPACTION_THRESHOLD_PCT = 85 # Trigger compaction at 85% of usable context MAX_ACTION_TOKENS = 10_000 # Hard cap per action result +# Hard cap on a persisted transcript message (BSON-safety backstop, well under +# Mongo's 16MB doc limit) -- not primary truncation, which happens upstream via truncate_to_tokens. +MAX_PERSISTED_MESSAGE_CHARS = 12000 + + +def _cap(text, max_chars): + if isinstance(text, str) and len(text) > max_chars: + return text[:max_chars] + '…[capped]' + return text + + +def cap_message(msg: dict, max_chars: int = MAX_PERSISTED_MESSAGE_CHARS) -> dict: + """Return a copy of a litellm message with content + tool-call arguments + capped to max_chars (BSON-safety backstop). Non-string/short fields untouched.""" + out = dict(msg) + if 'content' in out: + out['content'] = _cap(out['content'], max_chars) + if out.get('tool_calls'): + out['tool_calls'] = [ + {**tc, 'function': {**tc.get('function', {}), + 'arguments': _cap(tc.get('function', {}).get('arguments'), max_chars)}} + if tc.get('function') else tc + for tc in out['tool_calls'] + ] + return out + def get_context_window(model: str) -> int: """Get model's context window size from litellm. @@ -47,8 +72,6 @@ def truncate_to_tokens( max_tokens: int, model: str, fallback_path: Path = None, - output_dir: Path = None, - result_name: str = "result" ) -> str: """Truncate content to fit within token budget, with file fallback. @@ -57,8 +80,6 @@ def truncate_to_tokens( max_tokens: Maximum tokens allowed model: LLM model name for token counting fallback_path: Existing file to reference (task/workflow report.json) - output_dir: Directory to save shell output (creates file) - result_name: Prefix for saved filename Returns: Original content if under budget, or truncated with [TRUNCATED] marker @@ -75,13 +96,6 @@ def truncate_to_tokens( if fallback_path and fallback_path.exists(): file_hint = f"\nFull output: {fallback_path}" debug(f'using existing fallback: {fallback_path}', sub='runner.ai.context') - elif output_dir: - output_dir.mkdir(parents=True, exist_ok=True) - timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") - fallback_path = output_dir / f"{result_name}_{timestamp}.txt" - fallback_path.write_text(content) - file_hint = f"\nFull output saved to: {fallback_path}" - debug(f'saved output to: {fallback_path}', sub='runner.ai.context') else: file_hint = "" @@ -93,6 +107,11 @@ def truncate_to_tokens( return content[:truncate_at] + f"\n\n[TRUNCATED]{file_hint}" +def _usable_tokens(model: str) -> int: + """Model's context window minus the reserved output allowance.""" + return get_context_window(model) - OUTPUT_TOKEN_RESERVATION + + SUMMARIZATION_PROMPT = """Summarize the following attack session history into a compact context. Keep ONLY the essential information: - Key findings (vulnerabilities, open ports, services, credentials) @@ -124,9 +143,13 @@ class ChatHistory: messages: List[Dict[str, str]] = field(default_factory=list) model: Optional[str] = None - - def add_system(self, content: str) -> None: - self.messages.append({"role": "system", "content": content}) + # Billed token/cost usage accrued by LLM calls this object makes internally + # (history summarization/compaction). The owning `ai` task drains these into + # context.ai_tokens so summarization is billed alongside the main loop. + billed_tokens: int = 0 + billed_prompt_tokens: int = 0 + billed_completion_tokens: int = 0 + billed_cost: float = 0.0 def set_system(self, content: str) -> None: """Replace the first system message, or insert one at the start. @@ -172,22 +195,36 @@ def add_tool_result(self, name: str, tool_call_id: str, content: str) -> None: msg["name"] = name self.messages.append(msg) - def add_tool(self, content: str) -> None: - self.messages.append({"role": "tool", "content": content}) - def to_messages(self, max_tokens_total: int = 0) -> List[Dict[str, str]]: - """Return a copy of the messages list, trimming if over max_tokens_total. + """Return a copy of the messages list, trimming if over the effective budget. Uses litellm's trim_messages which preserves system messages and recent context while removing oldest messages first. Args: - max_tokens_total: Hard token limit. If > 0, trim messages to fit. + max_tokens_total: Requested hard token limit (0 = no explicit cap). """ - if max_tokens_total > 0: - return self.trim(max_tokens_total) + budget = self._trim_budget(max_tokens_total) + if budget > 0: + return self.trim(budget) return self.messages.copy() + def _trim_budget(self, max_tokens_total: int = 0) -> int: + """Effective trim budget, capped to the model's real context window. + + A flat max_tokens_total (e.g. 100k) ignores the model window and + fails with context_length_exceeded on smaller-window models. Cap it to + get_context_window(model) - OUTPUT_TOKEN_RESERVATION (headroom for the + response), and use that window-derived budget even when no explicit cap + is set. With no model known, keep the legacy caller-driven behavior. + """ + if not self.model: + return max_tokens_total + window_budget = max(_usable_tokens(self.model), 1) + if max_tokens_total > 0: + return min(max_tokens_total, window_budget) + return window_budget + def trim(self, max_tokens: int) -> List[Dict[str, str]]: """Trim messages to fit under max_tokens using litellm's trim_messages. @@ -201,11 +238,27 @@ def trim(self, max_tokens: int) -> List[Dict[str, str]]: Trimmed list of messages. """ from litellm.utils import trim_messages + from secator.ai.utils import _strip_leading_orphan_tools from secator.rich import console from secator.output_types import Warning original_count = len(self.messages) - trimmed = trim_messages(self.messages, max_tokens=max_tokens) + # litellm's trim_messages does len(msg["content"]), which raises TypeError when an + # assistant turn carries only tool_calls (content=None/absent) -- coerce to "" first. + # Wrap the call so a trimmer bug degrades to untrimmed history instead of crashing. + sanitized = [dict(m, content="") if m.get("content") is None else m for m in self.messages] + try: + trimmed = trim_messages(sanitized, max_tokens=max_tokens) + except Exception as e: # noqa: BLE001 - a token-trimming crash must never kill the AI loop + console.print(Warning( + message=f'Chat history trim failed ({type(e).__name__}: {e}); using untrimmed history.' + )) + trimmed = sanitized + + # litellm drops the OLDEST messages with no tool-pairing awareness, so the + # kept window can START with an orphan tool_result whose assistant(tool_calls) + # parent was dropped — Anthropic/OpenAI reject that. Drop leading orphans. + _strip_leading_orphan_tools(trimmed) dropped = original_count - len(trimmed) if dropped: @@ -285,7 +338,7 @@ def get_available_tokens(self, model: str) -> int: Available tokens (context - reservation - used) """ context_window = get_context_window(model) - usable = context_window - OUTPUT_TOKEN_RESERVATION + usable = _usable_tokens(model) used = self.count_tokens(model) available = usable - used debug( @@ -294,20 +347,18 @@ def get_available_tokens(self, model: str) -> int: ) return available - def should_compact(self, model: str, threshold_pct: int = COMPACTION_THRESHOLD_PCT) -> bool: + def should_compact(self, model: str) -> bool: """Check if compaction needed based on % of context used. Args: model: LLM model name - threshold_pct: Percentage threshold (default 85) Returns: True if compaction needed """ - context_window = get_context_window(model) - usable = context_window - OUTPUT_TOKEN_RESERVATION + usable = _usable_tokens(model) used = self.count_tokens(model) - threshold = usable * threshold_pct / 100 + threshold = usable * COMPACTION_THRESHOLD_PCT / 100 should = used > threshold pct_used = (used / usable * 100) if usable > 0 else 0 debug( @@ -343,7 +394,7 @@ def maybe_summarize(self, model: str, api_base: Optional[str] = None, return True, old_tokens, new_tokens def compact(self, model: str, api_base: Optional[str] = None, - api_key: Optional[str] = None, keep_last: int = 4) -> None: + api_key: Optional[str] = None) -> None: """Summarize non-system messages using an LLM, keeping the initial system prompt and the last few messages intact so the LLM retains recent context. @@ -351,8 +402,8 @@ def compact(self, model: str, api_base: Optional[str] = None, model: LLM model name api_base: Optional API base URL api_key: Optional API key - keep_last: Number of recent non-system messages to preserve (default 4) """ + keep_last = 4 if len(self.messages) <= 2: return @@ -374,13 +425,17 @@ def compact(self, model: str, api_base: Optional[str] = None, to_summarize = rest to_keep = [] - from secator.ai.utils import call_llm + from secator.ai.utils import call_llm, _strip_leading_orphan_tools from secator.rich import console from secator.utils import format_token_count + # The blind keep_last tail cut can leave to_keep STARTING with a tool_result + # whose assistant(tool_calls) parent fell into to_summarize — strip those so + # the rebuilt window never begins on an orphan tool_result. + _strip_leading_orphan_tools(to_keep) + # Calculate target summary size based on available context - context_window = get_context_window(model) - usable = context_window - OUTPUT_TOKEN_RESERVATION + usable = _usable_tokens(model) target_tokens = int(usable * 0.3) # Target 30% of usable context max_words = target_tokens // 2 # Rough tokens-to-words ratio @@ -390,6 +445,20 @@ def compact(self, model: str, api_base: Optional[str] = None, with console.status(f"[bold orange3]Compacting chat history...[/] [gray42] • {token_str}[/]", spinner="dots"): result = call_llm([{"role": "user", "content": prompt}], model, 0.3, api_base, api_key) + # Record billed usage of the summarization call so the owning task can + # roll it into context.ai_tokens. Missing usage counts as 0. + usage = result.get("usage") or {} + for attr, key, cast in ( + ("billed_tokens", "tokens", int), + ("billed_prompt_tokens", "prompt_tokens", int), + ("billed_completion_tokens", "completion_tokens", int), + ("billed_cost", "cost", float), + ): + try: + setattr(self, attr, getattr(self, attr) + cast(usage.get(key) or 0)) + except (TypeError, ValueError): + pass + self.messages = [] if initial_system: self.messages.append(initial_system) diff --git a/secator/ai/interactivity.py b/secator/ai/interactivity.py index 7744ae2e1..32a84e2ed 100644 --- a/secator/ai/interactivity.py +++ b/secator/ai/interactivity.py @@ -33,8 +33,12 @@ def ask_user(self, question: str, choices: List[str], session_id: str, raise NotImplementedError def get_excluded_tools(self) -> set: - """Return tool names to exclude from the LLM's available tools.""" - return set() + """Return tool names to exclude from the LLM's available tools. + + Excludes "stop" by default: only AutoBackend (no user to hand control back + to) needs the LLM to have an explicit stop tool. + """ + return {"stop"} def get_extra_tools(self) -> list: """Return additional tool schemas (not in TOOL_SCHEMAS) to inject.""" @@ -44,9 +48,6 @@ def get_extra_tools(self) -> list: class CLIBackend(InteractivityBackend): """Local terminal interactive backend.""" - def get_excluded_tools(self) -> set: - return {"stop"} - def ask_user(self, question, choices, session_id, prompt_type="follow_up", **context): if prompt_type == "permission": return self._handle_permission(**context) @@ -79,7 +80,6 @@ def _handle_follow_up(self, choices, **context): return None return prompt_user( history, - encryptor=context.get("encryptor"), max_iterations=context.get("max_iterations", 10), choices=choices, mode=context.get("mode", "chat"), @@ -95,67 +95,199 @@ def __init__(self, timeout: int = 600, query_engine: Any = None, poll_interval: self.query_engine = query_engine self.poll_interval = poll_interval - def get_excluded_tools(self) -> set: - return {"stop"} - def build_pending_prompt(self, question, choices, session_id, prompt_type="follow_up", **context): """Build a pending Ai finding for the remote user to see and answer. The caller must yield this item so it gets stored in the workspace (via runner hooks) before calling ask_user(), which will poll for the answer. + + ``prompt_uuid`` (from context) is stamped into ``extra_data`` so the poll + can match THIS exact prompt, not a stale earlier answer. """ from secator.output_types import Ai + extra_data = { + "permission_type": context.get("permission_type", ""), + "value": context.get("value", ""), + } + prompt_uuid = context.get("prompt_uuid") + if prompt_uuid: + extra_data["prompt_uuid"] = prompt_uuid + # A new prompt for this session supersedes any older still-pending one + # (e.g. a worker that died mid-poll). Expire them BEFORE this doc is + # persisted so only the current prompt stays live. + self._expire_stale_pending(session_id) + # The conversation id rides on `_context.session_id` (auto-stamped from the + # runner context on persist) — the poll + restore + secator-api all key on + # that, so this pending doc needs no top-level session_id field. return Ai( content=question, ai_type=prompt_type, status="pending", choices=choices, - session_id=session_id, - extra_data={ - "permission_type": context.get("permission_type", ""), - "value": context.get("value", ""), - }, + extra_data=extra_data, _timestamp=time.time(), ) def ask_user(self, question, choices, session_id, prompt_type="follow_up", **context): - answer = self._poll_for_answer(session_id, prompt_type) + answer = self._poll_for_answer(session_id, prompt_type, prompt_uuid=context.get("prompt_uuid")) if answer is None: return None if prompt_type == "permission": engine = context.get("engine") - if answer in ("allow", "allow_all") and engine: - ptype = context.get("permission_type") - value = context.get("value", "") - self._add_permission_rules(engine, ptype, value) + if answer in ("allow", "allow_all"): + # allow_all persists a session-scoped allow rule; single allow is + # a true one-shot that adds NO rule — next match re-prompts. + if answer == "allow_all" and engine: + ptype = context.get("permission_type") + value = context.get("value", "") + self._add_permission_rules(engine, ptype, value) return {"answer": "allow"} return {"answer": "deny"} # follow_up: return the answer text return {"answer": answer} - def _poll_for_answer(self, session_id, prompt_type): - """Poll DB for user answer until timeout.""" + def poll_steers(self, session_id): + """Drain pending steer docs for ``session_id`` and mark them consumed. + + A "steer" is a mid-flight user message: it's written into the channel + (``_type:"ai"``, ``ai_type:"steer"``, ``status:"pending"``) WHILE the agent + is running, and the worker picks it up at the next loop checkpoint to + redirect the next turn. This is distinct from a follow-up ``answer`` (which + the worker is *blocked* waiting on) and from a hard Stop (which revokes the + Celery task). + + Returns a list of steer content strings (oldest-first). Each returned doc is + flipped to ``status:"consumed"`` so it's injected exactly once. Robust by + design: any backend error returns ``[]`` so a steer can never crash the run. + """ + if self.query_engine is None: + return [] + base = { + "_type": "ai", + "ai_type": "steer", + # Correlate by the runner context's session_id, auto-stamped on every + # persisted item (item._context = self.context) — see _poll_for_answer. + "_context.session_id": session_id, + "status": "pending", + } + try: + results = self.query_engine.search(base, limit=50) + except Exception: # noqa: BLE001 - a steer must never crash the run + return [] + if not results: + return [] + # Oldest-first so multiple queued steers are injected in send order. + results = sorted(results, key=lambda r: r.get("_timestamp", 0)) + contents = [] + for doc in results: + content = doc.get("content") or doc.get("answer") or "" + if content: + contents.append(content) + # Mark this session's pending steers consumed so they inject exactly once. + try: + self.query_engine.update( + {**base}, + {"$set": {"status": "consumed"}}, + ) + except Exception: # noqa: BLE001 - consume failure must not crash the run + pass + return contents + + def _poll_for_answer(self, session_id, prompt_type, prompt_uuid=None): + """Poll the DB for the answer to THIS specific prompt until timeout. + + The query MUST be scoped to the exact prompt the worker is currently + blocked on — identified by ``prompt_uuid`` (stamped into the pending doc's + ``extra_data.prompt_uuid`` before it was persisted). Matching only on + ``{session_id, status:"answered"}`` is a bug: a multi-turn conversation + accumulates *previously* answered follow-up docs, so an unscoped query + returns a STALE answer immediately, the worker re-injects that old answer + as a brand-new prompt, re-runs the whole turn, asks again, re-matches the + same stale doc — an infinite respawn loop that re-runs scans and burns + tokens. Scoping on ``prompt_uuid`` makes the poll resolve only THIS + prompt's own answer (and time out only THIS prompt's doc). + + A steer (mid-flight user message) breaks the wait: if a pending steer + arrives for this session while we're blocked on a follow-up, we return its + content as the "answer" so the loop redirects immediately instead of + stalling until the follow-up is explicitly answered (or times out). This + keeps follow-up semantics intact for the no-steer case. + """ + base = { + "_type": "ai", + "ai_type": prompt_type, + # session_id is auto-stamped on every persisted item (item._context) + "_context.session_id": session_id, + } + if prompt_uuid: + base["extra_data.prompt_uuid"] = prompt_uuid + + answered_query = {**base, "status": "answered"} elapsed = 0 while elapsed < self.timeout: - results = self.query_engine.search({ - "_type": "ai", - "ai_type": prompt_type, - "session_id": session_id, - "status": "answered" - }, limit=1) - if results: - return results[0].get("answer") + answer = self._resolve_answer(answered_query) + if answer is not None: + return answer + # A steer breaks the wait: treat the steer as the user's answer so the + # blocked follow-up resolves and the next turn redirects. + steers = self.poll_steers(session_id) + if steers: + return "\n".join(steers) sleep(self.poll_interval) elapsed += self.poll_interval - # Timeout: update finding status - self.query_engine.update( - {"_type": "ai", "ai_type": prompt_type, "session_id": session_id, "status": "pending"}, + # One final search before giving up: the user may have answered during + # the last sleep (or between the last search and now). Without this the + # answer is silently stranded. + answer = self._resolve_answer(answered_query) + if answer is not None: + return answer + # Timeout: atomically flip ONLY a doc that is STILL pending, so a concurrent + # older pending doc isn't disturbed. If the answer landed in the race window + # the doc is already 'answered' and this no-ops -- re-read rather than abandon it. + modified = self.query_engine.update( + {**base, "status": "pending"}, {"$set": {"status": "timed_out"}} ) + if not modified: + answer = self._resolve_answer(answered_query) + if answer is not None: + return answer return None + def _resolve_answer(self, answered_query): + """Return the newest answered doc's answer, or None if none answered. + + Resolving against the newest by ``_timestamp`` is a backstop against + stale answers. + """ + results = self.query_engine.search(answered_query) + if not results: + return None + newest = max(results, key=lambda r: r.get("_timestamp", 0)) + return newest.get("answer") + + def _expire_stale_pending(self, session_id): + """Mark any older still-pending prompt for this session as timed_out. + + Called when a NEW prompt starts (before it is persisted), so it only + affects prior prompts. Stops stale 'pending' docs from accumulating — + a worker that dies mid-poll otherwise leaves the UI 'thinking' forever + and lets crud.answer_ai_prompt's "latest pending" collide. + FLAG: a DB-layer TTL index on pending Ai docs is the durable follow-up. + """ + if not self.query_engine: + return + self.query_engine.update( + { + "_type": "ai", + "_context.session_id": session_id, + "status": "pending", + }, + {"$set": {"status": "timed_out"}}, + ) + @staticmethod def _add_permission_rules(engine, ptype, value): """Add runtime allow rules after a remote permission approval.""" diff --git a/secator/ai/prompts.py b/secator/ai/prompts.py index 5e51595ac..8d7812997 100644 --- a/secator/ai/prompts.py +++ b/secator/ai/prompts.py @@ -241,13 +241,14 @@ def get_system_prompt(mode: str, workspace_path: str = "", backend=None) -> str: system_prompt = mode_config["system_prompt"] ws = workspace_path or "" - path_vars = dict(tasks_path=str(TASKS_PATH), workflows_path=str(WORKFLOWS_PATH), profiles_path=str(PROFILES_PATH)) - if mode == "attack": - result = system_prompt.safe_substitute(library_reference=build_library_reference(), **path_vars) - elif mode == "exploit": - result = system_prompt.safe_substitute(library_reference=build_library_reference(), **path_vars) - else: # chat mode - result = system_prompt.safe_substitute(output_types_reference=build_output_types_reference()) + # The queries.txt constraint (included by every mode) references $query_types and + # $output_types_reference, so they must be substituted for all modes — derive both + # from FINDING_TYPES so they never drift from the registry. + subst = dict(query_types=build_query_types(), output_types_reference=build_output_types_reference()) + if mode in ("attack", "exploit"): + path_vars = dict(tasks_path=str(TASKS_PATH), workflows_path=str(WORKFLOWS_PATH), profiles_path=str(PROFILES_PATH)) + subst.update(library_reference=build_library_reference(), **path_vars) + result = system_prompt.safe_substitute(**subst) # Determine interaction rules based on backend # The mode templates already include ${follow_up} for interactive modes. @@ -260,27 +261,6 @@ def get_system_prompt(mode: str, workspace_path: str = "", backend=None) -> str: return result.replace("$workspace_path", ws) -# def format_user_initial(targets: List[str], instructions: str, previous_results: List[Dict] = None) -> str: -# """Format initial user message as compact JSON. - -# Args: -# targets: List of target hosts/URLs -# instructions: User instructions for the task -# previous_results: Optional list of result dicts from upstream tasks - -# Returns: -# Compact JSON string (no whitespace) -# """ -# results_str = json.dumps(previous_results, default=str) -# instructions_str = json.dumps(instructions or "Analyze the previous results first") -# return f""" -# -# {instructions_str} -# {results_str} -# -# """ - - def format_tool_result(name: str, status: str, count: int, results: Any, max_items: int = 100) -> str: """Format tool result as compact JSON, truncating results if too many. diff --git a/secator/ai/prompts/constraints/common.txt b/secator/ai/prompts/constraints/common.txt index 3fce07169..fd0a6de6e 100644 --- a/secator/ai/prompts/constraints/common.txt +++ b/secator/ai/prompts/constraints/common.txt @@ -12,7 +12,7 @@ run_task(name="httpx", targets=["target3.com"], opts={"rate_limit": 30, "proxy": run_workflow(name="domain_recon", targets=["example.com"]) run_shell(command="curl -sk https://10.0.0.1/ | head -50") run_task(name="ai", targets=["example.com"], opts={"prompt": "Enumerate subdomains", "mode": "attack", "session_name": "Subdomain enumeration on example.com", "max_iterations": 5}) -run_query(query={'vulnerability': {'severity': {'$in': ['high', 'critical']}) +query_workspace(query={"_type": "vulnerability", "severity": {"$in": ["high", "critical"]}}) add_finding(name="XSS vuln", matched_at=["http://testphp.vulnweb.com/hpp/?pp=1"], ) @@ -73,8 +73,8 @@ When getting denied to run a command many times, you can also try it to run it i -The runner folder is: $workspace_path. This is where we store all inputs / outputs from the current run. -You can request to read or write files outside the workspace but this require user approval so when possible prefer to read / write to the runner folder. +Write any files you generate (e.g. a markdown report) to the runner folder's outputs directory: $workspace_path/.outputs/. Prefer this location; reading or writing files elsewhere requires user approval. +To find existing findings/results, ALWAYS use the query_workspace tool — it is the single source of truth for the workspace and already covers this run's live findings. Do NOT read local report files (e.g. via cat/jq) to look up findings. diff --git a/secator/ai/session.py b/secator/ai/session.py index 3023b533a..85fedc3d3 100644 --- a/secator/ai/session.py +++ b/secator/ai/session.py @@ -53,13 +53,26 @@ def list_sessions(max_sessions=20): ai_items = data.get('results', {}).get('ai', []) if not ai_items: continue - # Find first user prompt content and session name + # Find first user prompt content + session name, and the first non-empty + # `_context.session_id` across ALL ai docs (every persisted item stamps it, + # letting a resumed run adopt this session's id) -- single pass, stopping + # once both have been found. first_prompt = '' session_name = '' + session_id = '' + found_prompt = False + found_session_id = False for item in ai_items: - if item.get('ai_type') == 'prompt': + if not found_prompt and item.get('ai_type') == 'prompt': first_prompt = item.get('content', '') session_name = (item.get('_context') or {}).get('session_name', '') or (item.get('_context') or {}).get('name', '') + found_prompt = True + if not found_session_id: + sid = (item.get('_context') or {}).get('session_id', '') + if sid: + session_id = sid + found_session_id = True + if found_prompt and found_session_id: break info = data.get('info', {}) sessions.append({ @@ -68,6 +81,7 @@ def list_sessions(max_sessions=20): 'report_path': str(report_path), 'name': session_name, 'prompt': first_prompt, + 'session_id': session_id, 'targets': info.get('targets', []), 'timestamp': info.get('end_time') or info.get('start_time') or 0, 'mtime': history_path.stat().st_mtime, @@ -125,6 +139,42 @@ def show_session_picker(): return sessions[idx] +def print_session_results(session): + """Print a prior session's persisted results (findings + ai turns) to the + console in ``_timestamp`` order — the visible "here's where you left off" + replay shown on resume. Reads the session's ``report.json``; best-effort + (never raises), so a resume is never blocked by a display error. + + Args: + session: Session dict from show_session_picker (uses ``report_path``). + """ + from secator.output_types import OUTPUT_TYPES + + report_path = session.get('report_path') + if not report_path: + return + type_map = {cls.__name__.lower(): cls for cls in OUTPUT_TYPES} + try: + with open(report_path) as f: + data = json.load(f) + except (json.JSONDecodeError, OSError): + return + # Flatten all items with their type class, then print in timestamp order + all_items = [] + for type_name, items in data.get('results', {}).items(): + cls = type_map.get(type_name) + if not cls: + continue + for item_data in items: + all_items.append((item_data, cls)) + all_items.sort(key=lambda x: x[0].get('_timestamp', 0)) + for item_data, cls in all_items: + try: + console.print(cls.load(item_data), highlight=False) + except Exception: + continue + + def replay_session(session): """Replay all results from a previous session and restore history. @@ -135,36 +185,9 @@ def replay_session(session): ChatHistory: Restored history, or None on error. """ from secator.ai.history import ChatHistory - from secator.output_types import OUTPUT_TYPES - # Build type map for loading items - type_map = {cls.__name__.lower(): cls for cls in OUTPUT_TYPES} - - # Load and replay all results from report.json, sorted by timestamp - report_path = session.get('report_path') - if report_path: - try: - with open(report_path) as f: - data = json.load(f) - results = data.get('results', {}) - # Flatten all items with their type class - all_items = [] - for type_name, items in results.items(): - cls = type_map.get(type_name) - if not cls: - continue - for item_data in items: - all_items.append((item_data, cls)) - # Sort by _timestamp - all_items.sort(key=lambda x: x[0].get('_timestamp', 0)) - for item_data, cls in all_items: - try: - item = cls.load(item_data) - console.print(item, highlight=False) - except Exception: - continue - except (json.JSONDecodeError, OSError): - pass + # Show the prior conversation + findings on the console + print_session_results(session) # Load history history_path = session['history_path'] @@ -177,3 +200,102 @@ def replay_session(session): except (json.JSONDecodeError, OSError) as e: console.print(Error(message=f'Failed to load history: {e}')) return None + + +def restore_history_from_db(session_id, query_engine, model=None, encryptor=None, system_prompt=None): + """Rebuild an in-memory ChatHistory from the workspace's `_type:"ai"` Mongo docs. + + Headless equivalent of ``replay_session`` for the remote (web) path: a + respawned ``ai`` task on a different worker pod has no local report files, so + the conversation is rebuilt from the channel docs themselves (queried by + ``session_id``, ordered by ``_timestamp``). + + This is a **faithful, valid litellm transcript continuation** for docs + carrying a raw litellm ``message`` dict (persisted by Tasks 2-3 for every + prompt/assistant/tool_result turn, including tool_calls and tool_call_id + pairing): each persisted message is appended verbatim, in ``_timestamp`` + order. Internal loop nudges (the synthetic "continue"/"retry" ``user`` + prompts the run appends to live history but never persists as docs) are not + restored and so are omitted here — the result is therefore NOT literally + byte-identical to the live in-memory history, but it stays a valid transcript + (a clean tool→assistant continuation the model can resume from). Persisted + ``message.content`` is already encrypted (the encryption happens at persist + time, not at read time), so it is NOT re-encrypted here — doing so would + double-encrypt it. + + Docs from before this feature shipped don't carry a ``message`` field at + all (only the human-readable ``content`` used for the channel/report + display). Those fall back to the legacy **text-only** reconstruction: only + ``ai_type="prompt"``/``"response"`` docs become ``user``/``assistant`` + messages (re-encrypted here, since their plaintext ``content`` was never + encrypted at persist time), and intermediate tool-call/tool-result activity + is collapsed away (it was never captured verbatim pre-upgrade). + + Ordering assumption: a single session is either entirely message-carrying + (post-upgrade) or entirely legacy (pre-upgrade) — sessions aren't upgraded + mid-conversation. So it is safe to restore all message-docs first (in + their own timestamp order) and then append any legacy docs (in their own + timestamp order); within a real session only one of the two groups will be + non-empty, so this two-pass split never reorders an actual transcript. + + Args: + session_id: The conversation's session id (UUID generated by the UI). + query_engine: A ``QueryEngine`` (must resolve to the workspace Mongo + backend for the docs to be visible). + model: Optional LLM model name to set on the returned history. + encryptor: Optional ``SensitiveDataEncryptor``, used only for the legacy + text-only fallback (message-docs are already encrypted verbatim). + system_prompt: Optional system prompt to set as the first message. + + Returns: + ChatHistory: The rebuilt history (possibly with only a system prompt if + no prior docs exist). + """ + from secator.ai.history import ChatHistory + from secator.ai.encryption import maybe_encrypt + from secator.ai.utils import _repair_orphan_tool_uses, _strip_leading_orphan_tools + + history = ChatHistory(model=model) + if system_prompt is not None: + history.set_system(maybe_encrypt(system_prompt, encryptor)) + + try: + docs = query_engine.search({'_type': 'ai', '_context.session_id': session_id}) + except Exception as e: # noqa: BLE001 - backend errors must not crash the worker + console.print(Warning(message=f'Failed to restore session from DB: {e}')) + return history + + docs = sorted(docs or [], key=lambda d: d.get('_timestamp', 0)) + legacy = [] # docs without a raw message (pre-upgrade) -> text-only fallback + for doc in docs: + msg = doc.get('message') + if isinstance(msg, dict) and msg.get('role'): + # Byte-exact: the persisted message already holds encrypted content, so + # append verbatim (no re-encryption) — mirrors a fresh run's history. + history.messages.append(dict(msg)) + else: + legacy.append(doc) + + # Legacy docs (no message field): fall back to text-only prompt/response. + for doc in legacy: + ai_type = doc.get('ai_type') + content = doc.get('content', '') + if not content: + continue + if ai_type == 'prompt': + history.add_user(maybe_encrypt(content, encryptor)) + elif ai_type == 'response': + history.add_assistant(maybe_encrypt(content, encryptor)) + elif ai_type == 'steer': + # A mid-flight steer is a real user turn: preserve it as a user message on + # respawn (mirroring the live-loop framing) so the redirect survives a restore. + history.add_user(maybe_encrypt(f'[User interjected]: {content}', encryptor)) + # All other ai_types (action displays, follow_up/permission prompts, + # shell_output, summaries) are channel/UX artifacts, not conversation + # turns — intentionally skipped for a valid litellm transcript. + + # Guard against a partially-persisted turn producing an orphan tool result. + _repair_orphan_tool_uses(history.messages) + _strip_leading_orphan_tools(history.messages) + + return history diff --git a/secator/ai/tools.py b/secator/ai/tools.py index e85394b4c..d3aebb438 100644 --- a/secator/ai/tools.py +++ b/secator/ai/tools.py @@ -1,5 +1,7 @@ """Tool schema definitions for native LLM tool calling.""" +import json + from secator.ai.prompts import get_mode_config # Map tool names to action types used by existing action handlers @@ -13,8 +15,12 @@ "stop": "stop", } -# Reverse mapping: action type -> tool name -ACTION_TOOL_MAP = {v: k for k, v in TOOL_ACTION_MAP.items()} +# Shared "targets" parameter schema (identical across run_task/run_workflow) +_TARGETS_SCHEMA = { + "type": "array", + "items": {"type": "string"}, + "description": "List of targets (hosts, URLs, IPs)." +} # OpenAI-format tool schemas keyed by tool name TOOL_SCHEMAS = { @@ -30,14 +36,10 @@ "type": "string", "description": "The task name (e.g. nmap, httpx, nuclei, ffuf)." }, - "targets": { - "type": "array", - "items": {"type": "string"}, - "description": "List of targets (hosts, URLs, IPs)." - }, + "targets": _TARGETS_SCHEMA, "opts": { "type": "object", - "description": "Optional task-specific options (e.g. ports, rate_limit, timeout)." + "description": "Optional task-specific options (e.g. ports, rate_limit). Control/security flags are ignored." } }, "required": ["name", "targets"] @@ -56,14 +58,10 @@ "type": "string", "description": "The workflow name." }, - "targets": { - "type": "array", - "items": {"type": "string"}, - "description": "List of targets (hosts, URLs, IPs)." - }, + "targets": _TARGETS_SCHEMA, "opts": { "type": "object", - "description": "Optional workflow options (e.g. profiles)." + "description": "Optional workflow options (e.g. profiles). Control/security flags are ignored." } }, "required": ["name", "targets"] @@ -200,6 +198,33 @@ def build_tool_schemas(mode: str, is_subagent: bool = False, backend=None) -> li return schemas +def coerce_stringified_args(tool_name: str, arguments: dict) -> dict: + """Coerce args the model serialized as JSON strings back to their declared type. + + Some providers stringify nested object/array parameters even when the tool + schema says ``type: object`` / ``array`` (e.g. ``opts`` or ``query`` arriving + as a JSON string). Downstream handlers then call ``.get()`` / ``**opts`` / + ``.items()`` on a ``str`` and raise ``AttributeError`` — or silently drop the + value (``_sanitize_child_opts`` returns ``{}`` for a non-dict). Parse any such + arg once, here at the tool-call boundary, so every consumer gets the declared + type. Best-effort: an unparseable value is left as-is so the handler can return + a clean error rather than crash. + + Must run BEFORE arg decryption — ``_decrypt_dict`` would otherwise treat a + stringified object as a single encrypted value. + """ + if not isinstance(arguments, dict): + return arguments + props = TOOL_SCHEMAS.get(tool_name, {}).get("function", {}).get("parameters", {}).get("properties", {}) + for key, spec in props.items(): + if spec.get("type") in ("object", "array") and isinstance(arguments.get(key), str): + try: + arguments[key] = json.loads(arguments[key]) + except (json.JSONDecodeError, TypeError, ValueError): + pass + return arguments + + def tool_call_to_action(tool_name: str, arguments: dict) -> dict | None: """Convert a tool call to an action dict compatible with existing action handlers. @@ -215,6 +240,10 @@ def tool_call_to_action(tool_name: str, arguments: dict) -> dict | None: return None if not arguments: return None + # A model may emit non-object arguments (bare JSON int/array/string) -- `.items()` + # below would raise and abort the loop, so reject cleanly and let the caller retry. + if not isinstance(arguments, dict): + return None safe_arguments = {k: v for k, v in arguments.items() if k not in {"action", "description"}} descr = safe_arguments.get("name", "") or safe_arguments.get("query") or safe_arguments.get("command", "unknown") return {"action": action_type, "description": descr, **safe_arguments} diff --git a/secator/ai/utils.py b/secator/ai/utils.py index c73b28b82..c2b3315da 100644 --- a/secator/ai/utils.py +++ b/secator/ai/utils.py @@ -1,31 +1,458 @@ # secator/ai/utils.py """Utility functions for AI task - LLM initialization, calling, and response parsing.""" +import json import logging +import os import random -from typing import Dict, List, Optional +from dataclasses import fields +from typing import Any, Dict, List, Optional, Tuple from secator.definitions import LLM_SPINNER_MESSAGES from secator.config import CONFIG from secator.output_types import Warning, Error from secator.rich import console, maybe_status +from secator.runners import Task from secator.utils import format_token_count # Module-level state for litellm initialization _llm_initialized = False +SENSITIVE_ENV_PREFIXES = ( + "SECATOR_", + "ANTHROPIC_", "OPENAI_", "GOOGLE_", "AZURE_", "AWS_", "GCP_", + "GITHUB_TOKEN", "GITLAB_TOKEN", "SLACK_TOKEN", "DISCORD_TOKEN", + "SECRET_", "TOKEN_", "API_KEY", "PRIVATE_KEY", +) + + +def _sanitized_env() -> dict: + """Return a copy of os.environ with sensitive variables removed. + + Passed as the `env` run_opt to the AI shell `command` runner so an AI-run + `env`/`printenv` can't dump the LLM key + cloud creds into output that flows + back to the LLM and is persisted to Mongo. + """ + return {k: v for k, v in os.environ.items() + if not any(k.startswith(p) for p in SENSITIVE_ENV_PREFIXES) + and "KEY" not in k and "SECRET" not in k and "TOKEN" not in k and "PASSWORD" not in k} + + +def _build_action_display(action: Dict) -> str: + """Build a display string for the action being checked. + + Returns a concise description of the command/task/workflow for prompt context. + """ + action_type = action.get("action", "") + if action_type == "shell": + return action.get("command", "") + elif action_type in ("task", "workflow"): + name = action.get("name", "") + targets = action.get("targets", []) + opts = action.get("opts", {}) + parts = [f"{action_type}: {name}"] + if targets: + parts.append(f"targets={targets}") + if opts: + parts.append(f"opts={opts}") + return " ".join(parts) + return "" + + +def _is_approved(response) -> bool: + # Explicit allow-list: only a normalized "allow" answer approves. None, "deny", + # or any unexpected token denies (fail closed) — so a new backend or a refactored + # answer vocabulary can't silently approve via a "not deny" gap. + return bool(response) and response.get("answer") == "allow" + + +def _truncate(text: str, max_chars: int) -> str: + """Cap ``text`` to ~``max_chars``, keeping head + tail so both the start and the + final lines survive, with a clear marker for the dropped middle. Short text is + returned unchanged (no marker).""" + if len(text) <= max_chars: + return text + dropped = len(text) - max_chars + half = max_chars // 2 + return f"{text[:half]}\n…(truncated {dropped} chars)…\n{text[-(max_chars - half):]}" + + +def _format_action_error(e: Exception, max_chars: int = 400) -> str: + """Build a concise, LLM-facing error string for a failed action dispatch. + + Combines the exception type + message with the last few traceback frames (that's + where the actual failure is) so the model can see *where* it failed, then + truncates so a deep traceback can't blow up the next prompt's token budget. + """ + import traceback + + errtype = type(e).__name__ + msg = str(e) + head = f"{errtype}: {msg}" if msg else errtype + + tb_lines = traceback.format_exc().strip().splitlines() + tb_tail = "\n".join(tb_lines[-6:]) if tb_lines else "" + + detail = f"{head}\n{tb_tail}" if tb_tail else head + detail = _truncate(detail, max_chars) + return ( + f"Action failed with error: {detail}\n" + "Fix the issue and try again." + ) + + +_HEAVY_PROFILES = {'large', 'extra_large'} + + +def _is_heavy_runner(runner_type: str, name: str, opts: dict = None) -> bool: + """Whether a sub-runner is too heavy to run sync in-process inside the ai worker. + + Workflows/scans fan out across multiple pools, so they should always be + dispatched rather than run in-process. A task is heavy if its (possibly + opts-dependent) profile maps to a large worker pool (``large``/``extra_large``). + """ + if runner_type != 'task': + return True + try: + cls = Task.get_task_class(name) + except Exception: + return False + profile = getattr(cls, 'profile', 'small') + if callable(profile): + try: + profile = profile(opts or {}) # resolve dynamic profile (mirrors Command.s/si) + except Exception: + return True # can't resolve — be conservative and dispatch + return profile in _HEAVY_PROFILES + + +# Framework control/security keys the LLM must never set on a spawned sub-runner +# (esp. `dangerous`, which skips the permission engine). Task/workflow scan opts +# (nmap ports, httpx rate_limit, ...) are not control keys and pass through. +_FORBIDDEN_CHILD_OPT_KEYS = frozenset({ + "dangerous", + "interactive", + "hooks", + "sync", + "subagent", + "tty", + "dry_run", + "exporters", + "enable_reports", +}) + +# Cap a spawned subagent's iteration budget so it can't be told to loop unbounded. +_MAX_CHILD_ITERATIONS = 25 + + +def _sanitize_child_opts(opts: Any) -> Dict: + """Drop LLM-settable control/security keys from sub-runner opts; clamp max_iterations.""" + if not isinstance(opts, dict): + return {} + clean = {} + for key, value in opts.items(): + k = str(key) + if k in _FORBIDDEN_CHILD_OPT_KEYS or k.startswith("print_"): + continue + clean[key] = value + # Clamp the AI-subagent iteration budget (bool is an int subclass — drop it). + mi = clean.get("max_iterations") + if isinstance(mi, bool): + clean.pop("max_iterations", None) + elif isinstance(mi, (int, float)): + clean["max_iterations"] = max(1, min(int(mi), _MAX_CHILD_ITERATIONS)) + elif mi is not None: + clean.pop("max_iterations", None) + return clean + + +def build_subagent_prompt(objective: str, targets: list, evidence: str) -> str: + """Wrap the LLM-supplied subagent objective in a structured prompt. + + The `objective` is used verbatim (the parent LLM's intent). `targets` scopes + the work; `evidence` (auto-gathered, may be empty) is prior findings the + subagent should NOT re-discover. + """ + targets_str = ", ".join(str(t) for t in targets) if targets else "(inherit parent scope)" + evidence_block = evidence.strip() if evidence.strip() else "(none — no prior findings for this scope)" + return ( + f"## Objective\n{objective.strip() or '(no explicit objective given)'}\n\n" + f"## Scope\nWork ONLY within these target(s): {targets_str}\n\n" + f"## Already known (do not re-run tools that would re-discover these)\n{evidence_block}\n\n" + f"## Expected output\nInvestigate the objective, then report your findings concisely. " + f"Persist any new findings; do not repeat work already listed under 'Already known'." + ) + + +def _union_live_results(persisted: List[Dict], live_results: List[Dict], query_filter: Dict, limit: int) -> List[Dict]: + """Union backend results with this run's in-memory findings (local driver only). + + The live findings are filtered by the SAME query via an in-memory json backend, + then merged into the backend (disk) results and deduped by ``_uuid`` (backend wins), + respecting ``limit``. Makes query_workspace the single source of truth under the + local driver, whose JSON exporter only writes to disk at end-of-run. + """ + if not live_results: + return persisted + from secator.query import QueryEngine + # workspace_id "" + a `results` context => an in-memory json backend that filters + # the provided results by the query (no disk access). + live = QueryEngine("", context={"results": live_results}).search(query_filter, limit=limit or 0) + seen = {r.get("_uuid") for r in persisted if r.get("_uuid")} + for r in live: + u = r.get("_uuid") + if u and u in seen: + continue + persisted.append(r) + if u: + seen.add(u) + return persisted[:limit] if limit else persisted + + +def _resolve_field_type(f) -> Optional[type]: + """Resolve a dataclass field's declared type to a concrete builtin type. + + Mirrors ``OutputType.validate_fields``: ``f.type`` may be an actual type + (``bool``) or — under ``from __future__ import annotations`` — a string + annotation (``'bool'``). Returns the concrete type (``bool``/``int``/ + ``float``/``list``/``dict``/``str``) or ``None`` if it can't be resolved. + """ + t = f.type + # Actual type, e.g. bool / int / float / str + if isinstance(t, type): + return t + # Typing generic, e.g. List[str] -> list + origin = getattr(t, '__origin__', None) + if origin is not None: + return origin + # String annotation, e.g. 'bool', 'int', "List[str]" + if isinstance(t, str): + name = t.split('[', 1)[0].strip().lower() + return { + 'bool': bool, 'int': int, 'float': float, + 'str': str, 'list': list, 'dict': dict, + }.get(name) + return None + + +def _coerce_finding_fields(cls, data: Dict) -> Dict: + """Coerce AI-provided scalar values to a finding class's declared field types. + + LLMs frequently emit wrong-typed scalars (a ``bool`` field as the string + ``"true"``, an ``int`` as ``"3"``). This fixes *obvious* type mismatches + before validation so the finding isn't rejected for model type sloppiness. + + Only coerces when safe; unknown keys, already-correct values, and + unparseable values are left untouched (validation will still surface a real + error rather than silently dropping data). + """ + field_types = {f.name: _resolve_field_type(f) for f in fields(cls)} + for key, value in list(data.items()): + if key.startswith('_'): + continue + expected = field_types.get(key) + if expected is None or value is None: + continue + # Already the right type (note: bool is a subclass of int, so guard it). + if isinstance(value, expected) and not (expected is int and isinstance(value, bool)): + continue + + if expected is bool: + if isinstance(value, bool): + continue + if isinstance(value, int): + data[key] = bool(value) + elif isinstance(value, str): + s = value.strip().lower() + if s in ('true', '1', 'yes', 'on'): + data[key] = True + elif s in ('false', '0', 'no', 'off', ''): + data[key] = False + elif expected is int: + # Avoid coercing real bools into ints. + if isinstance(value, bool): + continue + if isinstance(value, float): + if value.is_integer(): + data[key] = int(value) + elif isinstance(value, str): + try: + data[key] = int(value) + except ValueError: + try: + f_val = float(value) + if f_val.is_integer(): + data[key] = int(f_val) + except ValueError: + pass + elif expected is float: + if isinstance(value, bool): + continue + if isinstance(value, int): + data[key] = float(value) + elif isinstance(value, str): + try: + data[key] = float(value) + except ValueError: + pass + elif expected is list: + if isinstance(value, str): + s = value.strip() + if s.startswith('['): + try: + parsed = json.loads(s) + if isinstance(parsed, list): + data[key] = parsed + except (json.JSONDecodeError, TypeError): + pass + # str fields: leave as-is (don't stringify); unknown types: leave untouched. + return data + + +def _get_action_label(action: Dict) -> str: + """Get a display label for an action.""" + act_type = action.get("action", "unknown") + if act_type in ("task", "workflow"): + name = action.get("name", "?") + opts = action.get("opts", {}) + # Defensive: a model may stringify `opts` (coerced at the tool-call boundary, + # but a malformed value can survive as a str) — never crash a display label. + session_name = opts.get("session_name", "") if isinstance(opts, dict) else "" + if session_name: + return session_name + targets = action.get("targets", []) + target_str = targets[0] if len(targets) == 1 else f"{len(targets)} targets" + return f"{name} on {target_str}" + elif act_type == "shell": + cmd = action.get("command", "")[:40] + return f"shell: {cmd}" + return act_type + + +def _decrypt_dict(d: Dict, encryptor: Any) -> Dict: + """Recursively decrypt all string values in a dict. + + Args: + d: Dictionary to decrypt + encryptor: SensitiveDataEncryptor instance + + Returns: + Decrypted dictionary + """ + # Backstop: callers should pass a dict, but a non-dict (e.g. an LLM that + # stringified an object arg) must not raise `.items()` here — return it + # unchanged rather than crash the whole action. + if not isinstance(d, dict): + return d + result = {} + for k, v in d.items(): + if isinstance(v, str): + result[k] = encryptor.decrypt(v) + elif isinstance(v, dict): + result[k] = _decrypt_dict(v, encryptor) + elif isinstance(v, list): + result[k] = [ + encryptor.decrypt(i) if isinstance(i, str) + else _decrypt_dict(i, encryptor) if isinstance(i, dict) + else i + for i in v + ] + else: + result[k] = v + return result + + +def _tool_call_fields(tc) -> Tuple[str, str]: + """Extract (name, arguments) from a tool_call, handling both the dict shape + (litellm/OpenAI JSON) and the SDK object shape (attribute access).""" + fn = tc.get("function", {}) if isinstance(tc, dict) else getattr(tc, "function", None) + if isinstance(fn, dict): + return fn.get("name", ""), fn.get("arguments", "") + if fn is not None: + return getattr(fn, "name", ""), getattr(fn, "arguments", "") + return "", "" + + +def _strip_leading_orphan_tools(messages: List[Dict]) -> int: + """Drop leading 'tool' (tool_result) messages with no preceding tool_use. + + Truncation/compaction drops the OLDEST messages with no tool-pairing + awareness, so the kept window can START with a tool_result whose + assistant(tool_calls) parent was dropped. Anthropic/OpenAI reject such a + leading orphan tool_result ("tool_result without matching tool_use"). + System messages are preserved; we scan past them and drop the run of + leading 'tool' messages that follows. Mutates `messages` in place. + + Args: + messages: List of message dicts in litellm/OpenAI format. + + Returns: + Number of leading orphan tool messages removed. + """ + i = 0 + while i < len(messages) and messages[i].get("role") == "system": + i += 1 + removed = 0 + while i < len(messages) and messages[i].get("role") == "tool": + messages.pop(i) + removed += 1 + return removed + + +def _dedupe_tool_results(messages: List[Dict]) -> int: + """Drop duplicate tool_result messages sharing a tool_call_id. + + Anthropic (and OpenRouter's providers) fold consecutive 'tool' messages into a + single user turn and reject more than one tool_result per tool_use id + ("each tool_use must have a single result. Found multiple tool_result blocks + with id X") — a NON-retryable 400. Duplicates arise when batch results are + grouped out of order (itertools.groupby only groups *consecutive* keys), or + when history trim/compaction restructures the window. Within each run of + consecutive 'tool' messages, keep the first result for each id and drop the + rest (in place). Returns the number removed. + """ + removed = 0 + i = 0 + while i < len(messages): + if messages[i].get("role") != "tool": + i += 1 + continue + seen = set() + j = i + while j < len(messages) and messages[j].get("role") == "tool": + tc_id = messages[j].get("tool_call_id") + if tc_id is not None and tc_id in seen: + del messages[j] + removed += 1 + continue # a message shifted into j; re-check without advancing + if tc_id is not None: + seen.add(tc_id) + j += 1 + i = j + return removed + + def _repair_orphan_tool_uses(messages: List[Dict]) -> int: - """Insert synthetic tool_result messages for orphan assistant tool_use blocks. + """Repair orphan tool_use/tool_result pairing for Anthropic/OpenAI. - Anthropic rejects requests where an assistant tool_use block is not - immediately followed by a matching tool_result. Mutates `messages` in place. + Two defects are fixed (both mutate `messages` in place): + - LEADING orphan tool_results: a kept window starting with a tool_result + whose assistant(tool_calls) parent was trimmed away (see + `_strip_leading_orphan_tools`). + - FORWARD orphan tool_uses: an assistant tool_use block not immediately + followed by a matching tool_result (synthesize an acknowledged result). Args: messages: List of message dicts in litellm/OpenAI format. Returns: - Number of synthetic tool_results inserted. + Number of messages removed or synthetic tool_results inserted. """ + # Leading orphan tool_results have no parent in this window — drop them. + repaired = _strip_leading_orphan_tools(messages) + # Duplicate tool_results for one id are rejected as a non-retryable 400 — drop + # extras so the request is valid (and, when hit as a 400, so the retry repairs it). + repaired += _dedupe_tool_results(messages) inserted = 0 i = 0 while i < len(messages): @@ -51,11 +478,7 @@ def _repair_orphan_tool_uses(messages: List[Dict]) -> int: tc_id = tc.get("id") if isinstance(tc, dict) else getattr(tc, "id", None) if not tc_id or tc_id in satisfied: continue - fn = tc.get("function", {}) if isinstance(tc, dict) else getattr(tc, "function", None) - if isinstance(fn, dict): - name = fn.get("name", "") - else: - name = getattr(fn, "name", "") if fn else "" + name, _ = _tool_call_fields(tc) to_insert.append({ "role": "tool", "tool_call_id": tc_id, @@ -69,7 +492,7 @@ def _repair_orphan_tool_uses(messages: List[Dict]) -> int: j += len(to_insert) i = j - return inserted + return repaired + inserted def init_llm(api_key: Optional[str] = None): @@ -146,9 +569,8 @@ def log_pre_api_call(self, model, messages, kwargs): tool_name = msg.get("name", msg.get("tool_call_id", "")) title_extra = f" [dim]{tool_name}[/]" try: - import json as _json from rich.pretty import Pretty - data = _json.loads(content) + data = json.loads(content) renderable = Pretty(data) except (ValueError, TypeError): pass @@ -192,19 +614,49 @@ def log_success_event(self, kwargs, response_obj, start_time, end_time): _llm_initialized = True +def _estimate_usage(model: str, messages: List[Dict], content: str, tool_calls) -> Dict: + """Estimate tokens when the provider omits `usage`, so calls are never unmetered. + + Uses litellm's own token counter for the model in use — prompt tokens from the + request messages, completion tokens from the response text (+ any tool-call + name/arguments). Returns the same shape as the real-usage dict (cost unknown). + """ + import litellm + + def _count(**kw): + try: + return litellm.token_counter(model=model, **kw) or 0 + except Exception: + return 0 + + prompt_tokens = _count(messages=messages) + completion_text = content or "" + for tc in tool_calls or []: + name, args = _tool_call_fields(tc) + completion_text += f" {name} {args}" + completion_tokens = _count(text=completion_text) + return { + "tokens": prompt_tokens + completion_tokens, + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "cost": None, + } + + def call_llm( messages: List[Dict], model: str, temperature: float = 0.7, api_base: Optional[str] = None, api_key: Optional[str] = None, - max_retries: int = 3, tools: Optional[List[Dict]] = None, ) -> Dict: """Call litellm completion and return response with usage.""" import time import litellm + max_retries = 3 + # Initialize litellm once (avoids callback accumulation) init_llm(api_key=api_key) @@ -230,25 +682,30 @@ def call_llm( # a matching tool_result). Safety net in case the caller bypassed ChatHistory. _repair_orphan_tool_uses(kwargs["messages"]) + # 400s are non-transient (malformed request, context_length_exceeded, ...) — + # handled separately below and NOT in this transient-retry tuple. retryable = ( litellm.InternalServerError, litellm.RateLimitError, - litellm.ServiceUnavailableError, litellm.APIConnectionError, litellm.BadRequestError, + litellm.ServiceUnavailableError, litellm.APIConnectionError, litellm.APIError ) for attempt in range(1, max_retries + 1): try: response = litellm.completion(**kwargs) break - except retryable as e: - # Detect the specific "orphan tool_use" error and repair before retry. + except litellm.BadRequestError as e: + # 400s fail fast, except the orphan tool_use case which we repair + # and retry (not counted as an attempt — the repair is the real fix). err_str = str(e) if 'tool_use' in err_str and 'tool_result' in err_str: repaired = _repair_orphan_tool_uses(kwargs["messages"]) if repaired: console.print(Warning( message=f"Repaired {repaired} orphan tool_use block(s); retrying LLM call.")) - # Don't count this as a retry attempt — the repair is the real fix. continue + console.print(Error(message=f"LLM call failed with non-retryable 400: {e}")) + raise + except retryable as e: if attempt < max_retries: wait = 2 ** attempt console.print(Warning( @@ -275,8 +732,16 @@ def call_llm( usage = { "tokens": response.usage.total_tokens, + "prompt_tokens": getattr(response.usage, "prompt_tokens", None), + "completion_tokens": getattr(response.usage, "completion_tokens", None), "cost": cost, } + else: + # usage missing/empty (streaming, some models) — estimate so the call + # is still metered instead of silently counting 0 tokens. + usage = _estimate_usage(model, kwargs["messages"], content, getattr(message, 'tool_calls', None)) + console.print(Warning( + message=f"LLM response missing usage; estimated ~{usage['tokens']} tokens for metering.")) # Get tool calls tool_calls = getattr(message, 'tool_calls', None) or [] @@ -292,8 +757,9 @@ def call_llm( ] -def format_llm_status(token_count, ctx_window, by_role): - """Format a rich status message for LLM calls with token counts and a spinner message.""" +def _format_token_breakdown(token_count, ctx_window, by_role): + """Format the token/context-window/per-role strings shared by the LLM status + spinner and the prompt_user title recap.""" token_str = format_token_count(token_count, icon='arrow_up', compact=True) ctx_str = format_token_count(ctx_window, compact=True) role_parts = [] @@ -301,6 +767,12 @@ def format_llm_status(token_count, ctx_window, by_role): if role in by_role: role_parts.append(f'[orange4]{role}[/]:{format_token_count(by_role[role], compact=True)}') role_str = ' | '.join(role_parts) + return token_str, ctx_str, role_str + + +def format_llm_status(token_count, ctx_window, by_role): + """Format a rich status message for LLM calls with token counts and a spinner message.""" + token_str, ctx_str, role_str = _format_token_breakdown(token_count, ctx_window, by_role) return ( f"[bold orange3]{random.choice(LLM_SPINNER_MESSAGES)}[/]" f" [gray42] • {token_str}/[dim red]{ctx_str}[/] ({role_str})[/]" @@ -313,7 +785,6 @@ def setup_ai(): from rich.prompt import Prompt # Load all models, sort, build color map - # all_models = sorted(litellm.model_list) # TODO: revise this, check why it doesn't list all models all_models = [] all_parts = set() for provider, model_names in litellm.models_by_provider.items(): @@ -327,15 +798,20 @@ def setup_ai(): all_parts.add(p) part_colors = {p: MODEL_COLORS[i % len(MODEL_COLORS)] for i, p in enumerate(sorted(all_parts))} - def _format_model(m, idx=None): + def _format_model(m, idx): parts = m.split('/') if len(parts) > 1: segments = [f"[bold {part_colors[p]}]{p}[/]" for p in parts[:-1]] colored = '/'.join(segments) + f"/[bold white]{parts[-1]}[/]" else: colored = f"[bold white]{m}[/]" - prefix = f"[dim]{idx:>4}[/] " if idx is not None else " " - return prefix + colored + return f"[dim]{idx:>4}[/] " + colored + + def _show_models(displayed, suffix, leading_newline=False): + prefix = "\n" if leading_newline else "" + console.print(f"{prefix}[bold] Found {len(displayed)} models{suffix}:[/]") + for i, m in enumerate(displayed, 1): + console.print(_format_model(m, idx=i), highlight=False) # Show current config current_model = CONFIG.addons.ai.default_model @@ -352,9 +828,7 @@ def _format_model(m, idx=None): # Display all models numbered displayed = all_models suffix = '' - console.print(f"[bold] Found {len(displayed)} models{suffix}:[/]") - for i, m in enumerate(displayed, 1): - console.print(_format_model(m, idx=i), highlight=False) + _show_models(displayed, suffix) # Enter prompt loop while True: @@ -363,9 +837,7 @@ def _format_model(m, idx=None): if not choice: # Empty input: re-show current list - console.print(f"\n[bold] Found {len(displayed)} models{suffix}:[/]") - for i, m in enumerate(displayed, 1): - console.print(_format_model(m, idx=i), highlight=False) + _show_models(displayed, suffix, leading_newline=True) continue if choice.lower() in ('q', 'quit', 'exit'): @@ -396,9 +868,7 @@ def _format_model(m, idx=None): else: displayed = filtered suffix = f' matching "{choice}"' - console.print(f"\n[bold] Found {len(displayed)} models{suffix}:[/]") - for i, m in enumerate(displayed, 1): - console.print(_format_model(m, idx=i), highlight=False) + _show_models(displayed, suffix, leading_newline=True) continue # Model selected - save config @@ -441,7 +911,7 @@ def _format_model(m, idx=None): return selected -def prompt_user(history, encryptor=None, max_iterations=10, choices=None, +def prompt_user(history, max_iterations=10, choices=None, mode="chat", model=None): """Prompt user for follow-up input via interactive menu. @@ -451,7 +921,6 @@ def prompt_user(history, encryptor=None, max_iterations=10, choices=None, Args: history: ChatHistory instance (read-only, used for token counts and compaction). - encryptor: Optional SensitiveDataEncryptor (unused, kept for compat). max_iterations: Current max iterations (used for continue message). choices: Optional list of choice strings from LLM follow_up action. model: Optional LLM model name for token count display. @@ -465,7 +934,6 @@ def prompt_user(history, encryptor=None, max_iterations=10, choices=None, return None from secator.rich import InteractiveMenu from secator.ai.prompts import format_continue - from secator.utils import format_token_count # Build title with token recap title = "What's next?" @@ -474,13 +942,7 @@ def prompt_user(history, encryptor=None, max_iterations=10, choices=None, from secator.ai.history import get_context_window by_role = history.count_tokens_by_role(model) ctx_window = get_context_window(model) - token_str = format_token_count(by_role['total'], icon='arrow_up', compact=True) - ctx_str = format_token_count(ctx_window, compact=True) - role_parts = [] - for role in ('system', 'user', 'assistant', 'tool'): - if role in by_role: - role_parts.append(f'[orange4]{role}[/]:{format_token_count(by_role[role], compact=True)}') - role_str = ' | '.join(role_parts) + token_str, ctx_str, role_str = _format_token_breakdown(by_role['total'], ctx_window, by_role) title += f" [gray42]• {token_str}/[dim red]{ctx_str}[/] ({role_str})[/]" except Exception: pass @@ -575,14 +1037,9 @@ def prompt_user(history, encryptor=None, max_iterations=10, choices=None, history.compact(model) new_tokens = history.count_tokens(model) console.print(f"[bold green]Compacted context: {old_tokens} -> {new_tokens} tokens[/]") - return prompt_user(history, encryptor, max_iterations, choices, mode, model) + return prompt_user(history, max_iterations, choices, mode, model) # exit return None except (KeyboardInterrupt, EOFError): return None - - -def _maybe_encrypt(text, encryptor): - """Encrypt text if encryptor is available, otherwise return as-is.""" - return encryptor.encrypt(text) if encryptor else text diff --git a/secator/celery_signals.py b/secator/celery_signals.py index e0cdf909a..9a96e8ca0 100644 --- a/secator/celery_signals.py +++ b/secator/celery_signals.py @@ -15,6 +15,32 @@ STATE_DIR = Path("/tmp/celery_state") STATE_DIR.mkdir(exist_ok=True, parents=True) +# Eviction flag. On worker shutdown (e.g. a K8s pod SIGTERM eviction) we raise this flag; the +# running task's monitor thread (in the prefork child, a separate process) polls it via a file +# and stops early, returning partial results so the surrounding chord proceeds — instead of the +# task hanging until the broker visibility timeout redelivers it (hours, on the long pool). +SHUTDOWN_FLAG = STATE_DIR / "worker_shutdown" + + +def is_worker_shutting_down(): + """True once the worker has begun shutting down (set by worker_shutting_down_handler).""" + return SHUTDOWN_FLAG.exists() + + +def clear_shutdown_flag(): + """Remove the eviction flag (called on worker boot to drop any stale flag).""" + if SHUTDOWN_FLAG.exists(): + SHUTDOWN_FLAG.unlink() + + +def worker_shutting_down_handler(**kwargs): + """Raise the eviction flag so the in-flight task's monitor stops it early and returns.""" + try: + SHUTDOWN_FLAG.parent.mkdir(parents=True, exist_ok=True) + SHUTDOWN_FLAG.write_text("1") + except OSError as e: + console.print(Info(message=f'Failed to raise worker shutdown flag: {e}')) + def get_lock_file_path(): worker_name = os.environ.get("WORKER_NAME", f"unknown_{os.getpid()}") @@ -126,6 +152,11 @@ def setup_handlers(): if CONFIG.celery.override_default_logging: signals.setup_logging.connect(setup_logging) + # Eviction handling (always on): clear any stale flag from a previous worker in this pod, + # and raise it on shutdown so the in-flight task stops early and lets its chord proceed. + clear_shutdown_flag() + signals.worker_shutting_down.connect(worker_shutting_down_handler) + # Register common handlers when either task‐ or idle‐based termination is enabled if CONFIG.celery.worker_kill_after_task or CONFIG.celery.worker_kill_after_idle_seconds != -1: signals.celeryd_after_setup.connect(capture_worker_name) diff --git a/secator/cli.py b/secator/cli.py index 1554d4111..eb8481abf 100644 --- a/secator/cli.py +++ b/secator/cli.py @@ -3035,8 +3035,8 @@ def task(name, verbose, check, system_exit): else: return False - # Run install - if hasattr(task, 'get_version_info'): + # Run install (skip for generic tasks with no static cmd — no external tool to install/version) + if hasattr(task, 'get_version_info') and getattr(task, 'cmd', None): cmd = f'secator install tools {task_name}' ret_code = Command.execute(cmd, name='install', quiet=not verbose, cwd=ROOT_FOLDER) version_info = task.get_version_info() diff --git a/secator/exporters/markdown.py b/secator/exporters/markdown.py index 0cdc8311b..2cd77bcd9 100644 --- a/secator/exporters/markdown.py +++ b/secator/exporters/markdown.py @@ -10,7 +10,14 @@ def send(self): if not ai_items: return - sections = [item.content for item in ai_items if item.content] + # report.data['results']['ai'] holds serialized dicts, not Ai objects, so + # read `content` defensively (handle both dict and OutputType forms). + def _content(item): + if isinstance(item, dict): + return item.get('content') + return getattr(item, 'content', None) + + sections = [c for item in ai_items if (c := _content(item))] if not sections: return diff --git a/secator/hooks/_dedup.py b/secator/hooks/_dedup.py index 0675eb8b9..0a4d1d5cb 100644 --- a/secator/hooks/_dedup.py +++ b/secator/hooks/_dedup.py @@ -1,6 +1,20 @@ # secator/hooks/_dedup.py +def _is_unset(field, value): + """Return True if `value` should be treated as "empty" for copy-forward purposes. + + For most fields, emptiness is the generic falsy check (`not value`). The `status` + field (Vulnerability) is special: its default value `'NEW'` is truthy but means + "untouched", so we treat `''` / `None` / `'NEW'` as unset. This lets a prior + `ACKNOWLEDGED` / `FIXED` status carry forward onto a re-found main whose status is + still the default `'NEW'`, while a never-touched vuln stays `'NEW'`. + """ + if field == 'status': + return not value or str(value).strip().upper() == 'NEW' + return not value + + def compute_duplicate_updates(workspace_findings, untagged_findings, copy_fields=None): """Compute duplicate-tagging updates for a set of findings (backend-agnostic). @@ -35,10 +49,12 @@ def compute_duplicate_updates(workspace_findings, untagged_findings, copy_fields if not hasattr(previous_item, field): continue value_prev = getattr(previous_item, field) - if not value_prev: + # Nothing meaningful to carry forward (handles `status='NEW'` as unset too). + if _is_unset(field, value_prev): continue value_curr = getattr(item, field, None) - if not value_curr and field not in copied_fields: + # Copy only onto an "empty" current value; for `status`, `'NEW'` counts as empty. + if _is_unset(field, value_curr) and field not in copied_fields: copied_fields[field] = value_prev related_ids = [] diff --git a/secator/hooks/api.py b/secator/hooks/api.py index 3499c1c89..88fae20e5 100644 --- a/secator/hooks/api.py +++ b/secator/hooks/api.py @@ -16,6 +16,8 @@ import requests from functools import cache +from ipaddress import ip_address +from urllib.parse import urlsplit from secator.config import CONFIG from secator.output_types import FINDING_TYPES, Error, Info @@ -44,9 +46,43 @@ def get_runner_dbg(runner): return {runner.unique_name: runner.status, 'type': runner.config.type, 'class': runner.__class__.__name__, 'caller': runner.config.name, **runner.context} # noqa: E501 +def _is_loopback_host(host): + """Return True if host is localhost or a loopback IP address.""" + if not host: + return False + host = host.strip('[]').lower() + if host == 'localhost': + return True + try: + return ip_address(host).is_loopback + except ValueError: + return False + + +def _check_transport_security(url): + """Refuse cleartext transport to a non-loopback host. + + `force_ssl` only controls TLS certificate *verification*; it must never be a + way to ship targets/output/the Bearer key over plaintext HTTP to a remote + host. Allow http:// only for loopback (localhost / 127.0.0.1 / ::1) dev use. + """ + parts = urlsplit(url) + if parts.scheme == 'https': + return + if parts.scheme == 'http' and _is_loopback_host(parts.hostname): + return + raise Exception( + f'Refusing to send API data over cleartext transport to "{url}": the api driver ' + f'transmits targets, command output and the Bearer API key. Use an https:// API_URL ' + f'(addons.api.url) for remote hosts; plaintext http:// is only allowed for loopback. ' + f'Note: `force_ssl:false` controls TLS cert verification only and does not enable this.' + ) + + def _make_request(method, endpoint, data=None): """Make HTTP request to external API endpoint.""" url = f'{API_URL.rstrip("/")}/{endpoint.lstrip("/")}' + _check_transport_security(url) headers = {'Content-Type': 'application/json'} if API_KEY: headers['Authorization'] = f'{API_HEADER_NAME} {API_KEY}' diff --git a/secator/installer.py b/secator/installer.py index 099fb401f..ff54460b9 100644 --- a/secator/installer.py +++ b/secator/installer.py @@ -65,7 +65,7 @@ def install(cls, tool_cls): name = tool_cls.__name__ console.print(Info(message=f'[bold yellow]:wrench: Installing {name} ...[/]')) status = InstallerStatus.UNKNOWN - has_cmd = hasattr(tool_cls, 'cmd') + has_cmd = bool(getattr(tool_cls, 'cmd', None)) # a generic task (cmd='') has no tool to install # For non-Command tasks (e.g. PythonRunner), only proceed if they have an install method if not has_cmd and not getattr(tool_cls, 'install_cmd', None) and not getattr(tool_cls, 'pypi_dependencies', None): diff --git a/secator/output_types/_base.py b/secator/output_types/_base.py index 0a7c08620..a14c77e48 100644 --- a/secator/output_types/_base.py +++ b/secator/output_types/_base.py @@ -147,6 +147,37 @@ def toDict(self, exclude=[]): return {k: v for k, v in data.items() if k not in exclude} return data + @classmethod + def field_types(cls) -> dict: + """Resolve each non-underscore field's declared type to a concrete builtin type. + + ``f.type`` may be an actual type (``bool``), a typing generic (``List[str]``, + resolved via ``__origin__``), or — under ``from __future__ import annotations`` + — a string annotation (``'bool'``, ``'List[str]'``). Returns + ``{field_name: concrete_type}``, omitting fields that can't be resolved. + """ + type_map = { + 'bool': bool, 'int': int, 'float': float, + 'str': str, 'list': list, 'dict': dict, + } + resolved = {} + for f in fields(cls): + if f.name.startswith('_'): + continue + t = f.type + if isinstance(t, type): + resolved[f.name] = t + continue + origin = getattr(t, '__origin__', None) + if origin is not None: + resolved[f.name] = origin + continue + if isinstance(t, str): + name = t.split('[', 1)[0].strip().lower() + if name in type_map: + resolved[f.name] = type_map[name] + return resolved + @classmethod def validate_fields(cls, data: dict) -> list: """Validate data types against dataclass field definitions. @@ -155,17 +186,14 @@ def validate_fields(cls, data: dict) -> list: """ errors = [] type_names = {str: 'str', int: 'int', float: 'float', dict: 'dict', list: 'list', bool: 'bool'} + expected_types = cls.field_types() for f in fields(cls): if f.name.startswith('_') or f.name not in data: continue value = data[f.name] if value is None: continue - expected_type = f.type if isinstance(f.type, type) else None - if expected_type is None: - origin = getattr(f.type, '__origin__', None) - if origin is not None: - expected_type = origin + expected_type = expected_types.get(f.name) if expected_type and not isinstance(value, expected_type): expected_name = type_names.get(expected_type, getattr(expected_type, '__name__', str(expected_type))) actual_name = type_names.get(type(value), type(value).__name__) diff --git a/secator/output_types/ai.py b/secator/output_types/ai.py index 192b2933a..92f7cc3e4 100644 --- a/secator/output_types/ai.py +++ b/secator/output_types/ai.py @@ -67,6 +67,7 @@ def render_markdown_for_rich(text: str, title: str = '') -> str: 'query': {'label': '🟢', 'color': 'magenta'}, 'stopped': {'label': '🛑', 'color': 'orange3'}, 'follow_up': {'label': '[FOLLOW UP]', 'color': 'orange3'}, + 'steer': {'label': '[STEER]', 'color': 'cyan'}, } ACTION_TYPES = ('task', 'workflow', 'shell', 'add_finding', 'query', 'stopped') @@ -84,7 +85,16 @@ class Ai(OutputType): status: str = field(default='', compare=False) answer: str = field(default='', compare=False) choices: list = field(default_factory=list, compare=False) - session_id: str = field(default='', compare=False) + # Raw litellm message for this transcript turn (role/content/tool_calls or + # role:tool/tool_call_id). Persisted so restore_history_from_db rebuilds the + # full conversation verbatim. Empty for non-transcript ai_types (action + # displays, follow_up, shell_output, etc.). compare=False: not an identity field. + message: dict = field(default_factory=dict, compare=False) + # NOTE: no top-level `session_id` field — the conversation id is carried by + # `_context.session_id`, auto-stamped on every persisted item from the runner + # context (see ai._init_options). restore_history_from_db, the remote answer + # poll, and secator-api all correlate on `_context.session_id`. Don't re-add a + # redundant top-level field. _source: str = field(default='', repr=True, compare=False) _type: str = field(default='ai', repr=True) _timestamp: int = field(default_factory=lambda: time.time(), compare=False) @@ -162,11 +172,10 @@ def __repr__(self) -> str: action_label_str = action_label.capitalize().replace('_', ' ') line = f'{s}[bold blue]{action_label_str}[/]' content = _s(self.content) - if self.ai_type in ['task', 'workflow', 'scan']: + if self.ai_type in ['task', 'workflow']: colors = { 'task': 'bold gold3', 'workflow': 'bold dark_orange3', - 'scan': 'bold red', } color = colors[self.ai_type] content = f'[{color}]{content}[/]' diff --git a/secator/runners/_base.py b/secator/runners/_base.py index f5a591af1..c7dc9314d 100644 --- a/secator/runners/_base.py +++ b/secator/runners/_base.py @@ -324,7 +324,7 @@ def collect(holder): @property def resolved_opts(self): - opts = {k: v for k, v in self.run_opts.items() if v is not None and not k.startswith('print_') and not k.endswith('_')} # noqa: E501 + opts = {k: v for k, v in self.run_opts.items() if v is not None and not k.startswith('print_') and not k.endswith('_') and k != 'env'} # noqa: E501 sensitive = self.sensitive_opt_names if sensitive: opts = {k: (REDACTED_OPT_VALUE if (k in sensitive and v) else v) for k, v in opts.items()} diff --git a/secator/runners/command.py b/secator/runners/command.py index 5eb71c416..11a3e1e53 100644 --- a/secator/runners/command.py +++ b/secator/runners/command.py @@ -27,6 +27,11 @@ logger = logging.getLogger(__name__) +# Upper bound on the monitor thread's poll interval (seconds). Keeps the shutdown/timeout checks +# responsive even when stat_update_frequency is large, so an evicted task stops well within the +# pod's termination grace period. +MONITOR_POLL_SECONDS = 5 + class Command(Runner): """Base class to execute an external command.""" @@ -550,8 +555,13 @@ def yielder(self): self.cwd = f'{self.reports_folder}/.outputs/{self.fqn}' os.makedirs(self.cwd, exist_ok=True) - # Run the command using subprocess - env = os.environ + # Run the command using subprocess. A caller may pass a sanitized/custom env + # via the `env` run_opt (e.g. the AI shell handler strips LLM/cloud secrets so + # an AI-run `env`/`printenv` can't dump them). Defaults to the process env when + # the opt is absent (existing behavior unchanged). Uses `.get(key, default)` + # (not `or`) so an explicit empty `env={}` — a deliberate empty environment — + # is honored rather than silently falling back to the full process env. + env = self.run_opts.get('env', os.environ) self.process = subprocess.Popen( command, stdin=subprocess.PIPE if sudo_password else None, @@ -738,6 +748,12 @@ def get_max_timeout(self): def _monitor_process(self): """Monitor thread that checks process health and kills if necessary.""" + from secator.celery_signals import clear_shutdown_flag, is_worker_shutting_down + # Only honour a shutdown raised *during* this run: clear any stale flag left by a previous + # worker/task sharing this machine's state dir. In prod each task gets a fresh pod (and /tmp), + # but tests and a long-lived `secator worker` reuse it, so a leftover flag would otherwise + # wrongly stop every later task. + clear_shutdown_flag() last_stats_time = 0 while not self.monitor_stop_event.is_set(): @@ -748,6 +764,16 @@ def _monitor_process(self): current_time = time() self.debug('Collecting monitor items', sub='monitor') + # Worker is shutting down (e.g. K8s pod eviction): stop early and save partial + # results so the surrounding chord proceeds, rather than hang until the broker + # visibility timeout redelivers this task. + if is_worker_shutting_down(): + warning = Warning(message='Worker shutting down (eviction): stopping task early, saving incomplete results') + if self.monitor_queue is not None: + self.monitor_queue.put(warning) + self.stop_process(exit_ok=True, sig=signal.SIGTERM) + break + # Collect and queue stats at regular intervals if (current_time - last_stats_time) >= CONFIG.runners.stat_update_frequency: stats_items = list(self._collect_stats()) @@ -796,8 +822,10 @@ def _monitor_process(self): self.monitor_queue.put(warning) break - # Sleep for a short interval before next check (stat update frequency) - self.monitor_stop_event.wait(CONFIG.runners.stat_update_frequency) + # Wake at least every MONITOR_POLL_SECONDS so the shutdown/timeout checks stay + # responsive (stats themselves are still gated to stat_update_frequency above), so + # an eviction is caught well within the pod's termination grace period. + self.monitor_stop_event.wait(min(CONFIG.runners.stat_update_frequency, MONITOR_POLL_SECONDS)) def _collect_stats(self): """Collect stats about the current running process, if any.""" diff --git a/secator/tasks/ai.py b/secator/tasks/ai.py index 9210321f1..a220f20cf 100644 --- a/secator/tasks/ai.py +++ b/secator/tasks/ai.py @@ -1,9 +1,8 @@ # secator/tasks/ai.py """AI-powered penetration testing task.""" import json -from itertools import groupby +import uuid from pathlib import Path -from time import sleep from typing import Generator from secator.config import CONFIG @@ -15,18 +14,118 @@ from secator.runners import PythonRunner from secator.rich import console, maybe_status from secator.ai.actions import ( - ActionContext, check_guardrails, dispatch_action, _run_batch, _decrypt_dict, _build_action_display + ActionContext, check_guardrails, safe_dispatch_action, _run_batch ) from secator.ai.guardrails import PermissionEngine from secator.ai.interactivity import create_backend, RemoteBackend from secator.ai.encryption import SensitiveDataEncryptor, maybe_encrypt -from secator.ai.history import ChatHistory, truncate_to_tokens, get_context_window +from secator.ai.history import ChatHistory, truncate_to_tokens, get_context_window, cap_message from secator.ai.prompts import ( - load_prompt, get_system_prompt, get_mode_config, format_tool_result, format_continue + load_prompt, get_system_prompt, get_mode_config, format_tool_result, format_continue, MODES ) -from secator.ai.tools import build_tool_schemas, tool_call_to_action, TOOL_SCHEMAS -from secator.ai.session import save_history, show_session_picker, replay_session -from secator.ai.utils import call_llm, init_llm, setup_ai, format_llm_status +from secator.ai.tools import build_tool_schemas, tool_call_to_action, coerce_stringified_args, TOOL_SCHEMAS +from secator.ai.session import ( + save_history, show_session_picker, replay_session, restore_history_from_db, print_session_results) +from secator.ai.utils import call_llm, init_llm, setup_ai, format_llm_status, _decrypt_dict, _build_action_display + + +# High-precision cues for the deterministic mode fast-path. Only unambiguous +# prompts (cues for exactly one of attack/chat, and no exploit-ish cue) are +# resolved here; everything else defers to the LLM classifier. +_ATTACK_CUES = ( + "scan", "pentest", "pen test", "enumerate", "enumeration", "recon", + "brute", "bruteforce", "fuzz", "attack", "nmap", "nuclei", "subdomain", "hack", +) +_CHAT_CUES = ( + "summarize", "summary", "explain", "what is", "what are", "what's", + "how do", "how does", "tell me", "describe", "list the", "show me", "?", +) +_EXPLOIT_CUES = ("exploit", "poc", "proof of concept", "cve-", "vulnerabilit") + + +def fast_detect_mode(prompt): + """Cheap deterministic pre-classifier. Returns 'attack'/'chat' for + unambiguous prompts, else None to defer to the LLM. Exploit-ish prompts + return None so the LLM keeps deciding those (no behavior change there).""" + text = (prompt or "").strip().lower() + if not text: + return "chat" + if any(cue in text for cue in _EXPLOIT_CUES): + return None + has_attack = any(cue in text for cue in _ATTACK_CUES) + has_chat = any(cue in text for cue in _CHAT_CUES) + if has_attack and not has_chat: + return "attack" + if has_chat and not has_attack: + return "chat" + return None + + +def _truncate_label(text, fallback): + """Truncate ``text`` to 80 chars with an ellipsis; empty text falls back to ``fallback``.""" + return (text[:80] + '...') if text and len(text) > 80 else (text or fallback) + + +def _reject_tool_call(runner, tool_name, tool_call_id, error_msg, reason): + """Shared body for rejecting a tool call: encrypt the error, record it as the + tool result in history, and return the ``tool_result`` Ai event to yield.""" + _error_content = maybe_encrypt(error_msg, runner.encryptor) + runner.history.add_tool_result(tool_name, tool_call_id, _error_content) + return Ai(content=f"[{tool_name}] {reason}", + ai_type="tool_result", + message=cap_message( + {"role": "tool", "tool_call_id": tool_call_id, "name": tool_name, "content": _error_content}), + _context=dict(runner.context)) + + +def _reject_malformed_tool_call(runner, name, tc_id, error, extra_fields, reason): + """Build the rejected-tool-call error JSON (schema-derived hint fields) and + reject it via ``_reject_tool_call``. Shared by the two ``_process_tool_calls`` + rejection paths (malformed JSON args, unknown tool/missing args), which only + differ in the extra schema-derived fields included in the error payload.""" + error_msg = json.dumps({"error": error, **extra_fields}, separators=(',', ':')) + return _reject_tool_call(runner, name, tc_id, error_msg, reason) + + +def _yield_tool_results(runner, collected): + """Group ``collected`` results by tool_call_id, add each group's tool result to + history, and yield a summary ``Ai(ai_type="tool_result")`` per group. Split out + of ``_dispatch_and_collect``; kept a plain function (not a method, like + ``_reject_tool_call`` above) so mocked-``self`` unit tests for that method still + hit this real code instead of an auto-mocked attribute. Order-preserving dict, + NOT itertools.groupby: batch results interleave by id and groupby only groups + consecutive keys, which would emit multiple tool_result messages per tool_use + (rejected by providers). + """ + budget = runner.history.get_action_budget(runner.model) + fallback_path = Path(runner.reports_folder) / "report.json" if runner.reports_folder else None + grouped = {} + for r in collected: + grouped.setdefault(r["_context"]['tool_call_id'], []).append(r) + for tc_id, group_results in grouped.items(): + tc_name = group_results[0]["_context"]['tool_call_name'] + has_errors = any(r["_type"] == "error" for r in group_results) + serialized = [ + {k: v for k, v in r.items() if k not in INTERNAL_FIELDS} + for r in group_results + ] + tool_result_str = format_tool_result( + tc_name, "error" if has_errors else "success", + len(serialized), serialized) + tool_result_str = truncate_to_tokens( + tool_result_str, budget, runner.model, fallback_path=fallback_path) + tool_result_str = maybe_encrypt(tool_result_str, runner.encryptor) + runner.history.add_tool_result(tc_name, tc_id, tool_result_str) + _tool_msg = {"role": "tool", "tool_call_id": tc_id, "name": tc_name, "content": tool_result_str} + _runner_id = next((r.get("_context", {}).get("task_id") + or r.get("_context", {}).get("workflow_id") + or r.get("_context", {}).get("scan_id") + for r in group_results if isinstance(r, dict)), "") + yield Ai(content=f"[{tc_name}] {len(serialized)} result(s)", + ai_type="tool_result", + message=cap_message(_tool_msg), + extra_data={"runner_id": _runner_id}, + _context=dict(runner.context)) @task() @@ -39,22 +138,17 @@ class ai(PythonRunner): opts = { "name": {"type": str, "default": "", "short": "n", "internal_name": "session_name", "help": "Name for the AI session or subagent"}, # noqa: E501 "prompt": {"type": str, "default": "", "short": "p", "help": "Prompt"}, - "mode": {"type": str, "default": "", "help": "Mode: attack or chat"}, + "mode": {"type": str, "default": "", "help": f"Mode: {', '.join(MODES)}"}, # derive from MODES, don't drift "model": {"type": str, "default": CONFIG.addons.ai.default_model, "help": "LLM model"}, - # Never set a secret/CONFIG value as a task-option `default`: secator-api - # serves task opts (including defaults) to the UI, so a CONFIG default - # would leak the platform's LLM API key into the runner form. Default to - # empty; the task falls back to CONFIG.addons.ai.* at runtime in - # _init_options (api_key = passed or CONFIG.addons.ai.api_key). The - # user-supplied value is still `sensitive` so it's redacted from serialized - # runner state (run_opts/cmd) even though it's never a default. + # Never default this to CONFIG.addons.ai.api_key: secator-api serves task opts + # (incl. defaults) to the UI, which would leak the key into the runner form. + # Falls back to CONFIG at runtime instead; still `sensitive` so it's redacted. "api_key": {"type": str, "default": "", "sensitive": True, "help": "API key for LLM provider (defaults to configured key)"}, # noqa: E501 "api_base": {"type": str, "default": "", "help": "API base URL (defaults to configured base)"}, "sensitive": {"is_flag": True, "default": True, "help": "Encrypt sensitive data"}, "max_iterations": {"type": int, "default": 10, "help": "Max iterations"}, "temperature": {"type": float, "default": 0.7, "help": "LLM temperature"}, "dry_run": {"is_flag": True, "default": False, "help": "Show without executing"}, - "yes": {"is_flag": True, "default": False, "short": "y", "help": "Auto-accept"}, "intent_model": {"type": str, "default": CONFIG.addons.ai.intent_model, "help": "Model for intent detection"}, "max_tokens_total": { "type": int, "default": CONFIG.addons.ai.max_tokens_total, @@ -67,6 +161,18 @@ class ai(PythonRunner): "internal": True, "help": "Context to pass to AI (findings, scope, objective)" }, + "allowed_targets": { + "type": list, + "default": None, + "internal": True, + "help": "Platform-set allow-list of target strings/regexes the AI must stay within (e.g. validated mandates)" # noqa: E501 + }, + "denied_targets": { + "type": list, + "default": None, + "internal": True, + "help": "Platform-set deny-list of target strings/regexes the AI must never touch (deny wins over allowed_targets)" # noqa: E501 + }, "subagent": { "is_flag": True, "default": False, @@ -151,21 +257,49 @@ def yielder(self) -> Generator: if not self.model: return + # Remote (web) resume: a respawned chat task restores its history from the + # workspace Mongo `_type:"ai"` docs (headless — no local files, no TUI). + if self.interactive == "remote": + restored = yield from self._maybe_resume_remote() + if restored: + return + # Resume session if self.resume and not self.is_subagent: session = show_session_picker() if session is None: return self.session_name = session["name"] - self.history = replay_session(session) + self._reports_folder = session['folder'] + if session.get("session_id"): + # New-format session: adopt the prior session_id (instead of a fresh + # str(self.id)) so appended docs continue the same `_context.session_id`. + self.session_id = session["session_id"] + self.context["session_id"] = self.session_id + # restore_history_from_db seeds system_prompt; no prompt exists yet (asked + # interactively below), so seed the same "chat" default `_detect_mode()` + # uses — the user's next answer re-detects the real mode and overwrites it. + self.mode = self.mode or "chat" + self._rebuild_prompt_and_tools() + self.history = restore_history_from_db( + self.session_id, self._get_query_engine(), + model=self.model, encryptor=self.encryptor, system_prompt=self.system_prompt) + # Show the prior conversation + findings on the console (the unified + # restore only rebuilds in-memory history; replay_session did this for + # the legacy path, so print it here to keep resume UX consistent). + print_session_results(session) + else: + # Legacy session: its docs carry no `_context.session_id`, so the unified + # restore would rebuild empty history. Fall back to the local + # history.json replay instead. + self.history = replay_session(session) if self.history is None: yield Error(message="Failed to restore session.") return self.history.model = self.model - self._reports_folder = session['folder'] result = self._prompt_and_redetect([]) if result is None: - save_history(self.history, self.reports_folder, debug_fn=self.debug) + self._save_history() return self.context["session_name"] = self.session_name yield from result @@ -174,9 +308,7 @@ def yielder(self) -> Generator: return # Get user prompt - self.prompt = self.run_opts.get("prompt", "") - if self.prompt and Path(self.prompt).is_file(): - self.prompt = Path(self.prompt).read_text().strip() + self.prompt = self._resolve_prompt() if not self.prompt and not self.is_subagent: from secator.definitions import IN_WORKER if not IN_WORKER: @@ -191,9 +323,8 @@ def yielder(self) -> Generator: return # Setup session metadata - prompt_label = (self.prompt[:80] + '...') if self.prompt and len(self.prompt) > 80 else (self.prompt or self.mode) if not self.session_name: - self.session_name = prompt_label + self.session_name = _truncate_label(self.prompt, self.mode) if self.encryptor: self.session_name = self.encryptor.decrypt(self.session_name) self.context["session_name"] = self.session_name @@ -204,14 +335,177 @@ def yielder(self) -> Generator: self._detect_mode() # Build system prompt + start history - self.system_prompt = get_system_prompt(self.mode, workspace_path=str(self.reports_folder), backend=self.backend) + self.system_prompt = self._system_prompt_for(self.mode) self.history.set_system(maybe_encrypt(self.system_prompt, self.encryptor)) - self.history.add_user(maybe_encrypt(self.prompt, self.encryptor)) - yield Ai(content=self.prompt, ai_type="prompt") + yield self._emit_user_prompt(self.prompt) yield Info(message=f"Using model: {self.model}, mode: {self.mode}") # Run loop yield from self._run_loop() + self._mark_turn_completed() # record this turn as done so a redelivery won't replay it + + # ------------------------------------------------------------------------- + # Remote (web) session restore + # ------------------------------------------------------------------------- + + def _system_prompt_for(self, mode): + """Compute the system prompt for ``mode`` using this runner's workspace + backend.""" + return get_system_prompt(mode, workspace_path=str(self.reports_folder), backend=self.backend) + + def _rebuild_prompt_and_tools(self): + """Rebuild system_prompt + tool_schemas for the current mode and store them. + + Returns the ``(system_prompt, tool_schemas)`` pair for callers that want it.""" + self.system_prompt = self._system_prompt_for(self.mode) + self.tool_schemas = build_tool_schemas(self.mode, is_subagent=self.is_subagent, backend=self.backend) + return self.system_prompt, self.tool_schemas + + def _resolve_prompt(self): + """Resolve the ``prompt`` run option, reading it from a file if it names one.""" + prompt = self.run_opts.get("prompt", "") + if prompt and Path(prompt).is_file(): + prompt = Path(prompt).read_text().strip() + return prompt + + def _emit_user_prompt(self, prompt): + """Add ``prompt`` to history (encrypted once) and return the user-prompt Ai event to yield.""" + encrypted = maybe_encrypt(prompt, self.encryptor) + self.history.add_user(encrypted) + return Ai(content=prompt, ai_type="prompt", message={"role": "user", "content": encrypted}) + + def _get_query_engine(self): + """Build a workspace-scoped QueryEngine from the runner context. + + The backend (mongodb/api/local) is resolved from ``context['drivers']`` + via ``QueryEngine._select_backend``. For the remote channel the API + appends the ``mongodb`` driver on dispatch, so this resolves to the + workspace Mongo backend. + """ + from secator.query import QueryEngine + return QueryEngine(self.context.get("workspace_id", ""), context=dict(self.context)) + + def _maybe_resume_remote(self): + """Restore chat history from Mongo when a remote session has prior docs. + + Returns True (via generator return) if this turn was fully handled as a + respawn (history restored, loop run), False to fall through to a fresh + conversation. Yields any items produced along the way. + """ + query_engine = self._get_query_engine() + + # Guard: remote interactivity requires a Mongo-backed query engine, else + # the RemoteBackend poll can never see the web answer (and restore can't + # read the channel docs). Warn loudly but don't hard-fail a fresh run. + backend_name = getattr(query_engine.backend, "name", "") + if backend_name not in ("mongodb", "api"): + yield Warning( + message=f'interactive="remote" but query engine resolved to "{backend_name}" backend ' + '(expected mongodb/api). The web answer channel will not work — check that the ' + '`mongodb` driver is in the runner context.' + ) + + # Skip replay of an already-completed turn — acks_late can redeliver the + # same celery_id after a worker crash; without this marker we'd re-run every + # tool action and re-bill tokens. + turn_uuid = self._turn_uuid() + if turn_uuid and self._turn_completed_marker(turn_uuid, query_engine): + self.debug(f'idempotency: turn {turn_uuid} already completed; skipping replay', sub='llm') + return True + + # Look for prior `_type:"ai"` docs for this session + try: + prior = query_engine.search({"_type": "ai", "_context.session_id": self.session_id}, limit=1) + except Exception as e: # noqa: BLE001 - backend errors must not crash the worker + self.debug(f'remote resume: failed to query prior docs: {e}', sub='llm') + prior = None + + if not prior: + # Fresh conversation: nothing to restore, fall through to normal start. + return False + + # Resolve the user's new prompt (the message that triggered this respawn) + self.prompt = self._resolve_prompt() + + # Session metadata + if not self.session_name: + self.session_name = _truncate_label(self.prompt, self.prompt) + self.context["session_name"] = self.session_name + + # Detect mode (defaults to chat) and build the system prompt + tools + self._detect_mode() + self.system_prompt = self._system_prompt_for(self.mode) + + # Rebuild history from the channel docs (text-only; see restore_history_from_db) + self.history = restore_history_from_db( + self.session_id, query_engine, model=self.model, + encryptor=self.encryptor, system_prompt=self.system_prompt) + self.history.model = self.model + + # Append the new user message that respawned the conversation + if self.prompt: + yield self._emit_user_prompt(self.prompt) + + yield Info(message=f"Resumed session from DB ({len(self.history.messages)} messages), model: {self.model}, mode: {self.mode}") # noqa: E501 + yield from self._run_loop() + self._mark_turn_completed() # record this turn as done so a redelivery won't replay it + return True + + def _save_history(self): + """Persist chat history to the local reports folder, unless on the remote path. + + For the remote (web) channel the workspace Mongo `_type:"ai"` docs are the + source of truth, so the local `history.json` write is skipped. + """ + if self.interactive == "remote": + return + save_history(self.history, self.reports_folder, debug_fn=self.debug) + + # ------------------------------------------------------------------------- + # Turn-level idempotency (remote/Celery redelivery) + # ------------------------------------------------------------------------- + + def _turn_uuid(self): + """Stable id naming THIS delivery's turn for idempotency. + + ``celery_id`` (the Celery request id) is stamped on the runner context by + the worker entrypoint (``run_command``) and is the SAME across an acks_late + worker-loss redelivery, so it uniquely and idempotently names one turn. + """ + return (self.context or {}).get("celery_id") + + def _turn_completed_marker(self, turn_uuid, query_engine): + """Return the persisted completion marker for ``turn_uuid``, or None.""" + try: + docs = query_engine.search({ + "_type": "ai", + "ai_type": "turn_completed", + "_context.session_id": self.session_id, + "extra_data.turn_uuid": turn_uuid, + }, limit=1) + except Exception as e: # noqa: BLE001 - a marker query must not crash the worker + self.debug(f'idempotency: marker query failed: {e}', sub='llm') + return None + return docs[0] if docs else None + + def _mark_turn_completed(self): + """Persist a turn-completion marker once the turn is durably done. + + Remote channel only. Reuses the workspace `_type:"ai"` docs (no new + collection); restore_history_from_db skips this ai_type so it never enters + the transcript. Called by the caller AFTER `_run_loop` returns, so a crash + mid-turn leaves no marker and the partial turn still resumes. + """ + if self.interactive != "remote": + return + turn_uuid = self._turn_uuid() + if not turn_uuid: + return + self.add_result(Ai( + content="", + ai_type="turn_completed", + status="completed", + extra_data={"turn_uuid": turn_uuid}, + ), print=False) # ------------------------------------------------------------------------- # _run_loop: main LLM interaction loop @@ -225,6 +519,8 @@ def _run_loop(self) -> Generator: ctx = ActionContext( targets=self.inputs, model=self.model, + api_key=self.api_key, + api_base=self.api_base, encryptor=self.encryptor, dry_run=self.dry_run, verbose=self.verbose, @@ -247,18 +543,27 @@ def _run_loop(self) -> Generator: iteration = 0 query_extensions = 0 empty_streak = 0 + rate_limit_streak = 0 self._context_warnings_shown = set() while iteration < self.max_iterations: iteration += 1 try: + # Mid-flight steering: drain any user messages sent WHILE the agent + # was running and inject them into history so the next turn redirects. + # Cheap query per iteration; robust (never crashes the loop). + yield from self._drain_steers() + # Auto-summarize when context > 85% threshold yield from self._summarize_auto() # Prompt user when context is filling up (local only) yield from self._summarize_user() + # Roll any billed summarization usage into context.ai_tokens + self._drain_history_usage() + # Subagent token usage (for batch progress tracking) if self.is_subagent: by_role = self.history.count_tokens_by_role(self.model) @@ -278,10 +583,18 @@ def _run_loop(self) -> Generator: with maybe_status(msg, spinner="dots"): result = call_llm(messages, self.model, self.temp, self.api_base, self.api_key, tools=self.tool_schemas) + # reset rate-limit guard on success + rate_limit_streak = 0 + content = result["content"] tool_calls = result.get("tool_calls", []) usage = result.get("usage", {}) + # Accumulate billed tokens for this run (read by the billing chore + # as context.ai_tokens). Done here, before any empty-response + # `continue`, so every billed call is counted exactly once. + self._account_usage(usage) + self.debug(f'content: {content[:200] if content else "(empty)"}', sub='llm') # Empty response @@ -290,36 +603,36 @@ def _run_loop(self) -> Generator: yield Warning(message="LLM returned empty response") if empty_streak >= 3: yield Error(message="3 consecutive empty responses - the model may not support tool calling. Stopping.") - save_history(self.history, self.reports_folder, debug_fn=self.debug) + self._save_history() return continue empty_streak = 0 - # Add assistant message to history - self._add_assistant_to_history(content, tool_calls) - - # Yield response content - if content: - display_content = self.encryptor.decrypt(content) if self.encryptor else content - yield Ai( - content=display_content, - ai_type="response", - mode=self.mode, - model=self.intent_model, - summary=not tool_calls, - extra_data={ - "iteration": iteration, - "max_iterations": self.max_iterations, - "tokens": usage.get("tokens") if usage else None, - "cost": usage.get("cost") if usage else None, - }, - ) + # Add assistant message to history and capture it for persistence + assistant_msg = self._add_assistant_to_history(content, tool_calls) + + # Persist the assistant turn (even tool-call-only turns carry tool_calls) + display_content = (self.encryptor.decrypt(content) if (self.encryptor and content) else (content or '')) + yield Ai( + content=display_content, + ai_type="response", + mode=self.mode, + model=self.intent_model, + summary=not tool_calls, + message=cap_message(assistant_msg), + extra_data={ + "iteration": iteration, + "max_iterations": self.max_iterations, + "tokens": usage.get("tokens") if usage else None, + "cost": usage.get("cost") if usage else None, + }, + ) # Process tool calls → validated actions follow_up_choices = None stop_reason = None - follow_up_ai = None + follow_up_prompt_uuid = None if tool_calls: actions = yield from self._process_tool_calls(tool_calls, ctx) @@ -335,7 +648,7 @@ def _run_loop(self) -> Generator: dispatch_result = yield from self._dispatch_and_collect(actions, ctx) follow_up_choices = dispatch_result.get("follow_up_choices") stop_reason = dispatch_result.get("stop_reason") - follow_up_ai = dispatch_result.get("follow_up_ai") + follow_up_prompt_uuid = dispatch_result.get("follow_up_prompt_uuid") if len(actions) > 1: yield Info(message=f"Executed {len(actions)} actions.") @@ -347,20 +660,25 @@ def _run_loop(self) -> Generator: # Stop tool → save and exit if stop_reason is not None: - save_history(self.history, self.reports_folder, debug_fn=self.debug) + self._save_history() return # Follow-up / content-only / max_iter → prompt user if follow_up_choices is not None or not tool_calls or iteration == self.max_iterations: - # For remote follow-up, yield the pending Ai so frontend can show it - if follow_up_ai and isinstance(self.backend, RemoteBackend): - follow_up_ai.status = "pending" - follow_up_ai.session_id = self.session_id - yield follow_up_ai - - result = self._prompt_and_redetect(follow_up_choices or []) + # Remote follow-up: the pending Ai doc was already stamped + persisted + # in _dispatch_and_collect (dedup by _uuid) — nothing to re-yield here. + + # Remote max-iter after tool work is a terminal turn (no further + # user input expected) — don't block-poll on prompt_uuid=None with no + # answerable pending doc; end cleanly via the loop tail (save + Info). + if (isinstance(self.backend, RemoteBackend) + and iteration == self.max_iterations + and follow_up_choices is None and tool_calls): + break + + result = self._prompt_and_redetect(follow_up_choices or [], prompt_uuid=follow_up_prompt_uuid) if result is None: - save_history(self.history, self.reports_folder, debug_fn=self.debug) + self._save_history() return yield from result continue @@ -374,38 +692,42 @@ def _run_loop(self) -> Generator: yield Warning(message="Interrupted by user.") result = self._prompt_and_redetect([]) if result is None: - save_history(self.history, self.reports_folder, debug_fn=self.debug) + self._save_history() return yield from result continue except Exception as e: if isinstance(e, litellm.RateLimitError): - yield Warning(message="Rate limit exceeded - waiting 5s and retry in the next iteration") - iteration -= 1 - sleep(5) + # call_llm already backed off (~2/4/8s); don't re-sleep. Bound consecutive + # 429s so a persistent rate limit can't spin forever; let iteration advance. + rate_limit_streak += 1 + if rate_limit_streak >= 4: + yield Error(message="Rate limit exceeded on 4 consecutive attempts - aborting. Check your provider quota/billing.") # noqa: E501 + self._save_history() + return + yield Warning(message=f"Rate limit exceeded (attempt {rate_limit_streak}/4) - retrying in the next iteration") continue elif isinstance(e, litellm.AuthenticationError): yield Error(message=str(e)) yield Error(message='Please set a valid API key with `secator config set addons.ai.api_key `') - save_history(self.history, self.reports_folder, debug_fn=self.debug) + self._save_history() return elif isinstance(e, litellm.APIConnectionError) or ( isinstance(e, litellm.InternalServerError) and 'connection error' in str(e).lower() ): - # Genuine connectivity failures (connection refused, DNS failure) surface in - # some litellm versions as InternalServerError("Connection error.") rather than - # APIConnectionError, so catch both and gate the latter on the connection message - # to avoid swallowing unrelated upstream 500 errors. + # Some litellm versions surface connectivity failures as InternalServerError + # instead of APIConnectionError, so catch both, gated by message to avoid + # swallowing unrelated upstream 500s. yield Error(message=f"Cannot connect to model '{self.model}': {e}") yield Error(message='Check api_base and connectivity: `secator config set addons.ai.api_base `') - save_history(self.history, self.reports_folder, debug_fn=self.debug) + self._save_history() return yield Error.from_exception(e) - save_history(self.history, self.reports_folder, debug_fn=self.debug) + self._save_history() return - save_history(self.history, self.reports_folder, debug_fn=self.debug) + self._save_history() yield Info(message=f"Reached max iterations ({iteration}/{self.max_iterations})") # ------------------------------------------------------------------------- @@ -433,6 +755,8 @@ def _init_options(self): self.passed_context = self.run_opts.get("context") or {} self.async_tasks = self.get_opt_value("async_tasks") self.dangerous = self.get_opt_value("dangerous") + self.allowed_targets = self.get_opt_value("allowed_targets") or [] + self.denied_targets = self.get_opt_value("denied_targets") or [] # Interactive mode: "local" / "remote" / "auto" interactive = self.get_opt_value("interactive") @@ -452,15 +776,43 @@ def _init_options(self): self.history = ChatHistory() self.encryptor = SensitiveDataEncryptor() if self.sensitive else None self.has_previous_results = len(self.results) > 0 - self.scope = "current" if self.has_previous_results > 0 else "workspace" + self.scope = "current" if self.has_previous_results else "workspace" self.permission_engine = PermissionEngine( CONFIG.addons.ai.permissions, targets=self.inputs, - workspace=self.reports_folder or "" + workspace=self.reports_folder or "", + allowed_targets=self.allowed_targets, + denied_targets=self.denied_targets, ) - # Create interactivity backend - self.session_id = self.session_name or str(self.id) + # Per-run billed-token accounting (AI analog of context.scan_hours), read + # by the platform billing chore. Init so it persists even with zero LLM calls. + self.context.setdefault("ai_tokens", 0) + self.context.setdefault("ai_prompt_tokens", 0) + self.context.setdefault("ai_completion_tokens", 0) + self.context.setdefault("ai_cost", 0.0) + + # Record the resolved model id so the metering chore can price tokens against + # the model registry. Set unconditionally (not setdefault) — records the + # configured model even if the user switches mid-session. + self.context["ai_model"] = self.model + + # Create interactivity backend. For remote (web), the UI reuses a stable + # session_id on respawn so a respawned task finds its prior docs; it arrives + # via self.context (authoritative — the dispatcher pops run_opts['context']). + self.session_id = ( + self.passed_context.get("session_id") + or (self.context or {}).get("session_id") + or self.session_name + or str(self.id) + ) + # Write session_id back onto the context: every persisted item copies + # self.context into `_context`, so this stamps `_context.session_id` on all + # `_type:"ai"` docs (incl. prompt/response turns yielded directly here). + # restore_history_from_db + the remote poll key on it, so skipping this + # would leave the transcript unqueryable and resume would restore nothing. + if self.context is not None: + self.context["session_id"] = self.session_id self.backend = create_backend(self.interactive, timeout=CONFIG.addons.ai.user_response_timeout) # Auto-approve workspace targets @@ -507,31 +859,37 @@ def _detect_mode(self, force=False): old_mode = self.mode if old_mode and not force: if not hasattr(self, 'tool_schemas'): - self.system_prompt = get_system_prompt(self.mode, workspace_path=str(self.reports_folder), backend=self.backend) - self.tool_schemas = build_tool_schemas(self.mode, is_subagent=self.is_subagent, backend=self.backend) + self._rebuild_prompt_and_tools() return if not self.prompt: self.mode = "chat" return - try: - selection_prompt = load_prompt("modes/_selection.txt") - messages = [{"role": "user", "content": f"{selection_prompt}\n{self.prompt}"}] - with maybe_status("[bold orange3]Detecting intent...[/]", spinner="dots"): - result = call_llm(messages, self.intent_model, temperature=0.3, api_base=self.api_base, api_key=self.api_key) - mode = result["content"].strip().lower() - if mode in ("attack", "chat"): - console.print(rf"[bold green]\[INF][/] Detected intent: [bold]{mode}[/]") - self.mode = mode - else: - self.mode = old_mode or "chat" - except Exception: - console.print(Warning(message='Could not detect mode using LLM. Falling back to "chat" mode.')) - self.mode = "chat" + # Resolve unambiguous prompts deterministically; skip the intent LLM round-trip. + fast_mode = fast_detect_mode(self.prompt) + if fast_mode: + console.print(rf"[bold green]\[INF][/] Detected intent: [bold]{fast_mode}[/] (fast-path)") + self.mode = fast_mode + else: + try: + selection_prompt = load_prompt("modes/_selection.txt") + messages = [{"role": "user", "content": f"{selection_prompt}\n{self.prompt}"}] + with maybe_status("[bold orange3]Detecting intent...[/]", spinner="dots"): + result = call_llm(messages, self.intent_model, temperature=0.3, api_base=self.api_base, api_key=self.api_key) # noqa: E501 + self._account_usage(result.get("usage")) + mode = result["content"].strip().lower() + if mode in MODES: # honor any real mode (incl. exploit), don't discard it + console.print(rf"[bold green]\[INF][/] Detected intent: [bold]{mode}[/]") + self.mode = mode + else: + self.mode = old_mode or "chat" + except Exception: + console.print(Warning(message='Could not detect mode using LLM. Falling back to "chat" mode.')) + self.mode = "chat" if not self.mode: self.mode = "chat" mode_max = get_mode_config(self.mode).get("max_iterations", self.max_iterations) self.max_iterations = max(self.max_iterations, mode_max) - self.system_prompt = get_system_prompt(self.mode, workspace_path=str(self.reports_folder), backend=self.backend) + self.system_prompt = self._system_prompt_for(self.mode) if not hasattr(self, 'tool_schemas') or not old_mode or old_mode != self.mode: self.tool_schemas = build_tool_schemas(self.mode, is_subagent=self.is_subagent, backend=self.backend) @@ -545,8 +903,7 @@ def _auto_approve_workspace_targets(self): if not workspace_id: return try: - from secator.query import QueryEngine - engine = QueryEngine(workspace_id, context=dict(self.context)) + engine = self._get_query_engine() results = engine.search({"_type": "target"}, limit=1000) target_names = {r.get("name") or r.get("_name", "") for r in results if r} target_names.discard("") @@ -558,6 +915,48 @@ def _auto_approve_workspace_targets(self): except Exception as e: self.debug(f'[workspace] failed to query targets: {e}', sub='guardrail') + # ------------------------------------------------------------------------- + # Mid-flight steering + # ------------------------------------------------------------------------- + + def _drain_steers(self): + """Drain pending mid-flight steers and inject them into the LLM history. + + A "steer" is a user message sent WHILE the agent is running (over the + remote/web channel: a pending ``_type:"ai", ai_type:"steer"`` doc written by + ``POST /ai/conversations/{id}/steer``). At the top of each loop iteration we + drain any pending steers for this session and append each to the history as + a ``[User interjected]: …`` user message so the model sees them on the next + turn. Cooperative — not a hard cancel (Stop already does that). + + The steer doc the API wrote is itself the persisted transcript entry (it + carries ``_context.session_id``, so the UI's transcript poll surfaces it as + an "interjected" user bubble). We deliberately do NOT yield a second + ``Ai(ai_type="steer")`` echo here — that would persist a duplicate doc with + the same content and double-render in the UI. ``poll_steers`` flips the + drained doc to ``status:"consumed"`` so it injects exactly once. + + Only the RemoteBackend has a channel to poll; for every other backend this + is a no-op. Robust: a steer must never crash the run, so all backend access + is best-effort and swallowed. + + Generator (``yield from``-compatible with the loop) — currently yields no + items, but kept a generator so future transcript echoes can be added without + changing the call site. + """ + if not isinstance(self.backend, RemoteBackend): + return + try: + steers = self.backend.poll_steers(self.session_id) + except Exception as e: # noqa: BLE001 - a steer must never crash the run + self.debug(f'steer: failed to poll steers: {e}', sub='llm') + return + for content in steers: + self.debug(f'steer: injecting user interjection: {content[:120]}', sub='llm') + self.history.add_user(maybe_encrypt(f"[User interjected]: {content}", self.encryptor)) + return + yield # noqa: unreachable - keeps this a generator for `yield from` + # ------------------------------------------------------------------------- # Summarization / compaction # ------------------------------------------------------------------------- @@ -631,15 +1030,19 @@ def _process_tool_calls(self, tool_calls, ctx): self.debug(f'[tool_call] {name}: failed to parse: {tc.function.arguments[:200]}', sub='llm') schema = TOOL_SCHEMAS.get(name, {}).get("function", {}) properties = schema.get("parameters", {}).get("properties", {}) - error_msg = json.dumps({ - "error": f"Tool call '{tc_id}' rejected: malformed JSON arguments ({e})", - "raw_arguments": tc.function.arguments[:200], - "expected_schema": {k: v.get("type", "any") for k, v in properties.items()}, - "hint": "Retry with properly formatted JSON arguments.", - }, separators=(',', ':')) - self.history.add_tool_result(name, tc_id, maybe_encrypt(error_msg, self.encryptor)) + yield _reject_malformed_tool_call( + self, name, tc_id, f"Tool call '{tc_id}' rejected: malformed JSON arguments ({e})", + { + "raw_arguments": tc.function.arguments[:200], + "expected_schema": {k: v.get("type", "any") for k, v in properties.items()}, + "hint": "Retry with properly formatted JSON arguments.", + }, "malformed arguments") continue + # Coerce object/array args the model stringified (provider quirk) BEFORE + # decrypt/convert, so handlers get the declared type not a JSON string. + args = coerce_stringified_args(name, args) + # Decrypt args if self.encryptor: args = _decrypt_dict(args, self.encryptor) @@ -652,13 +1055,13 @@ def _process_tool_calls(self, tool_calls, ctx): self.debug(f'[tool_call] skipping {name}: {reason}', sub='llm') schema = TOOL_SCHEMAS.get(name, {}).get("function", {}) params = schema.get("parameters", {}) - error_msg = json.dumps({ - "error": f"Tool call '{tc_id}' rejected: {reason}", - "required_fields": params.get("required", []), - "schema": {k: v.get("type", "any") for k, v in params.get("properties", {}).items()}, - "hint": "Provide all required fields. Retry with a complete arguments object.", - }, separators=(',', ':')) - self.history.add_tool_result(name, tc_id, maybe_encrypt(error_msg, self.encryptor)) + yield _reject_malformed_tool_call( + self, name, tc_id, f"Tool call '{tc_id}' rejected: {reason}", + { + "required_fields": params.get("required", []), + "schema": {k: v.get("type", "any") for k, v in params.get("properties", {}).items()}, + "hint": "Provide all required fields. Retry with a complete arguments object.", + }, f"rejected: {reason}") continue action["tool_call_id"] = tc_id @@ -678,7 +1081,7 @@ def _process_tool_calls(self, tool_calls, ctx): denial_display = f"{denial}\n[gray42]{cmd_display}[/gray42]" if cmd_display else denial yield Warning(message=denial_display) error_msg = json.dumps({"error": denial}, separators=(',', ':')) - self.history.add_tool_result(name, tc_id, maybe_encrypt(error_msg, self.encryptor)) + yield _reject_tool_call(self, name, tc_id, error_msg, "denied") continue actions.append(action) @@ -699,9 +1102,13 @@ def _dispatch_and_collect(self, actions, ctx): follow_up_choices = None stop_reason = None follow_up_ai = None + follow_up_prompt_uuid = None is_batch = len(actions) > 1 - action_iter = _run_batch(actions, ctx) if is_batch else dispatch_action(actions[0], ctx) + # safe_dispatch_action wraps dispatch so a handler error becomes an Error item + # fed back to the LLM, instead of killing the main loop (_run_batch does the + # same internally for each of its actions). + action_iter = _run_batch(actions, ctx) if is_batch else safe_dispatch_action(actions[0], ctx) collected = [] for result in action_iter: @@ -712,12 +1119,27 @@ def _dispatch_and_collect(self, actions, ctx): is_from_subagent = isinstance(result, OutputType) and bool(result._context.get('subagent')) if isinstance(result, Ai): - self.add_result(result, print=not is_from_subagent) if result.ai_type == "follow_up": follow_up_ai = result follow_up_choices = result.choices or (result.extra_data or {}).get("choices", []) + # Persist the follow-up doc in its FINAL renderable state (add_result + # dedupes by _uuid, so the later `yield follow_up_ai` is dropped). For a + # remote run, stamp status="pending" + choices + session_id first. + if isinstance(self.backend, RemoteBackend): + follow_up_ai.status = "pending" + follow_up_ai.session_id = self.session_id + if not follow_up_ai.choices and follow_up_choices: + follow_up_ai.choices = list(follow_up_choices) + # Stamp a unique correlation id so the poll resolves only this prompt's + # answer, not a stale one from a prior turn (not reusing _uuid, which + # mongo may reassign on insert). + follow_up_prompt_uuid = str(uuid.uuid4()) + follow_up_ai.extra_data = { + **(follow_up_ai.extra_data or {}), "prompt_uuid": follow_up_prompt_uuid} + self.add_result(result, print=not is_from_subagent) continue - elif result.ai_type == "stopped": + self.add_result(result, print=not is_from_subagent) + if result.ai_type == "stopped": stop_reason = result.content continue if result.ai_type not in ("shell_output", "response"): @@ -739,33 +1161,69 @@ def _dispatch_and_collect(self, actions, ctx): collected.append(result) ctx.results.append(result) - # Group results by tool_call_id and add to history - budget = self.history.get_action_budget(self.model) - fallback_path = Path(self.reports_folder) / "report.json" if self.reports_folder else None - for tc_id, group in groupby(collected, key=lambda r: r["_context"]['tool_call_id']): - group_results = list(group) - tc_name = group_results[0]["_context"]['tool_call_name'] - has_errors = any(r["_type"] == "error" for r in group_results) - serialized = [ - {k: v for k, v in r.items() if k not in INTERNAL_FIELDS} - for r in group_results - ] - tool_result_str = format_tool_result( - tc_name, "error" if has_errors else "success", - len(serialized), serialized) - tool_result_str = truncate_to_tokens( - tool_result_str, budget, self.model, fallback_path=fallback_path) - tool_result_str = maybe_encrypt(tool_result_str, self.encryptor) - self.history.add_tool_result(tc_name, tc_id, tool_result_str) - - return {"follow_up_choices": follow_up_choices, "stop_reason": stop_reason, "follow_up_ai": follow_up_ai} + yield from _yield_tool_results(self, collected) + + return { + "follow_up_choices": follow_up_choices, + "stop_reason": stop_reason, + "follow_up_ai": follow_up_ai, + "follow_up_prompt_uuid": follow_up_prompt_uuid, + } # ------------------------------------------------------------------------- # History helpers # ------------------------------------------------------------------------- + def _account_usage(self, usage): + """Accumulate billed token/cost usage from a single LLM call onto the runner context. + + `usage` is the dict returned by `call_llm` + (`{"tokens", "prompt_tokens", "completion_tokens", "cost"}`) or None. + Missing/None usage counts as 0 so accounting never crashes the run. The + running total lives on `self.context["ai_tokens"]` (int, cumulative) which + is persisted onto the task doc and read by the platform billing chore. + `context["ai_prompt_tokens"]`/`["ai_completion_tokens"]` carry the split. + """ + if not usage: + return + for usage_key, ctx_key, cast in ( + ("tokens", "ai_tokens", int), + ("prompt_tokens", "ai_prompt_tokens", int), + ("completion_tokens", "ai_completion_tokens", int), + ("cost", "ai_cost", float), + ): + try: + self.context[ctx_key] = cast(self.context.get(ctx_key, 0) or 0) + cast(usage.get(usage_key) or 0) + except (TypeError, ValueError): + pass + + def _drain_history_usage(self): + """Roll billed usage accrued by history summarization into context.ai_tokens. + + `ChatHistory.compact` makes its own LLM calls and stashes their billed + usage on the history object; drain it here so it is counted exactly once. + """ + history = getattr(self, "history", None) + if history is None: + return + tokens = getattr(history, "billed_tokens", 0) or 0 + prompt_tokens = getattr(history, "billed_prompt_tokens", 0) or 0 + completion_tokens = getattr(history, "billed_completion_tokens", 0) or 0 + cost = getattr(history, "billed_cost", 0.0) or 0.0 + if tokens: + self._account_usage({ + "tokens": tokens, + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "cost": cost, + }) + history.billed_tokens = 0 + history.billed_prompt_tokens = 0 + history.billed_completion_tokens = 0 + history.billed_cost = 0.0 + def _add_assistant_to_history(self, content, tool_calls): - """Add assistant message (with optional tool calls) to chat history.""" + """Add assistant message (with optional tool calls) to chat history; return the message.""" if tool_calls: litellm_tool_calls = [{ "id": tc.id, @@ -776,24 +1234,44 @@ def _add_assistant_to_history(self, content, tool_calls): else json.dumps(tc.function.arguments)), }, } for tc in tool_calls] - self.history.add_assistant_with_tool_calls( - maybe_encrypt(content, self.encryptor) if content else None, - litellm_tool_calls) - else: - self.history.add_assistant(maybe_encrypt(content, self.encryptor)) + enc = maybe_encrypt(content, self.encryptor) if content else None + msg = {"role": "assistant", "tool_calls": litellm_tool_calls} + if enc is not None: + msg["content"] = enc + self.history.messages.append(msg) + return msg + enc = maybe_encrypt(content, self.encryptor) + msg = {"role": "assistant", "content": enc} + self.history.messages.append(msg) + return msg # ------------------------------------------------------------------------- # Follow-up / prompt # ------------------------------------------------------------------------- - def _prompt_and_redetect(self, choices): + def _prompt_and_redetect(self, choices, prompt_uuid=None): """Prompt user via backend and re-detect intent. Works for all backends: CLIBackend shows rich menus, RemoteBackend polls DB, AutoBackend returns None (exits). + ``prompt_uuid`` correlates the (remote) poll to the SPECIFIC pending + follow_up doc this call raised, so a stale answered follow_up from a prior + turn can't resolve it (which would re-inject the old prompt and loop). + Returns list of items to yield, or None to exit. """ + # Plain-chat remote turns reach here with no pre-persisted pending doc, + # so persist one now with a real prompt_uuid (never poll on prompt_uuid=None). + if isinstance(self.backend, RemoteBackend) and not prompt_uuid: + prompt_uuid = str(uuid.uuid4()) + self.add_result(self.backend.build_pending_prompt( + question="What's next?", + choices=choices, + session_id=self.session_id, + prompt_type="follow_up", + prompt_uuid=prompt_uuid, + )) response = self.backend.ask_user( question="What's next?", choices=choices, @@ -804,6 +1282,7 @@ def _prompt_and_redetect(self, choices): max_iterations=self.max_iterations, mode=self.mode, model=self.model, + prompt_uuid=prompt_uuid, ) if response is None: return None @@ -817,18 +1296,16 @@ def _prompt_and_redetect(self, choices): self.history.add_user(maybe_encrypt(answer, self.encryptor)) # Handle explicit mode switch (e.g. summarize → chat) + self.max_iterations += extra_iters if response.get("switch_mode"): self.mode = response["switch_mode"] - self.system_prompt = get_system_prompt(self.mode, workspace_path=str(self.reports_folder), backend=self.backend) - self.tool_schemas = build_tool_schemas(self.mode, is_subagent=self.is_subagent, backend=self.backend) + self._rebuild_prompt_and_tools() self.history.set_system(maybe_encrypt(self.system_prompt, self.encryptor)) - self.max_iterations += extra_iters items.append(Info(message=f"Switched to {self.mode} mode")) else: # Re-detect mode (user may switch from chat to attack, etc.) previous_mode = self.mode self._detect_mode(force=True) - self.max_iterations += extra_iters if self.mode != previous_mode: self.history.set_system(maybe_encrypt(self.system_prompt, self.encryptor)) items.append(Info(message=f"Switched to {self.mode} mode")) @@ -836,5 +1313,8 @@ def _prompt_and_redetect(self, choices): # Token breakdown for prompt display by_role = self.history.count_tokens_by_role(self.model) extra_data = {"tokens": by_role["total"], "context_window": get_context_window(self.model), "by_role": by_role} - items.append(Ai(content=answer, ai_type="prompt", extra_data=extra_data)) + items.append(Ai( + content=answer, ai_type="prompt", extra_data=extra_data, + message={"role": "user", "content": maybe_encrypt(answer, self.encryptor)}, + )) return items diff --git a/secator/tasks/command.py b/secator/tasks/command.py new file mode 100644 index 000000000..ae4e068cd --- /dev/null +++ b/secator/tasks/command.py @@ -0,0 +1,97 @@ +from datetime import datetime, timezone +from time import time + +from secator.decorators import task +from secator.output_types import Error +from secator.runners import Command + + +@task() +class command(Command): + """Run an arbitrary shell command verbatim.""" + cmd = '' + shell = True + input_flag = None + # NOTE: input_types MUST stay empty. A non-empty list makes _validate_inputs() + # autodetect each input's type and drop mismatches (e.g. "whoami" autodetects as + # 'slug', not 'str', so [STRING] would strip most bare commands -> empty cmd -> FAILURE). + input_types = [] + output_types = [] + + def _build_cmd(self): + """Set the command to the raw input verbatim (no flag/opt append, no quoting).""" + self.cmd = self.inputs[0] if self.inputs else '' + self.cmd_options = {} + # Command.__init__ runs _build_cmd_input() BEFORE _build_cmd(), and it clobbers + # self.shell to (' | ' in self.cmd) — i.e. False for most commands. Restore the + # intended shell mode so &&, ;, redirects, $VAR, globbing are interpreted. + self.shell = True + + def is_installed(self): + """Arbitrary shell commands have no fixed binary to `which`/auto-install (the base + Command.is_installed() derives cmd_name from the class-level `cmd`, which is '' here). + Always report installed so the base yielder runs the input verbatim instead of trying + (and failing) to auto-install an empty command name. + """ + return True + + @classmethod + def from_result(cls, command_line, output, return_code, *, start_time=None, end_time=None, context=None, hooks=None): + """Build a `command` runner from an ALREADY-RUN command's result, without executing it. + + This is the "import" path (as opposed to the "execute" path exercised by + `run()`/`yielder()`): it never spawns a subprocess, it just populates the runner's + state fields from a result that was captured elsewhere, then fires the same + `on_start`/`on_end` hooks a normal run would fire so the imported command persists + like any other runner (e.g. via an `update_runner` hook passed in `hooks`). This is + the forward-looking seam for importing externally-run commands into Secator Cloud. + + Args: + command_line (str): The command line that was run, verbatim. It becomes `self.cmd` + via the constructor -> `_build_cmd()`, same as the live-execution path (with + `input_types = []`, inputs are never type-filtered, so this holds for every + command line, including bare single-word ones like "whoami"). + output (str): Captured stdout of the already-run command. + return_code (int): Process return code of the already-run command. 0 means + success; anything else marks the runner FAILURE (an `Error` result is added + so `self_errors`, which `status` derives from, is non-empty). + start_time (datetime, optional): When the command started (tz-aware). Defaults + to now if omitted. + end_time (datetime, optional): When the command finished (tz-aware). Defaults to + now if omitted. + context (dict, optional): Runner context (workspace, etc), same as the live path. + hooks (dict, optional): Runner hooks (e.g. `on_end: [update_runner]`), same as the + live path -- this is how the imported result gets persisted. + + Returns: + command: the populated runner, in SUCCESS or FAILURE status. `yielder()` / + `run()` are never called, so no subprocess is ever spawned. + """ + runner = cls(inputs=[command_line], context=context or {}, hooks=hooks or {}) + + # mark_started() fires the on_start hook. It also stamps start_time = now(), so + # apply the caller-supplied start_time right after (mark_started() unconditionally + # overwrites it, there's no way to seed it beforehand). + runner.mark_started() + runner.start_time = start_time or datetime.fromtimestamp(time(), timezone.utc) + + # Populate the captured result. + runner.output = output + runner.return_code = return_code + if return_code != 0: + # `status` derives FAILURE from `self_errors` being non-empty; add_result() stamps + # `_source` for `_owns_error()` matching. output=False is REQUIRED -- the default + # would append this synthetic Error's ANSI repr onto the captured stdout, corrupting it. + runner.add_result( + Error(message=f'Command exited with return code {return_code}'), + print=False, + output=False, + ) + + # mark_completed() fires the on_end hook (the persistence path). Same caveat as + # start_time: it unconditionally stamps end_time = now(), so apply the + # caller-supplied end_time right after. + runner.mark_completed() + runner.end_time = end_time or datetime.fromtimestamp(time(), timezone.utc) + + return runner diff --git a/tests/unit/test_ai_actions.py b/tests/unit/test_ai_actions.py index 05ff6c99b..4b16dc5cf 100644 --- a/tests/unit/test_ai_actions.py +++ b/tests/unit/test_ai_actions.py @@ -9,8 +9,14 @@ if ADDONS_ENABLED['ai']: from secator.ai.actions import ( ActionContext, dispatch_action, _handle_follow_up, _handle_shell, - _handle_query, _handle_add_finding, _run_runner, _decrypt_dict + _handle_query, _handle_add_finding, _run_runner, _decrypt_dict, + _build_hooks_from_context, _coerce_finding_fields, _sanitize_child_opts, + _build_child_hooks_or_denial, + _MAX_SUBAGENT_DEPTH, _MAX_SUBAGENTS_PER_TURN, + _MAX_SHELL_OUTPUT_CHARS, _truncate, ) + from secator.ai.utils import _MAX_CHILD_ITERATIONS + from secator.runners import Task from secator.output_types import Ai, Error, Info, Warning, Vulnerability, Url @@ -35,6 +41,14 @@ def test_decrypt_nested_dict(self): self.assertEqual(result['outer']['inner'], 'VALUE') + def test_decrypt_non_dict_returned_unchanged(self): + """Backstop: a non-dict (e.g. a stringified query arg) must not raise + `.items()` — it is returned unchanged instead of crashing the action.""" + encryptor = MagicMock() + self.assertEqual(_decrypt_dict('{"_type": "url"}', encryptor), '{"_type": "url"}') + self.assertEqual(_decrypt_dict(['a', 'b'], encryptor), ['a', 'b']) + encryptor.decrypt.assert_not_called() + def test_decrypt_list_values(self): encryptor = MagicMock() encryptor.decrypt.side_effect = lambda x: x.upper() @@ -119,6 +133,9 @@ def test_follow_up_with_choices(self): self.assertEqual(results[0].ai_type, 'follow_up') self.assertEqual(results[0].content, 'What next?') self.assertEqual(results[0].extra_data['choices'], ['Scan deeper', 'Try SQL injection']) + # Choices must also land on the top-level `choices` field (what the web UI reads), + # not only in extra_data — otherwise the persisted follow-up doc renders no buttons. + self.assertEqual(results[0].choices, ['Scan deeper', 'Try SQL injection']) @unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') @@ -134,52 +151,163 @@ def test_shell_dry_run(self): self.assertIn('DRY RUN', results[0].message) self.assertIn('whoami', results[0].message) - @patch('secator.ai.actions.subprocess.run') - def test_shell_execution(self, mock_run): - mock_run.return_value = MagicMock(stdout='root\n', stderr='') + def test_shell_execution(self): + """Real integration test: `_handle_shell` dispatches the actual `command` task + (no driver -> no persistence, but it still runs) and surfaces its stdout.""" ctx = ActionContext(targets=['t.com'], model='m') - results = list(_handle_shell({'action': 'shell', 'command': 'whoami'}, ctx)) + results = list(_handle_shell({'action': 'shell', 'command': 'echo hello'}, ctx)) self.assertEqual(len(results), 2) # First: the command being run self.assertIsInstance(results[0], Ai) self.assertEqual(results[0].ai_type, 'shell') - self.assertEqual(results[0].content, 'whoami') + self.assertEqual(results[0].content, 'echo hello') # Second: the output self.assertIsInstance(results[1], Ai) self.assertEqual(results[1].ai_type, 'shell_output') - self.assertIn('root', results[1].content) + self.assertIn('hello', results[1].content) + + @patch('secator.tasks.command.command') + def test_shell_dispatches_command_task_with_parent_context(self, mock_task_cls): + """The shell command must run as a `command` task carrying the parent's context + (notably `session_id`), so it persists nested under the conversation.""" + fake_runner = MagicMock() + fake_runner.id = 'task_abc' + fake_runner.output = 'hello' + mock_task_cls.return_value = fake_runner + + ctx = ActionContext(targets=['t.com'], model='m', session_id='sess-123') + results = list(_handle_shell({'action': 'shell', 'command': 'echo hello'}, ctx)) + + # CommandTask constructed with the raw command line as its sole input (first + # positional arg, since we instantiate the concrete class directly). + call_args = mock_task_cls.call_args + self.assertEqual(call_args[0][0], ['echo hello']) + captured_context = call_args[1]['context'] + self.assertEqual(captured_context.get('session_id'), 'sess-123') + + fake_runner.run.assert_called_once() + self.assertEqual(fake_runner.max_timeout, 60) - @patch('secator.ai.actions.subprocess.run') - def test_shell_stderr(self, mock_run): - mock_run.return_value = MagicMock(stdout='', stderr='error msg') - ctx = ActionContext(targets=['t.com'], model='m') + self.assertEqual(results[0].ai_type, 'shell') + # The shell Ai must carry the UI-linking extra_data (runner_type + runner_id) + # so a future regression in that wiring is caught. + self.assertEqual(results[0].extra_data.get('runner_type'), 'task') + self.assertTrue(results[0].extra_data.get('runner_id')) + self.assertEqual(results[-1].ai_type, 'shell_output') + self.assertIn('hello', results[-1].content) + + @patch('secator.tasks.command.command') + def test_shell_passes_sanitized_env(self, mock_task_cls): + """SECURITY: the sanitized process env is passed via the `env` run_opt so an + AI-run `env`/`printenv` can't dump the LLM key / cloud creds into output that + reaches the LLM + Mongo.""" + import os as _os + fake_runner = MagicMock() + fake_runner.id = 'task_abc' + fake_runner.output = '' + mock_task_cls.return_value = fake_runner + + with patch.dict(_os.environ, {'ANTHROPIC_API_KEY': 'sk-secret', 'HOME': '/home/x'}): + ctx = ActionContext(targets=['t.com'], model='m') + list(_handle_shell({'action': 'shell', 'command': 'env'}, ctx)) + + # run_opts are spread as kwargs (command takes **run_opts), so `env` is a + # top-level kwarg, not nested under a `run_opts` kwarg. + passed_env = mock_task_cls.call_args[1]['env'] + self.assertNotIn('ANTHROPIC_API_KEY', passed_env) + # a benign var still passes through so the command can actually run + self.assertEqual(passed_env.get('HOME'), '/home/x') + + @patch('secator.ai.actions._build_child_hooks_or_denial') + def test_shell_extracts_task_hooks_and_they_fire(self, mock_hooks): + """REGRESSION GUARD (live-Mongo bug): _build_child_hooks_or_denial returns a + CLASS-keyed dict ({Task: {on_end: [...]}, ...}). Direct `command` instantiation + bypasses the Task wrapper that would extract the Task sub-dict, so _handle_shell + must extract `hooks.get(Task, {})` itself — otherwise register_hooks resolves + nothing and the mongodb persistence hooks never fire (command runs SUCCESS but + no doc is persisted). This asserts the extraction happened AND the hook actually + fired during the real run — a flat-dict mock cannot catch the extraction bug. + """ + fired = [] + + def sentinel(runner): + fired.append(runner) + return runner # on_end convention: return the runner + + mock_hooks.return_value = ({Task: {'on_end': [sentinel]}}, None) - results = list(_handle_shell({'action': 'shell', 'command': 'bad'}, ctx)) + ctx = ActionContext(targets=['t.com'], model='m') + results = list(_handle_shell({'action': 'shell', 'command': 'echo hi'}, ctx)) - self.assertEqual(results[1].content, 'error msg') + # The class-keyed dict was extracted to the Task sub-dict, so the on_end hook + # resolved and fired against the real command runner. + self.assertTrue(fired, "on_end hook did not fire — Task sub-dict was not extracted") + self.assertEqual(results[-1].ai_type, 'shell_output') + self.assertIn('hi', results[-1].content) - @patch('secator.ai.actions.subprocess.run') - def test_shell_no_output(self, mock_run): - mock_run.return_value = MagicMock(stdout='', stderr='') + def test_shell_no_output(self): ctx = ActionContext(targets=['t.com'], model='m') results = list(_handle_shell({'action': 'shell', 'command': 'true'}, ctx)) self.assertEqual(results[1].content, '(no output)') - @patch('secator.ai.actions.subprocess.run') - def test_shell_exception(self, mock_run): - mock_run.side_effect = Exception('Command timed out') + @patch('secator.tasks.command.command') + def test_shell_exception(self, mock_task_cls): + mock_task_cls.side_effect = Exception('Command timed out') ctx = ActionContext(targets=['t.com'], model='m') results = list(_handle_shell({'action': 'shell', 'command': 'slow'}, ctx)) - # shell Ai + Error - self.assertEqual(len(results), 2) - self.assertIsInstance(results[1], Error) - self.assertIn('failed', results[1].message) + # Only the Error (the shell Ai is emitted AFTER construction, which never + # succeeds here). + self.assertEqual(len(results), 1) + self.assertIsInstance(results[0], Error) + self.assertIn('failed', results[0].message) + + @patch('secator.tasks.command.command') + def test_shell_output_capped_when_over_limit(self, mock_task_cls): + # M1: huge stdout must be truncated to <= cap + marker and carry the marker. + big = "HEAD_LINE\n" + ("x" * (_MAX_SHELL_OUTPUT_CHARS * 3)) + "\nTAIL_LINE" + fake_runner = MagicMock() + fake_runner.id = 'task_abc' + fake_runner.output = big + mock_task_cls.return_value = fake_runner + ctx = ActionContext(targets=['t.com'], model='m') + + results = list(_handle_shell({'action': 'shell', 'command': 'dump'}, ctx)) + + content = results[1].content + self.assertLess(len(content), len(big)) + # body is bounded by the cap (plus the short marker line) + self.assertLessEqual(len(content), _MAX_SHELL_OUTPUT_CHARS + 40) + self.assertIn('truncated', content) + # head + tail preserved so the model sees the start AND the final lines + self.assertIn('HEAD_LINE', content) + self.assertIn('TAIL_LINE', content) + + def test_shell_output_short_passes_through_unchanged(self): + # M1: short output must pass through untouched (no marker). The `command` + # runner rstrips its captured output, so a trailing newline is not expected. + ctx = ActionContext(targets=['t.com'], model='m') + + results = list(_handle_shell({'action': 'shell', 'command': 'echo root'}, ctx)) + + self.assertEqual(results[1].content, 'root') + self.assertNotIn('truncated', results[1].content) + + def test_truncate_short_text_unchanged(self): + self.assertEqual(_truncate('short', 100), 'short') + + def test_truncate_keeps_head_and_tail(self): + text = 'START' + ('m' * 500) + 'END' + out = _truncate(text, 100) + self.assertLessEqual(len(out), 100 + 40) + self.assertTrue(out.startswith('START')) + self.assertTrue(out.endswith('END')) + self.assertIn('truncated', out) def test_shell_decrypts_command(self): encryptor = MagicMock() @@ -197,8 +325,14 @@ class TestHandleQuery(unittest.TestCase): """Tests for the _handle_query action handler.""" def test_query_no_workspace(self): + """A NON-local backend (mongodb/api) without a workspace_id yields the + 'No workspace' guard. The local driver is exempt (it answers from in-memory + results — see test_query_local_driver_exempt_from_workspace_guard).""" + mock_engine = MagicMock() + mock_engine.backend.name = "mongodb" ctx = ActionContext(targets=['t.com'], model='m', context={}) - results = list(_handle_query({'action': 'query', 'query': {}}, ctx)) + with patch.object(ctx, 'get_query_engine', return_value=mock_engine): + results = list(_handle_query({'action': 'query', 'query': {}}, ctx)) self.assertEqual(len(results), 1) self.assertIsInstance(results[0], Warning) @@ -244,6 +378,41 @@ def test_query_success(self, mock_get_engine): for r in result_dicts: self.assertTrue(r['_context'].get('ai_query_result')) + @patch('secator.ai.actions.ActionContext.get_query_engine') + def test_query_stringified_json_is_coerced(self, mock_get_engine): + """A model that passes `query` as a JSON *string* (schema says object) must + still work — coerced to a dict, then searched. Regression for the + AttributeError('str' object has no attribute 'items') in _decrypt_dict.""" + mock_engine = MagicMock() + mock_engine.search.return_value = [{'_type': 'url', '_context': {}}] + mock_get_engine.return_value = mock_engine + # Encryptor active is the exact condition that made the original crash fire. + encryptor = MagicMock() + encryptor.decrypt.side_effect = lambda s: s + ctx = ActionContext(targets=['t.com'], model='m', context={'workspace_id': 'ws1'}, encryptor=encryptor) + + results = list(_handle_query( + {'action': 'query', 'query': '{"_type": "url", "verified": true}'}, ctx)) + + self.assertFalse([r for r in results if isinstance(r, Error)], 'stringified query must not error') + mock_engine.search.assert_called_once_with({'_type': 'url', 'verified': True}, limit=100) + + def test_query_unparseable_string_returns_clean_error(self): + """A non-JSON string yields an Error the LLM can act on — not a crash.""" + ctx = ActionContext(targets=['t.com'], model='m', context={'workspace_id': 'ws1'}) + results = list(_handle_query({'action': 'query', 'query': 'not json at all'}, ctx)) + errors = [r for r in results if isinstance(r, Error)] + self.assertEqual(len(errors), 1) + self.assertIn('JSON object', errors[0].message) + + def test_query_non_dict_returns_clean_error(self): + """A non-dict, non-str query (e.g. a list) yields a clean Error, not a crash.""" + ctx = ActionContext(targets=['t.com'], model='m', context={'workspace_id': 'ws1'}) + results = list(_handle_query({'action': 'query', 'query': ['_type', 'url']}, ctx)) + errors = [r for r in results if isinstance(r, Error)] + self.assertEqual(len(errors), 1) + self.assertIn('JSON object', errors[0].message) + @patch('secator.ai.actions.ActionContext.get_query_engine') def test_query_failure(self, mock_get_engine): mock_engine = MagicMock() @@ -273,6 +442,82 @@ def test_query_decrypts_filter(self): call_args = mock_engine.search.call_args[0][0] self.assertEqual(call_args['host'], 'example.com') + def test_union_live_results_dedup_and_filter(self): + """_union_live_results filters live by the query, merges into backend results, + and dedupes by _uuid (backend wins).""" + from secator.ai.actions import _union_live_results + persisted = [{"_uuid": "a", "_type": "port"}] + live = [{"_uuid": "a", "_type": "port"}, # dup -> deduped + {"_uuid": "b", "_type": "port"}, # new -> included + {"_uuid": "c", "_type": "url"}] # wrong type -> filtered out by the query + out = _union_live_results(list(persisted), live, {"_type": "port"}, 100) + self.assertEqual(sorted(r["_uuid"] for r in out), ["a", "b"]) + # no live results -> persisted returned unchanged + self.assertEqual(_union_live_results([{"_uuid": "z"}], [], {}, 0), [{"_uuid": "z"}]) + + def test_query_local_driver_unions_live_results(self): + """Local (json) driver: query_workspace unions this run's in-memory findings + with the backend (JSON exporter only writes to disk at end-of-run).""" + mock_engine = MagicMock() + mock_engine.backend.name = "json" + mock_engine.search.return_value = [{"_uuid": "disk1", "_type": "port", "_context": {}}] + ctx = ActionContext(targets=['t'], model='m', context={'workspace_id': 'ws1'}, + results=[{"_uuid": "live1", "_type": "port", "_context": {}}]) + with patch.object(ctx, 'get_query_engine', return_value=mock_engine): + results = list(_handle_query({'action': 'query', 'query': {'_type': 'port'}}, ctx)) + uuids = {r.get('_uuid') for r in results if isinstance(r, dict)} + self.assertIn('disk1', uuids) # backend result + self.assertIn('live1', uuids) # unioned live in-memory result + + def test_query_mongodb_driver_does_not_union(self): + """Non-local backend (mongodb) is queried normally — live self.results are NOT unioned.""" + mock_engine = MagicMock() + mock_engine.backend.name = "mongodb" + mock_engine.search.return_value = [{"_uuid": "db1", "_type": "port", "_context": {}}] + ctx = ActionContext(targets=['t'], model='m', context={'workspace_id': 'ws1'}, + results=[{"_uuid": "live1", "_type": "port", "_context": {}}]) + with patch.object(ctx, 'get_query_engine', return_value=mock_engine): + results = list(_handle_query({'action': 'query', 'query': {'_type': 'port'}}, ctx)) + uuids = {r.get('_uuid') for r in results if isinstance(r, dict)} + self.assertIn('db1', uuids) + self.assertNotIn('live1', uuids) # not unioned for non-local backends + + def test_query_local_driver_exempt_from_workspace_guard(self): + """Local driver with NO workspace_id is not blocked by the 'No workspace' guard — + it answers from in-memory results.""" + mock_engine = MagicMock() + mock_engine.backend.name = "json" + mock_engine.search.return_value = [] + ctx = ActionContext(targets=['t'], model='m', context={}, # no workspace_id + results=[{"_uuid": "live1", "_type": "port", "_context": {}}]) + with patch.object(ctx, 'get_query_engine', return_value=mock_engine): + results = list(_handle_query({'action': 'query', 'query': {'_type': 'port'}}, ctx)) + warnings = [r for r in results if isinstance(r, Warning)] + self.assertFalse(any('No workspace' in getattr(w, 'message', '') for w in warnings)) + uuids = {r.get('_uuid') for r in results if isinstance(r, dict)} + self.assertIn('live1', uuids) + + def test_query_stringified_limit_coerced_to_int(self): + """A model-supplied string limit ('10') is coerced to int before the backend + (a str limit raises TypeError: '>=' not supported between int and str).""" + mock_engine = MagicMock() + mock_engine.backend.name = "mongodb" + mock_engine.search.return_value = [] + ctx = ActionContext(targets=['t'], model='m', context={'workspace_id': 'ws1'}) + with patch.object(ctx, 'get_query_engine', return_value=mock_engine): + list(_handle_query({'action': 'query', 'query': {'_type': 'ip'}, 'limit': '10'}, ctx)) + self.assertEqual(mock_engine.search.call_args.kwargs.get('limit'), 10) + + def test_query_bad_limit_falls_back_to_default(self): + """A non-numeric limit falls back to the default (100), not a crash.""" + mock_engine = MagicMock() + mock_engine.backend.name = "mongodb" + mock_engine.search.return_value = [] + ctx = ActionContext(targets=['t'], model='m', context={'workspace_id': 'ws1'}) + with patch.object(ctx, 'get_query_engine', return_value=mock_engine): + list(_handle_query({'action': 'query', 'query': {}, 'limit': 'notanumber'}, ctx)) + self.assertEqual(mock_engine.search.call_args.kwargs.get('limit'), 100) + @unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') class TestRunRunner(unittest.TestCase): @@ -316,6 +561,363 @@ def test_run_runner_uses_ctx_targets_as_default(self): self.assertIn('default.com', results[0].message) + @patch('secator.ai.actions.TemplateLoader') + @patch('secator.ai.actions.Task') + @patch('secator.ai.actions._build_hooks_from_context') + def test_run_runner_propagates_hooks_and_emits_runner_id(self, mock_build_hooks, mock_task_cls, _mock_tpl): + """Sub-runner must receive driver hooks (so its results persist) and the + emitted action Ai must carry the created runner's id + type for the UI.""" + sentinel_hooks = {'fake': ['hook']} + mock_build_hooks.return_value = sentinel_hooks + + # Fake runner: an iterable whose id is populated (mimics on_init stamping it) + mock_runner = MagicMock() + mock_runner.id = 'runner123' + mock_runner.reports_folder = None + mock_runner.__iter__.return_value = iter([]) + mock_task_cls.return_value = mock_runner + + ctx = ActionContext( + targets=['t.com'], model='m', + context={'workspace_id': 'ws1', 'drivers': ['mongodb']}, + ) + action = {'action': 'task', 'name': 'nmap', 'targets': ['10.0.0.1']} + + results = list(_run_runner(action, ctx, 'task')) + + # Runner constructed with hooks= from the context drivers + _, kwargs = mock_task_cls.call_args + self.assertEqual(kwargs.get('hooks'), sentinel_hooks) + self.assertEqual(kwargs.get('context', {}).get('workspace_id'), 'ws1') + + # Action Ai item carries runner_id + runner_type + ai_items = [r for r in results if isinstance(r, Ai) and r.ai_type == 'task'] + self.assertEqual(len(ai_items), 1) + self.assertEqual(ai_items[0].extra_data.get('runner_id'), 'runner123') + self.assertEqual(ai_items[0].extra_data.get('runner_type'), 'task') + + @patch('secator.ai.actions.TemplateLoader') + @patch('secator.ai.actions.Task') + @patch('secator.ai.actions._build_hooks_from_context') + def test_run_runner_propagates_session_id(self, mock_build_hooks, mock_task_cls, _mock_tpl): + """The dispatched sub-runner's context must carry the ai task's session_id + (the conversation id) so its persisted runner doc is queryable by the + conversation. session_id may be derived (not already in ctx.context), so + it must be stamped from ctx.session_id.""" + # non-empty: context has drivers, so empty hooks would trip the M2 guard + mock_build_hooks.return_value = {'fake': ['hook']} + mock_runner = MagicMock() + mock_runner.id = 'runner123' + mock_runner.reports_folder = None + mock_runner.__iter__.return_value = iter([]) + mock_task_cls.return_value = mock_runner + + # session_id lives on the ActionContext but NOT in context (it is derived) + ctx = ActionContext( + targets=['t.com'], model='m', + context={'workspace_id': 'ws1', 'drivers': ['mongodb']}, + session_id='conv-abc-123', + ) + action = {'action': 'task', 'name': 'nmap', 'targets': ['10.0.0.1']} + + list(_run_runner(action, ctx, 'task')) + + _, kwargs = mock_task_cls.call_args + sub_context = kwargs.get('context', {}) + self.assertEqual(sub_context.get('session_id'), 'conv-abc-123') + self.assertEqual(sub_context.get('workspace_id'), 'ws1') + + @patch('secator.ai.actions.TemplateLoader') + @patch('secator.ai.actions.Task') + @patch('secator.ai.actions._build_hooks_from_context') + def test_run_runner_preserves_existing_session_id(self, mock_build_hooks, mock_task_cls, _mock_tpl): + """A session_id already present in ctx.context must not be overwritten.""" + mock_build_hooks.return_value = {} + mock_runner = MagicMock() + mock_runner.id = 'runner123' + mock_runner.reports_folder = None + mock_runner.__iter__.return_value = iter([]) + mock_task_cls.return_value = mock_runner + + ctx = ActionContext( + targets=['t.com'], model='m', + context={'workspace_id': 'ws1', 'session_id': 'from-context'}, + session_id='from-ctx-field', + ) + action = {'action': 'task', 'name': 'nmap', 'targets': ['10.0.0.1']} + + list(_run_runner(action, ctx, 'task')) + + _, kwargs = mock_task_cls.call_args + self.assertEqual(kwargs.get('context', {}).get('session_id'), 'from-context') + + @patch('secator.ai.actions.TemplateLoader') + @patch('secator.ai.actions.Task') + @patch('secator.ai.actions._build_hooks_from_context') + def test_run_runner_structures_subagent_prompt(self, mock_build_hooks, mock_task_cls, _tpl): + mock_build_hooks.return_value = {'fake': ['hook']} + mock_runner = MagicMock(); mock_runner.id = 'r1'; mock_runner.reports_folder = None + mock_runner.__iter__.return_value = iter([]); mock_task_cls.return_value = mock_runner + ctx = ActionContext(targets=['10.0.0.1'], model='m', + context={'workspace_id': 'ws1', 'drivers': ['mongodb']}) + with patch('secator.ai.actions._gather_subagent_evidence', return_value="- port 10.0.0.1:443"): + action = {'action': 'task', 'name': 'ai', 'targets': ['10.0.0.1'], + 'opts': {'prompt': 'Test auth on the API'}} + list(_run_runner(action, ctx, 'task')) + _, kwargs = mock_task_cls.call_args + prompt = kwargs.get('run_opts', {}).get('prompt', '') + self.assertIn('## Objective', prompt) + self.assertIn('Test auth on the API', prompt) + self.assertIn('- port 10.0.0.1:443', prompt) # evidence injected + + @patch('secator.ai.actions.TemplateLoader') + @patch('secator.ai.actions.Task') + @patch('secator.ai.actions._build_hooks_from_context') + def test_run_runner_subagent_inherits_parent_llm_config(self, mock_build_hooks, mock_task_cls, _tpl): + """A spawned subagent inherits the parent's resolved model/api_key/api_base so it + can actually run (else it falls back to CONFIG.default_model with no key).""" + mock_build_hooks.return_value = {'fake': ['hook']} + mock_runner = MagicMock(); mock_runner.id = 'r1'; mock_runner.reports_folder = None + mock_runner.__iter__.return_value = iter([]); mock_task_cls.return_value = mock_runner + ctx = ActionContext(targets=['10.0.0.1'], model='openrouter/anthropic/x', + api_key='PARENTKEY', api_base='https://base', + context={'workspace_id': 'ws1', 'drivers': ['mongodb']}) + with patch('secator.ai.actions._gather_subagent_evidence', return_value=""): + action = {'action': 'task', 'name': 'ai', 'targets': ['10.0.0.1'], 'opts': {'prompt': 'do x'}} + list(_run_runner(action, ctx, 'task')) + ro = mock_task_cls.call_args[1].get('run_opts', {}) + self.assertEqual(ro.get('model'), 'openrouter/anthropic/x') + self.assertEqual(ro.get('api_key'), 'PARENTKEY') + self.assertEqual(ro.get('api_base'), 'https://base') + + @patch('secator.ai.actions.TemplateLoader') + @patch('secator.ai.actions.Task') + @patch('secator.ai.actions._build_hooks_from_context') + def test_run_runner_subagent_explicit_model_wins(self, mock_build_hooks, mock_task_cls, _tpl): + """An explicit LLM-supplied model on the subagent opts is preserved (setdefault).""" + mock_build_hooks.return_value = {'fake': ['hook']} + mock_runner = MagicMock(); mock_runner.id = 'r1'; mock_runner.reports_folder = None + mock_runner.__iter__.return_value = iter([]); mock_task_cls.return_value = mock_runner + ctx = ActionContext(targets=['t'], model='parent/model', + context={'workspace_id': 'ws1', 'drivers': ['mongodb']}) + with patch('secator.ai.actions._gather_subagent_evidence', return_value=""): + action = {'action': 'task', 'name': 'ai', 'targets': ['t'], + 'opts': {'prompt': 'x', 'model': 'explicit/model'}} + list(_run_runner(action, ctx, 'task')) + ro = mock_task_cls.call_args[1].get('run_opts', {}) + self.assertEqual(ro.get('model'), 'explicit/model') + + +@unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') +class TestSanitizeChildOpts(unittest.TestCase): + """Tests for _sanitize_child_opts (C1: LLM-supplied subagent opts allow-list).""" + + def test_strips_dangerous_and_control_keys(self): + opts = { + 'dangerous': True, 'interactive': 'local', 'hooks': {'x': 1}, + 'sync': False, 'subagent': True, 'tty': True, 'dry_run': True, + 'exporters': ['csv'], 'enable_reports': False, + } + clean = _sanitize_child_opts(opts) + self.assertEqual(clean, {}) + + def test_strips_print_star_keys(self): + clean = _sanitize_child_opts({'print_cmd': True, 'print_item': False, 'print_anything': 1}) + self.assertEqual(clean, {}) + + def test_keeps_benign_task_opts(self): + clean = _sanitize_child_opts({'ports': '80,443', 'rate_limit': 100, 'mode': 'attack'}) + self.assertEqual(clean, {'ports': '80,443', 'rate_limit': 100, 'mode': 'attack'}) + + def test_clamps_max_iterations(self): + clean = _sanitize_child_opts({'max_iterations': 9999}) + self.assertEqual(clean['max_iterations'], _MAX_CHILD_ITERATIONS) + clean = _sanitize_child_opts({'max_iterations': 5}) + self.assertEqual(clean['max_iterations'], 5) + # bool / non-numeric max_iterations is dropped (bool is an int subclass) + self.assertNotIn('max_iterations', _sanitize_child_opts({'max_iterations': True})) + self.assertNotIn('max_iterations', _sanitize_child_opts({'max_iterations': 'lots'})) + + def test_non_dict_returns_empty(self): + self.assertEqual(_sanitize_child_opts(None), {}) + self.assertEqual(_sanitize_child_opts('dangerous'), {}) + + @patch('secator.ai.actions.TemplateLoader') + @patch('secator.ai.actions.Task') + @patch('secator.ai.actions._build_hooks_from_context') + def test_run_runner_neutralizes_dangerous_and_interactive(self, mock_build_hooks, mock_task_cls, _mock_tpl): + """C1: an LLM emitting opts={'dangerous': True, 'interactive': 'local'} must NOT + propagate either into the spawned child's run_opts — dangerous is forced False + and interactive (a control key) is stripped.""" + mock_build_hooks.return_value = {} + mock_runner = MagicMock() + mock_runner.id = 'runner123' + mock_runner.reports_folder = None + mock_runner.__iter__.return_value = iter([]) + mock_task_cls.return_value = mock_runner + + ctx = ActionContext(targets=['t.com'], model='m', context={'workspace_id': 'ws1'}) + action = { + 'action': 'task', 'name': 'nmap', 'targets': ['10.0.0.1'], + 'opts': {'dangerous': True, 'interactive': 'local', 'ports': '80'}, + } + + list(_run_runner(action, ctx, 'task')) + + _, kwargs = mock_task_cls.call_args + run_opts = kwargs.get('run_opts', {}) + self.assertEqual(run_opts.get('dangerous'), False) + self.assertNotEqual(run_opts.get('interactive'), 'local') + # benign task opt still passes through + self.assertEqual(run_opts.get('ports'), '80') + + @patch('secator.ai.actions.TemplateLoader') + @patch('secator.ai.actions.Task') + @patch('secator.ai.actions._build_hooks_from_context') + def test_run_runner_ai_subagent_forced_flags_over_llm_opts(self, mock_build_hooks, mock_task_cls, _mock_tpl): + """Spawning an `ai` subagent with hostile opts: subagent forced True, + interactive forced False, dangerous forced False regardless of LLM input.""" + mock_build_hooks.return_value = {} + mock_runner = MagicMock() + mock_runner.id = 'runner123' + mock_runner.reports_folder = None + mock_runner.__iter__.return_value = iter([]) + mock_task_cls.return_value = mock_runner + + ctx = ActionContext(targets=['t.com'], model='m', context={'workspace_id': 'ws1'}) + action = { + 'action': 'task', 'name': 'ai', 'targets': ['10.0.0.1'], + 'opts': {'dangerous': True, 'interactive': 'local', 'subagent': False}, + } + + list(_run_runner(action, ctx, 'task')) + + _, kwargs = mock_task_cls.call_args + run_opts = kwargs.get('run_opts', {}) + self.assertEqual(run_opts.get('dangerous'), False) + self.assertEqual(run_opts.get('interactive'), False) + self.assertEqual(run_opts.get('subagent'), True) + + +@unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') +class TestBuildHooksFromContext(unittest.TestCase): + """Tests for _build_hooks_from_context (driver name -> hooks dict).""" + + def test_no_drivers_returns_empty(self): + self.assertEqual(_build_hooks_from_context({}), {}) + self.assertEqual(_build_hooks_from_context({'drivers': []}), {}) + + @patch('secator.loader.get_available_drivers') + @patch('secator.loader.order_drivers') + @patch('secator.loader.discover_external_drivers') + @patch('secator.utils.import_dynamic') + def test_builds_hooks_from_driver_names(self, mock_import, _disc, mock_order, mock_avail): + from secator.runners import Task + mock_order.side_effect = lambda d: d + mock_avail.return_value = ['mongodb', 'api'] + mongo_hooks = {Task: {'on_init': ['update_runner']}} + mock_import.return_value = mongo_hooks + + hooks = _build_hooks_from_context({'drivers': ['mongodb']}) + + mock_import.assert_called_once_with('secator.hooks.mongodb', 'HOOKS') + self.assertIn(Task, hooks) + self.assertIn('on_init', hooks[Task]) + + @patch('secator.loader.get_available_drivers') + @patch('secator.loader.order_drivers') + @patch('secator.loader.discover_external_drivers') + @patch('secator.utils.import_dynamic') + def test_skips_unsupported_driver(self, mock_import, _disc, mock_order, mock_avail): + mock_order.side_effect = lambda d: d + mock_avail.return_value = ['mongodb'] + hooks = _build_hooks_from_context({'drivers': ['bogus']}) + self.assertEqual(hooks, {}) + mock_import.assert_not_called() + + +@unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') +class TestChildHooksOrDenial(unittest.TestCase): + """M2: refuse to spawn a persistence-less child when the parent has drivers.""" + + @patch('secator.ai.actions._build_hooks_from_context') + def test_parent_no_drivers_empty_hooks_allowed(self, mock_build): + # pure local/no-persistence run: empty hooks child is expected, no denial + mock_build.return_value = {} + hooks, denial = _build_child_hooks_or_denial({'workspace_id': 'ws1'}) + self.assertEqual(hooks, {}) + self.assertIsNone(denial) + + @patch('secator.ai.actions._build_hooks_from_context') + def test_parent_drivers_present_hooks_pass_through(self, mock_build): + # normal spawn: drivers present + non-empty hooks -> pass through unchanged + sentinel = {'fake': ['hook']} + mock_build.return_value = sentinel + hooks, denial = _build_child_hooks_or_denial({'drivers': ['mongodb']}) + self.assertEqual(hooks, sentinel) + self.assertIsNone(denial) + + @patch('secator.ai.actions._build_hooks_from_context') + def test_parent_drivers_but_empty_hooks_denied(self, mock_build): + # parent HAS drivers but rebuild produced no hooks -> refuse, surface Warning + mock_build.return_value = {} + hooks, denial = _build_child_hooks_or_denial({'drivers': ['mongodb']}) + self.assertEqual(hooks, {}) + self.assertIsInstance(denial, Warning) + self.assertIn('drop findings', denial.message) + + @patch('secator.ai.actions._build_hooks_from_context') + def test_rebuild_raise_with_drivers_denied_not_swallowed(self, mock_build): + # a raising rebuild must not degrade to hooks={} silently -> Warning + mock_build.side_effect = RuntimeError('boom') + hooks, denial = _build_child_hooks_or_denial({'drivers': ['mongodb']}) + self.assertEqual(hooks, {}) + self.assertIsInstance(denial, Warning) + self.assertIn('rebuild failed', denial.message) + + @patch('secator.ai.actions._build_hooks_from_context') + def test_rebuild_raise_no_drivers_allowed(self, mock_build): + # no parent drivers: a rebuild error still yields an allowed empty-hooks child + mock_build.side_effect = RuntimeError('boom') + hooks, denial = _build_child_hooks_or_denial({'workspace_id': 'ws1'}) + self.assertEqual(hooks, {}) + self.assertIsNone(denial) + + @patch('secator.ai.actions.TemplateLoader') + @patch('secator.ai.actions.Task') + @patch('secator.ai.actions._build_hooks_from_context') + def test_run_runner_refuses_spawn_on_lost_persistence(self, mock_build, mock_task_cls, _tpl): + # end-to-end: parent has drivers, rebuild empty -> _run_runner yields a + # Warning and never constructs the child runner + mock_build.return_value = {} + ctx = ActionContext( + targets=['t.com'], model='m', + context={'workspace_id': 'ws1', 'drivers': ['mongodb']}, + ) + action = {'action': 'task', 'name': 'nmap', 'targets': ['10.0.0.1']} + results = list(_run_runner(action, ctx, 'task')) + mock_task_cls.assert_not_called() + warnings = [r for r in results if isinstance(r, Warning)] + self.assertEqual(len(warnings), 1) + self.assertIn('denied', warnings[0].message) + + @patch('secator.ai.actions.TemplateLoader') + @patch('secator.ai.actions.Task') + @patch('secator.ai.actions._build_hooks_from_context') + def test_run_runner_no_drivers_spawns_normally(self, mock_build, mock_task_cls, _tpl): + # parent has NO drivers: empty-hooks child still spawns (no false alarm) + mock_build.return_value = {} + mock_runner = MagicMock() + mock_runner.id = 'runner123' + mock_runner.reports_folder = None + mock_runner.__iter__.return_value = iter([]) + mock_task_cls.return_value = mock_runner + ctx = ActionContext(targets=['t.com'], model='m', context={'workspace_id': 'ws1'}) + action = {'action': 'task', 'name': 'nmap', 'targets': ['10.0.0.1']} + results = list(_run_runner(action, ctx, 'task')) + mock_task_cls.assert_called_once() + self.assertFalse([r for r in results if isinstance(r, Warning)]) + @unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') class TestGetQueryEngine(unittest.TestCase): @@ -647,6 +1249,77 @@ def test_add_finding_decrypts_values(self): self.assertIsInstance(results[1], Vulnerability) self.assertEqual(results[1].matched_at, 'http://t.com/search') + def test_coerce_finding_fields_scalar_types(self): + # LLMs send wrong-typed scalars (bool as "true", float/int as strings). + # The coercion helper fixes them to the declared field types. + data = _coerce_finding_fields( + Vulnerability, + { + 'name': 'SQL Injection', + 'verified': 'true', + 'cvss_score': '7.5', + 'severity_nb': '3', + }, + ) + self.assertIs(data['verified'], True) + self.assertIsInstance(data['verified'], bool) + self.assertEqual(data['cvss_score'], 7.5) + self.assertIsInstance(data['cvss_score'], float) + self.assertEqual(data['severity_nb'], 3) + self.assertIsInstance(data['severity_nb'], int) + # str fields are left untouched. + self.assertEqual(data['name'], 'SQL Injection') + # Coerced data validates clean. + self.assertEqual(Vulnerability.validate_fields(data), []) + + def test_add_finding_coerces_scalar_types(self): + # End-to-end: wrong-typed scalars flow through the handler and validate + # clean, producing a Vulnerability with the coerced bool/float values. + ctx = ActionContext(targets=['t.com'], model='m') + results = list( + _handle_add_finding( + { + 'action': 'add_finding', + '_type': 'vulnerability', + 'name': 'SQL Injection', + 'matched_at': 'http://t.com/login', + 'verified': 'true', + 'cvss_score': '7.5', + 'severity_nb': '3', + }, + ctx, + ) + ) + + # No validation Error: the sloppy types were coerced before validation. + self.assertEqual(len(results), 2) + vuln = results[1] + self.assertIsInstance(vuln, Vulnerability) + self.assertIs(vuln.verified, True) + self.assertIsInstance(vuln.verified, bool) + self.assertEqual(vuln.cvss_score, 7.5) + self.assertIsInstance(vuln.cvss_score, float) + + def test_add_finding_unparseable_bool_surfaces_error(self): + # An unparseable value must NOT be silently dropped; validation reports it. + ctx = ActionContext(targets=['t.com'], model='m') + results = list( + _handle_add_finding( + { + 'action': 'add_finding', + '_type': 'vulnerability', + 'name': 'SQL Injection', + 'matched_at': 'http://t.com/login', + 'verified': 'maybe', + }, + ctx, + ) + ) + + self.assertEqual(len(results), 1) + self.assertIsInstance(results[0], Error) + self.assertIn('verified', results[0].message) + @unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') class TestRunBatch(unittest.TestCase): @@ -692,6 +1365,216 @@ def test_run_batch_empty_actions(self): self.assertEqual(len(results), 1) self.assertIsInstance(results[0], Warning) + def test_get_action_label_tolerates_stringified_opts(self): + """Regression: a str `opts` (model stringified it) must not crash the + batch label with AttributeError('str' object has no attribute 'get').""" + from secator.ai.actions import _get_action_label + label = _get_action_label( + {"action": "task", "name": "nmap", "targets": ["10.0.0.1"], "opts": '{"session_name": "x"}'}) + self.assertEqual(label, "nmap on 10.0.0.1") # falls back to name-on-target, no crash + + +@unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') +class TestCheckGuardrailsFailClosed(unittest.TestCase): + """H10: prompts exhausted with the decision still 'ask' must fail CLOSED (deny).""" + + def test_unresolved_after_max_rounds_denies(self): + from secator.ai.actions import check_guardrails_sync + # permission engine that never resolves: always 'ask', no shell/target/path layer + res = MagicMock(decision="ask", shell_command="", targets=[], paths=[], reason="needs approval") + engine = MagicMock() + engine.check_action.return_value = res + ctx = ActionContext(targets=['t.com'], model='m') + ctx.permission_engine = engine + denial, _items = check_guardrails_sync({"action": "shell", "command": "x"}, ctx) + self.assertIsNotNone(denial, "exhausted-but-unresolved guardrail must deny, not return None") + self.assertIn("unresolved", denial) + + +class TestApprovalAllowList(unittest.TestCase): + """Ask-loop approval must be an explicit allow-list: only "allow" proceeds.""" + + def _run(self, answer): + from secator.ai.actions import check_guardrails_sync + ask = MagicMock(decision="ask", shell_command="somecmd", targets=[], paths=[], reason="needs approval") + allow = MagicMock(decision="allow", shell_command="", targets=[], paths=[], reason="") + engine = MagicMock() + engine.check_action.side_effect = [ask, allow] + backend = MagicMock() + backend.ask_user.return_value = None if answer is None else {"answer": answer} + ctx = ActionContext(targets=['t.com'], model='m') + ctx.permission_engine = engine + ctx.backend = backend + denial, _items = check_guardrails_sync({"action": "shell", "command": "somecmd"}, ctx) + return denial + + def test_allow_proceeds(self): + self.assertIsNone(self._run("allow")) + + def test_deny_denies(self): + self.assertIsNotNone(self._run("deny")) + + def test_unexpected_answer_denies(self): + # an out-of-vocabulary token must NOT be treated as approval (fail closed) + self.assertIsNotNone(self._run("sure")) + self.assertIsNotNone(self._run("allow_all_typo")) + + def test_none_response_denies(self): + self.assertIsNotNone(self._run(None)) + + +@unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') +class TestSubagentFanoutCap(unittest.TestCase): + """H4: recursion depth + per-turn fan-out caps on AI-subagent spawns.""" + + def _mock_task(self, mock_task_cls): + runner = MagicMock() + runner.id = 'runner123' + runner.reports_folder = None + runner.__iter__.return_value = iter([]) + mock_task_cls.return_value = runner + return runner + + def test_depth_cap_refuses_spawn(self): + """Spawning an AI subagent at/over _MAX_SUBAGENT_DEPTH is denied.""" + ctx = ActionContext( + targets=['t.com'], model='m', + context={'ai_subagent_depth': _MAX_SUBAGENT_DEPTH}, + ) + action = {'action': 'task', 'name': 'ai', 'targets': ['t.com']} + + with patch('secator.ai.actions.Task') as mock_task_cls: + results = list(_run_runner(action, ctx, 'task')) + + mock_task_cls.assert_not_called() # denied before constructing the child + self.assertEqual(len(results), 1) + self.assertIsInstance(results[0], Warning) + self.assertIn('depth cap', results[0].message) + # no Ai task item emitted (spawn refused) + self.assertFalse([r for r in results if isinstance(r, Ai) and r.ai_type == 'task']) + + def test_per_turn_breadth_cap_refuses_spawn(self): + """In a batch, spawning past _MAX_SUBAGENTS_PER_TURN is denied.""" + ctx = ActionContext( + targets=['t.com'], model='m', in_batch=True, + context={'ai_subagent_turn_count': _MAX_SUBAGENTS_PER_TURN}, + ) + action = {'action': 'task', 'name': 'ai', 'targets': ['t.com']} + + with patch('secator.ai.actions.Task') as mock_task_cls: + results = list(_run_runner(action, ctx, 'task')) + + mock_task_cls.assert_not_called() + self.assertEqual(len(results), 1) + self.assertIsInstance(results[0], Warning) + self.assertIn('fan-out cap', results[0].message) + + @patch('secator.ai.actions.TemplateLoader') + @patch('secator.ai.actions.Task') + @patch('secator.ai.actions._build_hooks_from_context') + def test_normal_depth1_spawn_succeeds(self, mock_build_hooks, mock_task_cls, _mock_tpl): + """A first-level AI subagent (depth 0 -> 1) still spawns; child inherits depth+1.""" + mock_build_hooks.return_value = {} + self._mock_task(mock_task_cls) + + ctx = ActionContext(targets=['t.com'], model='m', context={}) # depth 0, not in a batch + action = {'action': 'task', 'name': 'ai', 'targets': ['t.com']} + + results = list(_run_runner(action, ctx, 'task')) + + mock_task_cls.assert_called_once() + ai_items = [r for r in results if isinstance(r, Ai) and r.ai_type == 'task'] + self.assertEqual(len(ai_items), 1) + self.assertFalse([r for r in results if isinstance(r, Warning)]) + # child context carries incremented depth + _, kwargs = mock_task_cls.call_args + self.assertEqual(kwargs.get('context', {}).get('ai_subagent_depth'), 1) + + +@unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') +class TestChildContextParenting(unittest.TestCase): + """A spawned sub-runner must get a CLEAN identity: it inherits the conversation + session_id + drivers but NOT the parent's runner-doc id, so it mints its own doc + (linked to the conversation by session_id), instead of clobbering the parent.""" + + def test_child_context_strips_parent_identity_keeps_session(self): + from secator.ai.actions import _get_result_context, ActionContext + ctx = ActionContext( + targets=['t.com'], model='m', + context={ + 'workspace_id': 'ws1', 'workspace_name': 'w', 'drivers': ['mongodb'], + 'task_id': 'PARENT_AI_ID', # the parent ai task's own doc id + 'session_id': 'conv-1', + }, + session_id='conv-1', + ) + action = {'action': 'task', 'name': 'nmap', 'tool_call_id': 'tc1', 'tool_call_name': 'run_task'} + child = _get_result_context(action, ctx) + # keeps the conversation link + drivers/workspace + self.assertEqual(child['session_id'], 'conv-1') + self.assertEqual(child['drivers'], ['mongodb']) + self.assertEqual(child['workspace_id'], 'ws1') + # marks it a child + self.assertTrue(child.get('has_parent')) + # does NOT inherit the parent's runner-doc identity (would clobber / suppress its own doc) + self.assertNotIn('task_id', child) + self.assertNotIn('workflow_id', child) + self.assertNotIn('scan_id', child) + + +@unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') +class TestBuildSubagentPrompt(unittest.TestCase): + def test_structure_sections_and_objective(self): + from secator.ai.actions import build_subagent_prompt + p = build_subagent_prompt("Test auth on the API", ["10.0.0.1", "app.x.com"], "- Port 443 open") + self.assertIn("## Objective", p) + self.assertIn("Test auth on the API", p) # objective verbatim + self.assertIn("## Scope", p) + self.assertIn("10.0.0.1", p) + self.assertIn("app.x.com", p) + self.assertIn("## Already known", p) + self.assertIn("- Port 443 open", p) # evidence injected + self.assertIn("## Expected output", p) + + def test_empty_evidence_renders_none(self): + from secator.ai.actions import build_subagent_prompt + p = build_subagent_prompt("Do X", ["t.com"], "") + self.assertIn("(none", p.lower()) # explicit "none" marker + + +@unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') +class TestGatherSubagentEvidence(unittest.TestCase): + def test_queries_targets_and_formats(self): + from secator.ai.actions import _gather_subagent_evidence, ActionContext + mock_engine = MagicMock() + mock_engine.search.return_value = [ + {"_type": "port", "ip": "10.0.0.1", "port": 443}, + {"_type": "url", "url": "http://app.x.com/login"}, + ] + ctx = ActionContext(targets=[], model='m', context={'workspace_id': 'ws1'}) + with patch.object(ctx, 'get_query_engine', return_value=mock_engine): + out = _gather_subagent_evidence(ctx, ["10.0.0.1", "app.x.com"], limit=40) + # queried by an $or over the targets + q = mock_engine.search.call_args[0][0] + self.assertIn("$or", q) + # formatted a compact summary + self.assertIn("port", out) + self.assertIn("10.0.0.1", out) + self.assertIn("url", out) + + def test_no_targets_returns_empty(self): + from secator.ai.actions import _gather_subagent_evidence, ActionContext + ctx = ActionContext(targets=[], model='m', context={}) + self.assertEqual(_gather_subagent_evidence(ctx, [], limit=40), "") + + def test_search_error_returns_empty(self): + from secator.ai.actions import _gather_subagent_evidence, ActionContext + mock_engine = MagicMock() + mock_engine.search.side_effect = Exception("boom") + ctx = ActionContext(targets=[], model='m', context={'workspace_id': 'ws1'}) + with patch.object(ctx, 'get_query_engine', return_value=mock_engine): + self.assertEqual(_gather_subagent_evidence(ctx, ["t"], limit=40), "") + if __name__ == '__main__': unittest.main() diff --git a/tests/unit/test_ai_guardrails.py b/tests/unit/test_ai_guardrails.py index a73f9559f..8391b7479 100644 --- a/tests/unit/test_ai_guardrails.py +++ b/tests/unit/test_ai_guardrails.py @@ -1,6 +1,6 @@ # tests/unit/test_ai_guardrails.py import unittest -from unittest.mock import patch +from unittest.mock import MagicMock, patch from secator.definitions import ADDONS_ENABLED @@ -10,7 +10,8 @@ from secator.ai.guardrails import ( parse_rule, match_rule, extract_command_targets, detect_paths, detect_paths_with_access, detect_sensitive_env_vars, classify_command, build_target_choices, PermissionEngine, - _is_file_path + _is_file_path, _normalize_ip, _peel_wrapper, _exec_wrappers, EXEC_WRAPPERS, + _parse_subcommands, ) from secator.output_types import Warning, Error @@ -88,6 +89,53 @@ def test_match_rule_basename_for_paths(self): self.assertFalse(match_rule("example.com", [".com"])) +@unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') +class TestEncodedIPDeny(unittest.TestCase): + """M8: alternate IP encodings must not evade an IP/CIDR deny rule.""" + + META = "169.254.169.254" + + def test_normalize_ip_encodings(self): + import ipaddress + expected = ipaddress.ip_address(self.META) + for enc in ("2852039166", "0xA9FEA9FE", "0xa9fea9fe", + "::ffff:169.254.169.254", "[::ffff:169.254.169.254]", + "0xA9.0xFE.0xA9.0xFE", "169.254.169.254"): + self.assertEqual(_normalize_ip(enc), expected, enc) + + def test_normalize_ip_non_ip(self): + # Hostnames and port-suffixed values are not IPs (no DNS resolution here) + self.assertIsNone(_normalize_ip("example.com")) + self.assertIsNone(_normalize_ip("10.0.0.1:8080")) + + def test_encoded_forms_denied(self): + deny = ["169.254.169.254"] + for enc in ("2852039166", "0xA9FEA9FE", "::ffff:169.254.169.254", "169.254.169.254"): + self.assertTrue(match_rule(enc, deny), enc) + + def test_public_ip_still_allowed(self): + # A normal public IP must not match the metadata deny rule + self.assertFalse(match_rule("8.8.8.8", ["169.254.169.254"])) + self.assertFalse(match_rule("93.184.216.34", ["169.254.169.254"])) + + def test_cidr_deny_membership(self): + # Encoded link-local addresses fall inside a CIDR deny rule + self.assertTrue(match_rule("2852039166", ["169.254.0.0/16"])) + self.assertFalse(match_rule("8.8.8.8", ["169.254.0.0/16"])) + + def test_check_value_denies_encoded_targets(self): + engine = PermissionEngine(config=dict(deny=["target(169.254.169.254)"], allow=["target(*)"])) + for enc in ("2852039166", "0xA9FEA9FE", "::ffff:169.254.169.254"): + self.assertEqual(engine._check_value("target", enc).decision, "deny", enc) + self.assertEqual(engine._check_value("target", "8.8.8.8").decision, "allow") + + def test_encoded_url_target_denied(self): + # curl http:/// resolves to the metadata IP → deny (via URL host extraction) + engine = PermissionEngine(config=dict(deny=["target(169.254.169.254)"], allow=["task(*)", "target(*)"])) + result = engine.check_action({"action": "task", "name": "nmap", "targets": ["http://2852039166/latest/meta-data/"]}) + self.assertEqual(result.decision, "deny") + + @unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') class TestDetection(unittest.TestCase): @@ -212,6 +260,23 @@ def test_follow_up_always_allowed(self): result = engine.check_action({"action": "follow_up", "reason": "test"}) self.assertEqual(result.decision, "allow") + # --- M7: add_finding privileged-type gating --- + + def test_add_finding_benign_allowed(self): + engine = self._make_engine() + result = engine.check_action({"action": "add_finding", "_type": "vulnerability", "name": "XSS"}) + self.assertEqual(result.decision, "allow") + + def test_add_finding_target_type_not_allowed(self): + engine = self._make_engine() + result = engine.check_action({"action": "add_finding", "_type": "target", "name": "evil.com"}) + self.assertEqual(result.decision, "ask") + + def test_add_finding_target_type_case_insensitive(self): + engine = self._make_engine() + result = engine.check_action({"action": "add_finding", "_type": " Target ", "name": "evil.com"}) + self.assertEqual(result.decision, "ask") + # --- Target checks --- def test_target_allowed_via_targets_variable(self): @@ -239,6 +304,13 @@ def test_target_ask_for_unknown(self): self.assertEqual(result.decision, "ask") self.assertIn("10.5.2.3", result.targets) + def test_target_no_catchall_asks_not_allows(self): + """M6: with no target rule/catch-all configured, an unknown target must ask (fail-safe), not silently allow.""" + engine = self._make_engine(allow=["task(*)"]) # no target(...) rule in any category + result = engine.check_action({"action": "task", "name": "nmap", "targets": ["10.5.2.3"]}) + self.assertEqual(result.decision, "ask") + self.assertIn("10.5.2.3", result.targets) + def test_task_target_validation(self): engine = self._make_engine( allow=["task(*)", "target({targets})"], @@ -341,6 +413,135 @@ def test_default_ask_when_no_rules_match(self): self.assertIn("nmap", result.shell_command) +@unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') +class TestAllowedTargets(unittest.TestCase): + """Platform-supplied allowed_targets (e.g. validated mandates) constrain target scope.""" + + def _make_engine(self, allow=None, deny=None, ask=None, targets=None, allowed_targets=None, workspace="/tmp/workspace"): # noqa: E501 + config = {"allow": allow or [], "deny": deny or [], "ask": ask or []} + return PermissionEngine( + config, targets=targets or [], workspace=workspace, allowed_targets=allowed_targets or []) + + def test_allowed_target_literal_match(self): + engine = self._make_engine(allow=["shell(nmap)"], allowed_targets=["example.com"]) + result = engine.check_action({"action": "shell", "command": "nmap example.com"}) + self.assertEqual(result.decision, "allow") + + def test_allowed_target_regex_match(self): + engine = self._make_engine(allow=["shell(nmap)"], allowed_targets=[r".*\.example\.com"]) + result = engine.check_action({"action": "shell", "command": "nmap api.example.com"}) + self.assertEqual(result.decision, "allow") + + def test_target_outside_allowed_targets_is_constrained(self): + """A target not matching any allowed_targets regex must NOT be silently allowed.""" + engine = self._make_engine(allow=["shell(nmap)"], allowed_targets=[r".*\.example\.com"]) + result = engine.check_action({"action": "shell", "command": "nmap evil.attacker.com"}) + self.assertNotEqual(result.decision, "allow") + + def test_allowed_targets_presence_forces_target_check(self): + """Even with no config target rules, allowed_targets makes the target step run.""" + engine = self._make_engine(allow=["task(*)"], allowed_targets=["10.0.0.1"]) + result = engine.check_action({"action": "task", "name": "nmap", "targets": ["8.8.8.8"]}) + self.assertNotEqual(result.decision, "allow") + + def test_allowed_target_task_in_scope(self): + engine = self._make_engine(allow=["task(*)"], allowed_targets=[r"10\.0\.0\.\d+"]) + result = engine.check_action({"action": "task", "name": "nmap", "targets": ["10.0.0.5"]}) + self.assertEqual(result.decision, "allow") + + def test_deny_still_wins_over_allowed_targets(self): + engine = self._make_engine( + allow=["shell(nmap)"], deny=["target(169.254.169.254)"], allowed_targets=[r".*"]) + result = engine.check_action({"action": "shell", "command": "nmap 169.254.169.254"}) + self.assertEqual(result.decision, "deny") + + def test_invalid_regex_falls_back_to_literal(self): + # '[' is invalid regex → treated as a literal string + engine = self._make_engine(allow=["shell(nmap)"], allowed_targets=["host[1"]) + self.assertEqual(len(engine.allowed_targets), 1) + + def test_allowed_target_url_host_component(self): + engine = self._make_engine(allow=["shell(curl)"], allowed_targets=["example.com"]) + result = engine.check_action({"action": "shell", "command": "curl https://example.com/path"}) + self.assertEqual(result.decision, "allow") + + +@unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') +class TestDeniedTargets(unittest.TestCase): + """Platform-supplied denied_targets (e.g. mandate deny scope) block target scope. Deny wins.""" + + def _make_engine(self, allow=None, deny=None, ask=None, targets=None, allowed_targets=None, # noqa: E501 + denied_targets=None, workspace="/tmp/workspace"): + config = {"allow": allow or [], "deny": deny or [], "ask": ask or []} + return PermissionEngine( + config, targets=targets or [], workspace=workspace, + allowed_targets=allowed_targets or [], denied_targets=denied_targets or []) + + def test_denied_target_literal_match(self): + # IP targets are extracted without DNS, so deny applies directly. + engine = self._make_engine(allow=["shell(nmap)"], denied_targets=["10.0.0.1"]) + result = engine.check_action({"action": "shell", "command": "nmap 10.0.0.1"}) + self.assertEqual(result.decision, "deny") + + def test_denied_target_regex_match(self): + # Hostnames are only extracted if they resolve — patch DNS so the host is seen. + engine = self._make_engine(allow=["shell(nmap)"], denied_targets=[r".*\.evil\.com"]) + with patch("secator.ai.guardrails._resolves", return_value=True): + result = engine.check_action({"action": "shell", "command": "nmap api.evil.com"}) + self.assertEqual(result.decision, "deny") + + def test_deny_wins_over_allow_when_target_matches_both(self): + """A target matching BOTH allowed_targets and denied_targets must be DENIED.""" + engine = self._make_engine( + allow=["shell(nmap)"], allowed_targets=[r".*\.example\.com"], denied_targets=[r"admin\.example\.com"]) + with patch("secator.ai.guardrails._resolves", return_value=True): + result = engine.check_action({"action": "shell", "command": "nmap admin.example.com"}) + self.assertEqual(result.decision, "deny") + + def test_deny_wins_over_allow_at_check_value_level(self): + """Unit-level deny-wins: _check_value denies a target in both allow + deny lists.""" + engine = self._make_engine( + allow=["shell(nmap)"], allowed_targets=[r".*"], denied_targets=[r"169\.254\.169\.254"]) + result = engine._check_value("target", "169.254.169.254") + self.assertEqual(result.decision, "deny") + + def test_allow_only_target_is_allowed(self): + """allow-only (in allowed, not in denied) → allowed.""" + engine = self._make_engine( + allow=["shell(nmap)"], allowed_targets=[r".*"], denied_targets=[r"169\.254\.169\.254"]) + result = engine._check_value("target", "10.0.0.5") + self.assertEqual(result.decision, "allow") + + def test_deny_only_target_is_denied(self): + """deny-only (matches denied, no allowed entries) → denied.""" + engine = self._make_engine(allow=["shell(nmap)"], denied_targets=[r"10\.0\.0\.1"]) + result = engine.check_action({"action": "shell", "command": "nmap 10.0.0.1"}) + self.assertEqual(result.decision, "deny") + + def test_denied_targets_presence_forces_target_check(self): + """Even with no config target rules, denied_targets makes the target step run.""" + engine = self._make_engine(allow=["task(*)"], denied_targets=["10.0.0.1"]) + result = engine.check_action({"action": "task", "name": "nmap", "targets": ["10.0.0.1"]}) + self.assertEqual(result.decision, "deny") + + def test_denied_target_task_in_deny_scope(self): + engine = self._make_engine( + allow=["task(*)"], allowed_targets=[r"10\.0\.0\.\d+"], denied_targets=[r"10\.0\.0\.1"]) + result = engine.check_action({"action": "task", "name": "nmap", "targets": ["10.0.0.1"]}) + self.assertEqual(result.decision, "deny") + + def test_denied_target_url_host_component(self): + engine = self._make_engine(allow=["shell(curl)"], denied_targets=["evil.com"]) + with patch("secator.ai.guardrails._resolves", return_value=True): + result = engine.check_action({"action": "shell", "command": "curl https://evil.com/path"}) + self.assertEqual(result.decision, "deny") + + def test_invalid_regex_falls_back_to_literal(self): + # '[' is invalid regex → treated as a literal string + engine = self._make_engine(allow=["shell(nmap)"], denied_targets=["host[1"]) + self.assertEqual(len(engine.denied_targets), 1) + + @unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') class TestTargetPrompt(unittest.TestCase): @@ -411,6 +612,54 @@ def test_prompt_target_deny_choice(self): self.assertEqual(result, "deny") +@unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') +class TestPromptShell(unittest.TestCase): + + def _make_engine(self, allow=None, deny=None, ask=None): + config = {"allow": allow or [], "deny": deny or [], "ask": ask or []} + return PermissionEngine(config) + + def _menu_returning(self, idx): + """Patch the rich menu so .show() yields (idx, label).""" + menu = MagicMock() + menu.return_value.show.return_value = (idx, "") + return menu + + def test_prompt_shell_non_interactive_returns_deny(self): + engine = self._make_engine(ask=["shell(*)"]) + self.assertEqual(engine.prompt_shell("curl https://x", interactive=False), "deny") + + def test_allow_this_command_is_one_shot(self): + """Option 0 approves ONLY this invocation — no session rule; the next call re-prompts (H9).""" + engine = self._make_engine(ask=["shell(*)"]) + with patch('secator.rich.InteractiveMenu', self._menu_returning(0)), \ + patch('secator.ai.guardrails._extract_cmd_names', return_value=["curl"]): + result = engine.prompt_shell("curl https://good.example") + self.assertEqual(result, "allow") + # No runtime rule was added, so a second, different-arg curl is NOT auto-allowed + self.assertEqual(engine.runtime_allow, []) + self.assertEqual(engine._check_value("shell", "curl").decision, "ask") + + def test_allow_all_commands_adds_session_rule(self): + """Option 1 persists a session-wide allow for the command name (unchanged).""" + engine = self._make_engine(ask=["shell(*)"]) + with patch('secator.rich.InteractiveMenu', self._menu_returning(1)), \ + patch('secator.ai.guardrails._extract_cmd_names', return_value=["curl"]): + result = engine.prompt_shell("curl https://good.example") + self.assertEqual(result, "allow") + self.assertEqual(engine.runtime_allow, [("shell", ["curl"])]) + # Now any curl is auto-allowed for the session + self.assertEqual(engine._check_value("shell", "curl").decision, "allow") + + def test_deny_choice_blocks(self): + engine = self._make_engine(ask=["shell(*)"]) + with patch('secator.rich.InteractiveMenu', self._menu_returning(2)), \ + patch('secator.ai.guardrails._extract_cmd_names', return_value=["curl"]): + result = engine.prompt_shell("curl https://good.example") + self.assertEqual(result, "deny") + self.assertEqual(engine.runtime_allow, []) + + @unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') class TestGuardrailsIntegration(unittest.TestCase): @@ -610,6 +859,55 @@ def test_execute_command_not_in_whitelist(self): self.assertEqual(result.decision, "ask") self.assertIn("rm", result.shell_command) + # === Exec-wrapper laundering (C2) + destructive deny (H6) === + + def test_timeout_wrapper_does_not_launder_destructive_rm(self): + """`timeout 60 rm -rf /` must NOT auto-allow via the timeout wrapper (C2 + H6).""" + engine = self._engine() + result = engine.check_action({"action": "shell", "command": "timeout 60 rm -rf /"}) + self.assertEqual(result.decision, "deny") + + def test_xargs_wrapper_does_not_launder_inner_command(self): + """`xargs ... rm ...` must not auto-allow via the xargs wrapper (C2).""" + engine = self._engine() + result = engine.check_action({"action": "shell", "command": "xargs -I{} rm -rf {}"}) + # Inner rm is unknown (not root-destructive) -> prompt, never silent allow. + self.assertEqual(result.decision, "ask") + + def test_timeout_wrapper_restores_interpreter_ask_gate(self): + """Wrapping an interpreter (`timeout 60 bash -c ...`) must keep the ask gate (C2).""" + engine = self._engine() + result = engine.check_action({"action": "shell", "command": "timeout 60 bash -c 'rm -rf /'"}) + self.assertEqual(result.decision, "ask") + + def test_sudo_wrapper_does_not_launder_destructive_rm(self): + """`sudo rm -rf /` must be denied, not laundered through sudo (C2 + H6).""" + engine = self._engine() + result = engine.check_action({"action": "shell", "command": "sudo rm -rf /"}) + self.assertEqual(result.decision, "deny") + + def test_destructive_root_rm_denied(self): + """Bare `rm -rf /` (and one level under /) must be denied (H6).""" + engine = self._engine() + self.assertEqual( + engine.check_action({"action": "shell", "command": "rm -rf /"}).decision, "deny") + self.assertEqual( + engine.check_action({"action": "shell", "command": "rm -rf /etc"}).decision, "deny") + + def test_scoped_rm_still_prompts_not_denied(self): + """Scoped `rm -rf /tmp/x` is not catastrophic -> prompt (not silent allow/deny) (H6).""" + engine = self._engine() + result = engine.check_action({"action": "shell", "command": "rm -rf /tmp/data/x"}) + self.assertEqual(result.decision, "ask") + + def test_wrapper_preserves_allowed_inner_command(self): + """`timeout 60 curl ...` must not regress: curl stays allowed at the action level.""" + engine = self._engine(targets=["10.0.0.1"]) + # Inner curl is allow-listed; the unknown URL target is what triggers the ask, + # proving the wrapper was peeled and curl recognised (not denied). + result = engine.check_action({"action": "shell", "command": "timeout 60 curl http://10.0.0.1/x"}) + self.assertEqual(result.decision, "allow") + # === Should NOT trigger approval (allow) === def test_read_file_in_workspace(self): @@ -695,6 +993,23 @@ def test_mixed_read_write_in_single_command(self): self.assertEqual(access_map["/etc/hosts"], "read") self.assertEqual(access_map["/tmp/copy.txt"], "write") + # --- M9: output-flag destinations are writes (shfmt-gated: need real shell parser) --- + + def test_curl_output_flag_classified_as_write(self): + """curl -o dest is a write, so `deny write(/etc/*)` fires (not a read).""" + paths = detect_paths_with_access("curl -o /etc/passwd http://x") + self.assertIn(("/etc/passwd", "write"), paths) + + def test_wget_output_flag_classified_as_write(self): + """wget -O dest is a write.""" + paths = detect_paths_with_access("wget -O /etc/passwd http://x") + self.assertIn(("/etc/passwd", "write"), paths) + + def test_curl_without_output_flag_stays_read(self): + """curl with no -o only reads (URL is not a file path); no write leaks in.""" + paths = detect_paths_with_access("curl http://x") + self.assertNotIn("write", [a for _, a in paths]) + def test_fd_redirect_2_to_1_not_detected_as_path(self): """2>&1 is a fd redirect, not a file path.""" paths = detect_paths('curl -sk "http://example.com" 2>&1 | head -100') @@ -965,5 +1280,184 @@ def test_runtime_allow_subdirectory_matching(self): self.assertEqual(result.decision, "allow") +class TestOutputFlagWrites(unittest.TestCase): + """M9: output-flag write classification, proven locally by stubbing the shell + parser (real shfmt is absent in CI-less envs, which makes the tests above no-ops).""" + + def _paths(self, argv, redirects=None): + """Run detect_paths_with_access with a stubbed extract_commands (no shfmt).""" + with patch('safecmd.bashxtract.extract_commands', + return_value=([argv], [], redirects or [])): + return detect_paths_with_access(" ".join(argv)) + + def test_curl_o_space_form_is_write(self): + paths = self._paths(["curl", "-o", "/etc/passwd", "http://x"]) + self.assertIn(("/etc/passwd", "write"), paths) + + def test_curl_long_output_equals_form_is_write(self): + paths = self._paths(["curl", "--output=/etc/passwd", "http://x"]) + self.assertIn(("/etc/passwd", "write"), paths) + + def test_curl_o_attached_short_form_is_write(self): + paths = self._paths(["curl", "-o/etc/passwd", "http://x"]) + self.assertIn(("/etc/passwd", "write"), paths) + + def test_wget_O_form_is_write(self): + paths = self._paths(["wget", "-O", "/etc/passwd", "http://x"]) + self.assertIn(("/etc/passwd", "write"), paths) + + def test_wget_output_document_equals_form_is_write(self): + paths = self._paths(["wget", "--output-document=/etc/passwd", "http://x"]) + self.assertIn(("/etc/passwd", "write"), paths) + + def test_curl_no_output_flag_has_no_write(self): + paths = self._paths(["curl", "http://x"]) + self.assertNotIn("write", [a for _, a in paths]) + + def test_curl_o_stdout_dash_not_treated_as_file(self): + paths = self._paths(["curl", "-o", "-", "http://x"]) + self.assertEqual(paths, []) + + def test_redirect_still_write_with_output_flag_cmd(self): + """Redirect classification is preserved alongside the new flag handling.""" + paths = self._paths(["echo", "x"], redirects=[("", "/etc/y")]) + self.assertIn(("/etc/y", "write"), paths) + + +@unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') +class TestWrapperPeelingM11(unittest.TestCase): + """M11: broadened + arg-grammar-aware exec-wrapper peeling closes the C2 laundering class. + + The `_peel_wrapper` assertions are PROVEN (they take a token list — no shfmt needed). + The `check_action` integration assertions stub `extract_commands` (real shfmt is absent + in CI-less envs, which no-ops every parser-dependent test), same pattern as + TestOutputFlagWrites.""" + + # --- PROVEN: peel locates the leaf command past each wrapper's own arg grammar --- + + def test_proxychains_peels_to_inner(self): + self.assertEqual(_peel_wrapper(["proxychains", "curl", "http://evil"]), ["curl", "http://evil"]) + + def test_proxychains_config_flag_consumed(self): + self.assertEqual(_peel_wrapper(["proxychains", "-f", "/etc/pc.conf", "dd"]), ["dd"]) + + def test_firejail_peels_past_long_opts(self): + self.assertEqual(_peel_wrapper(["firejail", "--net=none", "rm", "-rf", "/tmp/x"]), ["rm", "-rf", "/tmp/x"]) + + def test_flock_lockfile_positional_consumed(self): + self.assertEqual(_peel_wrapper(["flock", "/tmp/l", "curl", "http://evil"]), ["curl", "http://evil"]) + + def test_flock_value_opt_then_lockfile(self): + self.assertEqual(_peel_wrapper(["flock", "-w", "5", "/tmp/l", "dd"]), ["dd"]) + + def test_runuser_cmd_string_reparsed(self): + self.assertEqual(_peel_wrapper(["runuser", "-c", "curl http://evil"]), ["curl", "http://evil"]) + + def test_runuser_user_then_dashdash(self): + self.assertEqual(_peel_wrapper(["runuser", "-u", "bob", "--", "curl", "http://evil"]), ["curl", "http://evil"]) + + def test_su_user_positional_and_cmd_string(self): + self.assertEqual(_peel_wrapper(["su", "root", "-c", "dd if=/dev/zero"]), ["dd", "if=/dev/zero"]) + + def test_script_cmd_string_reparsed(self): + self.assertEqual(_peel_wrapper(["script", "-c", "curl http://evil", "/tmp/log"]), ["curl", "http://evil"]) + + def test_torsocks_peels_to_inner(self): + self.assertEqual(_peel_wrapper(["torsocks", "curl", "http://evil"]), ["curl", "http://evil"]) + + def test_sudo_user_value_opt_consumed(self): + # pre-existing C2 gap: `sudo -u bob` mis-read `bob` as the command; grammar now consumes it + self.assertEqual(_peel_wrapper(["sudo", "-u", "bob", "rm", "-rf", "/"]), ["rm", "-rf", "/"]) + + # --- PROVEN: C2-covered wrappers + normal commands unchanged --- + + def test_c2_timeout_still_peels(self): + self.assertEqual(_peel_wrapper(["timeout", "60", "rm", "-rf", "/"]), ["rm", "-rf", "/"]) + + def test_c2_interpreter_gate_preserved(self): + # bash is NOT a wrapper — it stays the leaf so its ask-gate still fires + self.assertEqual(_peel_wrapper(["timeout", "60", "bash", "-c", "rm -rf /"]), ["bash", "-c", "rm -rf /"]) + + def test_normal_command_untouched(self): + self.assertEqual(_peel_wrapper(["curl", "http://ok"]), ["curl", "http://ok"]) + + def test_bare_wrapper_checked_by_name(self): + self.assertEqual(_peel_wrapper(["sudo"]), ["sudo"]) + + # --- PROVEN: config EXTENDS the built-in baseline (never shrinks below it) --- + + def test_config_added_wrapper_honored(self): + try: + CONFIG.addons.ai.exec_wrappers = ["myrunner"] + self.assertIn("myrunner", _exec_wrappers()) + self.assertTrue(EXEC_WRAPPERS <= _exec_wrappers()) # baseline is the floor + self.assertEqual(_peel_wrapper(["myrunner", "dd", "if=/dev/zero"]), ["dd", "if=/dev/zero"]) + finally: + CONFIG.addons.ai.exec_wrappers = [] + + # --- PROVEN-via-stub: end-to-end deny/ask fires on the peeled leaf, not the wrapper name --- + + def _decide(self, argv): + engine = PermissionEngine(dict(CONFIG.addons.ai.permissions), targets=["10.0.0.1"], + workspace="/home/user/.secator/reports/test/tasks/ai_1") + with patch('safecmd.bashxtract.extract_commands', return_value=([argv], [], [])): + return engine.check_action({"action": "shell", "command": " ".join(argv)}).decision + + def test_proxychains_denied_inner_command(self): + self.assertEqual(self._decide(["proxychains", "dd", "if=/dev/zero"]), "deny") # dd is deny-listed + + def test_flock_launders_denied_command(self): + self.assertEqual(self._decide(["flock", "/tmp/l", "dd"]), "deny") + + def test_firejail_scoped_rm_asks(self): + self.assertEqual(self._decide(["firejail", "rm", "-rf", "/tmp/x"]), "ask") # not silent-allowed as firejail + + +@unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') +class TestShellParserFallback(unittest.TestCase): + """When the shfmt-based shell parser (safecmd/shfmt) is unavailable, the + guardrail must Warn (NOT claim 'Missing ai addon') and fall back to + whole-command approval — an empty sub-command list makes the caller `ask`.""" + + def setUp(self): + import secator.ai.guardrails as g + g._SHELL_PARSER_WARNED = False # reset warn-once flag per test + + def _run_and_capture(self): + printed = [] + with patch('secator.rich.console.print', side_effect=lambda x, *a, **k: printed.append(x)): + result = _parse_subcommands('curl -s https://x.com | head -5') + return result, printed + + def test_missing_safecmd_warns_not_ai_addon(self): + # Simulate safecmd not installed -> ImportError on the in-function import. + with patch.dict('sys.modules', {'safecmd.bashxtract': None}): + result, printed = self._run_and_capture() + self.assertEqual(result, []) # unparseable -> caller falls back to ask + self.assertEqual(len(printed), 1) + item = printed[0] + self.assertIsInstance(item, Warning) # a Warning, not an Error + self.assertIn('safecmd', item.message) + self.assertNotIn('ai addon', item.message.lower()) + + def test_missing_shfmt_binary_warns(self): + # safecmd imports, but the shfmt binary it shells out to is not on PATH. + with patch('safecmd.bashxtract.extract_commands', side_effect=FileNotFoundError('shfmt')): + result, printed = self._run_and_capture() + self.assertEqual(result, []) + self.assertEqual(len(printed), 1) + self.assertIsInstance(printed[0], Warning) + self.assertIn('shfmt', printed[0].message) + self.assertNotIn('ai addon', printed[0].message.lower()) + + def test_warns_only_once_across_commands(self): + printed = [] + with patch('safecmd.bashxtract.extract_commands', side_effect=FileNotFoundError): + with patch('secator.rich.console.print', side_effect=lambda x, *a, **k: printed.append(x)): + _parse_subcommands('a | b') + _parse_subcommands('c | d') + self.assertEqual(len(printed), 1) # warn-once, no per-command spam + + if __name__ == '__main__': unittest.main() diff --git a/tests/unit/test_ai_handlers.py b/tests/unit/test_ai_handlers.py index 2ad9413e8..aabba8d5d 100644 --- a/tests/unit/test_ai_handlers.py +++ b/tests/unit/test_ai_handlers.py @@ -150,7 +150,7 @@ def test_ai_task_has_required_opts(self): from secator.tasks.ai import ai required_opts = ['prompt', 'mode', 'model', 'api_base', 'sensitive', - 'max_iterations', 'temperature', 'dry_run', 'yes'] + 'max_iterations', 'temperature', 'dry_run'] for opt in required_opts: self.assertIn(opt, ai.opts, f"Missing opt: {opt}") diff --git a/tests/unit/test_ai_history.py b/tests/unit/test_ai_history.py index 983a9bcb5..22d149028 100644 --- a/tests/unit/test_ai_history.py +++ b/tests/unit/test_ai_history.py @@ -13,18 +13,9 @@ @unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') class TestChatHistory(unittest.TestCase): - def test_add_system(self): - history = ChatHistory() - history.add_system("You are an assistant.") - - messages = history.to_messages() - self.assertEqual(len(messages), 1) - self.assertEqual(messages[0]["role"], "system") - self.assertEqual(messages[0]["content"], "You are an assistant.") - def test_set_system_replaces_existing(self): history = ChatHistory() - history.add_system("old prompt") + history.set_system("old prompt") history.add_user("user msg") history.set_system("new prompt") @@ -62,7 +53,7 @@ def test_add_assistant(self): def test_to_messages_returns_list(self): history = ChatHistory() - history.add_system("sys") + history.set_system("sys") history.add_user("user") history.add_assistant("assistant") @@ -88,16 +79,6 @@ def test_to_messages_returns_copy(self): # Original should be unchanged self.assertEqual(len(history.to_messages()), 1) - - def test_add_tool(self): - history = ChatHistory() - history.add_tool("tool output here") - - messages = history.to_messages() - self.assertEqual(len(messages), 1) - self.assertEqual(messages[0]["role"], "tool") - self.assertEqual(messages[0]["content"], "tool output here") - @patch('secator.ai.history.get_context_window') @patch('litellm.token_counter') def test_maybe_summarize_below_threshold(self, mock_token_counter, mock_get_ctx): @@ -106,7 +87,7 @@ def test_maybe_summarize_below_threshold(self, mock_token_counter, mock_get_ctx) mock_token_counter.return_value = 1000 # Well under 85% history = ChatHistory() - history.add_system("system prompt") + history.set_system("system prompt") history.add_user("short message") summarized, old_tokens, new_tokens = history.maybe_summarize("test-model") @@ -130,7 +111,7 @@ def test_maybe_summarize_above_threshold(self, mock_token_counter, mock_call_llm mock_call_llm.return_value = {"content": "Summary of session.", "usage": None} history = ChatHistory() - history.add_system("system prompt") + history.set_system("system prompt") # Add enough content to exceed threshold for i in range(20): history.add_user("x" * 200) @@ -154,7 +135,7 @@ def test_summarize_preserves_system_prompt(self, mock_token_counter, mock_call_l mock_call_llm.return_value = {"content": "Compact summary.", "usage": None} history = ChatHistory() - history.add_system("You are an AI pentester.") + history.set_system("You are an AI pentester.") for i in range(10): history.add_user("x" * 200) history.add_assistant("y" * 200) @@ -175,7 +156,7 @@ def test_summarize_preserves_system_prompt(self, mock_token_counter, mock_call_l def test_trim_drops_oldest_messages(self): """Trim drops messages to fit under token limit.""" history = ChatHistory() - history.add_system("s" * 40) + history.set_system("s" * 40) history.add_user("u" * 40) for i in range(10): history.add_user(f"msg{i} " + "x" * 200) @@ -190,10 +171,52 @@ def test_trim_drops_oldest_messages(self): self.assertEqual(history.messages[0]["role"], "system") self.assertEqual(history.messages[0]["content"], "s" * 40) + def test_trim_sanitizes_none_content_before_trimmer(self): + """Messages with content=None (assistant tool-call turns) must never reach + trim_messages as None — litellm's shorten path does len(content) and crashes + with 'object of type NoneType has no len()'. trim() coerces them to "". + """ + history = ChatHistory() + history.set_system("s") + history.add_user("u") + # assistant turn carrying only tool_calls -> content is None + history.add_assistant_with_tool_calls(None, [ + {"id": "1", "type": "function", "function": {"name": "q", "arguments": "{}"}} + ]) + history.add_tool_result("q", "1", "result") + + captured = {} + + def fake_trim(messages, max_tokens): + captured["messages"] = messages + # replicate litellm's crashing operation to prove it no longer crashes + for m in messages: + if m.get("role") != "system": + _ = len(m["content"]) # would raise TypeError on None + return messages + + with patch("litellm.utils.trim_messages", side_effect=fake_trim): + history.trim(max_tokens=1000) # must not raise + + self.assertTrue(all(m.get("content") is not None for m in captured["messages"])) + + def test_trim_survives_trimmer_exception(self): + """A crash inside trim_messages must not propagate and kill the AI loop — + trim() degrades to the (sanitized) untrimmed history instead. + """ + history = ChatHistory() + history.set_system("s") + history.add_user("u") + + with patch("litellm.utils.trim_messages", side_effect=RuntimeError("boom")): + out = history.trim(max_tokens=1000) # must not raise + + self.assertEqual(len(out), 2) + def test_to_messages_with_max_tokens_total(self): """to_messages with max_tokens_total trims messages.""" history = ChatHistory() - history.add_system("s" * 40) + history.set_system("s" * 40) history.add_user("u" * 40) for i in range(20): history.add_user("x" * 400) @@ -212,16 +235,86 @@ def test_to_messages_with_max_tokens_total(self): def test_to_messages_no_truncation_when_under_limit(self): """to_messages with max_tokens_total does nothing when under limit.""" history = ChatHistory() - history.add_system("short") + history.set_system("short") history.add_user("msg") messages = history.to_messages(max_tokens_total=500) self.assertEqual(len(messages), 2) + @patch('secator.ai.history.get_context_window') + def test_to_messages_caps_budget_to_small_window(self, mock_get_ctx): + """M3: a flat max_tokens_total is capped to a small model's window.""" + mock_get_ctx.return_value = 8000 # small-window model + + history = ChatHistory() + history.model = "small-model" + history.set_system("s" * 40) + for i in range(40): + history.add_user("x" * 4000) # long history, well over 8k tokens + + original_count = len(history.messages) + # Flat 100k cap would NOT trim on a real 8k model without this fix. + messages = history.to_messages(max_tokens_total=100000) + + self.assertLess(len(messages), original_count) # trimmed to fit the window + self.assertEqual(messages[0]["role"], "system") + + @patch('secator.ai.history.get_context_window') + def test_to_messages_no_explicit_cap_uses_window(self, mock_get_ctx): + """M3: with max_tokens_total=0 and a model, trim to the window-derived budget.""" + mock_get_ctx.return_value = 8000 + + history = ChatHistory() + history.model = "small-model" + history.set_system("s" * 40) + for i in range(40): + history.add_user("x" * 4000) + + original_count = len(history.messages) + messages = history.to_messages() # no explicit cap + + self.assertLess(len(messages), original_count) + + @patch('secator.ai.history.get_context_window') + def test_to_messages_large_window_matches_flat_budget(self, mock_get_ctx): + """M3: on a large-window model the flat cap is honored (no over-trim).""" + mock_get_ctx.return_value = 200000 # window - reserve (191808) > 100k cap + + history = ChatHistory() + history.model = "big-model" + history.set_system("short") + history.add_user("small message") + + with patch.object(history, 'trim', wraps=history.trim) as spy: + history.to_messages(max_tokens_total=100000) + # Budget = min(100000, 200000 - 8192) = 100000, unchanged by the window. + spy.assert_called_once_with(100000) + + @patch('secator.ai.history.get_context_window') + def test_to_messages_window_cap_preserves_tool_pairs(self, mock_get_ctx): + """M3 + H2: window-capped trim never leaves a leading orphan tool_result.""" + mock_get_ctx.return_value = 8000 + + history = ChatHistory() + history.model = "small-model" + history.set_system("s" * 40) + for i in range(30): + tool_calls = [{"id": f"call_{i}", "type": "function", + "function": {"name": "nmap", "arguments": "{}"}}] + history.add_assistant_with_tool_calls("x" * 2000, tool_calls) + history.add_tool_result("nmap", f"call_{i}", "y" * 2000) + + messages = history.to_messages(max_tokens_total=100000) + + # First non-system message must not be an orphan tool result. + non_system = [m for m in messages if m["role"] != "system"] + if non_system: + self.assertNotEqual(non_system[0]["role"], "tool") + def test_to_messages_no_truncation_when_zero(self): """to_messages without max_tokens_total does not truncate.""" history = ChatHistory() - history.add_system("s" * 40) + history.set_system("s" * 40) history.add_user("u" * 40) for i in range(20): history.add_user("x" * 400) @@ -235,7 +328,7 @@ def test_to_messages_no_truncation_when_zero(self): def test_maybe_summarize_skips_when_not_needed(self, mock_token_counter, mock_get_model_info): history = ChatHistory() - history.add_system("system") + history.set_system("system") history.add_user("user") original_messages = history.to_messages() @@ -321,7 +414,7 @@ def test_set_system_invalidates_token_cache(self, mock_token_counter): mock_token_counter.return_value = 50 history = ChatHistory() - history.add_system("old prompt") + history.set_system("old prompt") # Count tokens - this caches the count history.count_tokens("gpt-4") @@ -485,32 +578,6 @@ def test_truncate_to_tokens_with_fallback_path(self, mock_token_counter): finally: fallback_path.unlink() - @patch('litellm.token_counter') - def test_truncate_to_tokens_saves_shell_output(self, mock_token_counter): - """truncate_to_tokens saves shell output to .outputs directory.""" - from secator.ai.history import truncate_to_tokens - - mock_token_counter.return_value = 1000 - content = "shell output " * 500 - - with tempfile.TemporaryDirectory() as tmpdir: - output_dir = Path(tmpdir) - - result = truncate_to_tokens( - content, 100, "gpt-4", - output_dir=output_dir, - result_name="shell" - ) - - self.assertIn("[TRUNCATED]", result) - self.assertIn("saved to:", result) - - # Verify file was created - saved_files = list(output_dir.glob("shell_*.txt")) - self.assertEqual(len(saved_files), 1) - self.assertEqual(saved_files[0].read_text(), content) - - @patch('secator.ai.history.get_context_window') @patch('secator.ai.utils.call_llm') @patch('litellm.token_counter') @@ -523,7 +590,7 @@ def test_maybe_summarize_uses_percentage_threshold(self, mock_token_counter, moc mock_call_llm.return_value = {"content": "Summary.", "usage": None} history = ChatHistory() - history.add_system("system") + history.set_system("system") history.add_user("user1") history.add_assistant("response1") @@ -540,7 +607,7 @@ def test_maybe_summarize_no_threshold_param(self, mock_token_counter, mock_get_c mock_token_counter.return_value = 1000 history = ChatHistory() - history.add_system("system") + history.set_system("system") # Should work without threshold param summarized, _, _ = history.maybe_summarize("gpt-4") @@ -609,7 +676,7 @@ def test_summarize_handles_tool_messages(self, mock_token_counter, mock_call_llm mock_call_llm.return_value = {"content": "Summary with tool results.", "usage": None} history = ChatHistory() - history.add_system("system prompt") + history.set_system("system prompt") history.add_user("scan target.com") # Add several rounds with tool calling messages @@ -635,7 +702,7 @@ def test_count_tokens_by_role(self, mock_token_counter): mock_token_counter.return_value = 100 history = ChatHistory() - history.add_system("system prompt") + history.set_system("system prompt") history.add_user("user message") history.add_assistant("assistant reply") history.add_tool_result("tool_func", "call_1", "tool result") @@ -665,5 +732,27 @@ def test_count_tokens_by_role_aggregates(self, mock_token_counter): self.assertNotIn("system", result) +@unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') +class TestCapMessage(unittest.TestCase): + + def test_cap_message_truncates_content_and_arguments(self): + from secator.ai.history import cap_message, MAX_PERSISTED_MESSAGE_CHARS + big = "x" * (MAX_PERSISTED_MESSAGE_CHARS + 500) + msg = {"role": "assistant", "content": big, + "tool_calls": [{"id": "1", "type": "function", + "function": {"name": "q", "arguments": big}}]} + out = cap_message(msg) + assert len(out["content"]) <= MAX_PERSISTED_MESSAGE_CHARS + 20 # marker slack + assert "[capped]" in out["content"] + assert len(out["tool_calls"][0]["function"]["arguments"]) <= MAX_PERSISTED_MESSAGE_CHARS + 20 + # original not mutated + assert len(msg["content"]) == MAX_PERSISTED_MESSAGE_CHARS + 500 + + def test_cap_message_leaves_small_messages_untouched(self): + from secator.ai.history import cap_message + msg = {"role": "tool", "tool_call_id": "1", "content": "small"} + assert cap_message(msg) == msg + + if __name__ == '__main__': unittest.main() diff --git a/tests/unit/test_ai_interactivity.py b/tests/unit/test_ai_interactivity.py index 6060c5f16..b9453731d 100644 --- a/tests/unit/test_ai_interactivity.py +++ b/tests/unit/test_ai_interactivity.py @@ -2,6 +2,9 @@ import unittest from unittest.mock import MagicMock, patch +from secator.definitions import ADDONS_ENABLED +HAS_AI = ADDONS_ENABLED.get('ai', False) + class TestInteractivityBackendBase(unittest.TestCase): """Verify base class interface.""" @@ -15,7 +18,7 @@ def test_base_ask_user_raises(self): def test_base_get_excluded_tools(self): from secator.ai.interactivity import InteractivityBackend backend = InteractivityBackend() - self.assertEqual(backend.get_excluded_tools(), set()) + self.assertEqual(backend.get_excluded_tools(), {"stop"}) def test_base_get_extra_tools(self): from secator.ai.interactivity import InteractivityBackend @@ -100,14 +103,125 @@ def test_ask_user_polls_until_timeout(self, mock_sleep): # Should have called update to set timed_out mock_engine.update.assert_called_once() + def test_poll_scopes_query_to_prompt_uuid(self): + """The poll must correlate on the specific prompt's uuid. + + Regression test for the infinite-respawn loop: without scoping on + prompt_uuid, a stale answered follow_up from a prior turn resolves the + current wait immediately, the worker re-injects that old answer as a new + prompt and re-runs the turn forever. The query MUST include + extra_data.prompt_uuid so only THIS prompt's own answer resolves it. + """ + from secator.ai.interactivity import RemoteBackend + mock_engine = MagicMock() + mock_engine.search.return_value = [{"answer": "the right answer"}] + backend = RemoteBackend(timeout=60, query_engine=mock_engine, poll_interval=0.01) + + result = backend.ask_user("What next?", [], "session1", prompt_uuid="abc-123") + + self.assertEqual(result["answer"], "the right answer") + # The search query must be scoped to this prompt's uuid (else a stale + # answered follow_up from a prior turn would match -> loop). + search_query = mock_engine.search.call_args[0][0] + self.assertEqual(search_query.get("extra_data.prompt_uuid"), "abc-123") + self.assertEqual(search_query.get("status"), "answered") + @patch('secator.ai.interactivity.sleep') - def test_ask_user_returns_on_second_poll(self, mock_sleep): + def test_timeout_update_scoped_to_prompt_uuid(self, mock_sleep): + """On timeout, only THIS prompt's pending doc is flipped to timed_out.""" from secator.ai.interactivity import RemoteBackend mock_engine = MagicMock() - mock_engine.search.side_effect = [ - [], # first poll: not answered - [{"answer": "option B"}], # second poll: answered + mock_engine.search.return_value = [] # never answered + mock_engine.update = MagicMock() + backend = RemoteBackend(timeout=5, query_engine=mock_engine, poll_interval=5) + + result = backend.ask_user("What next?", [], "session1", prompt_uuid="abc-123") + + self.assertIsNone(result) + mock_engine.update.assert_called_once() + update_query = mock_engine.update.call_args[0][0] + self.assertEqual(update_query.get("extra_data.prompt_uuid"), "abc-123") + self.assertEqual(update_query.get("status"), "pending") + + def test_build_pending_prompt_stamps_permission_prompt_uuid(self): + """A permission pending doc carries its prompt_uuid in extra_data.""" + from secator.ai.interactivity import RemoteBackend + backend = RemoteBackend(timeout=60, query_engine=MagicMock()) + item = backend.build_pending_prompt( + "Shell `nmap` requires approval", ["allow", "deny"], "session1", + prompt_type="permission", permission_type="shell", value="nmap", + prompt_uuid="uuid-shell", + ) + self.assertEqual(item.extra_data.get("prompt_uuid"), "uuid-shell") + self.assertEqual(item.extra_data.get("permission_type"), "shell") + self.assertEqual(item.ai_type, "permission") + self.assertEqual(item.status, "pending") + + @patch('secator.ai.interactivity.sleep') + def test_later_permission_layer_does_not_resolve_from_earlier_allow(self, mock_sleep): + """H7: a later guardrail layer must not auto-resolve from an earlier 'allow'.""" + from secator.ai.interactivity import RemoteBackend + + # Fake "DB": one answered doc from the FIRST (shell) layer only. + answered_db = [{ + "_type": "ai", "ai_type": "permission", "status": "answered", + "_context": {"session_id": "session1"}, + "extra_data": {"prompt_uuid": "uuid-shell"}, + "answer": "allow", "_timestamp": 100.0, + }] + + def fake_search(query, *args, **kwargs): + # Honor prompt_uuid scoping like a real backend would. + want_uuid = query.get("extra_data.prompt_uuid") + out = [] + for d in answered_db: + if d.get("status") != query.get("status"): + continue + if want_uuid is not None and d["extra_data"].get("prompt_uuid") != want_uuid: + continue + out.append(d) + return out + + mock_engine = MagicMock() + mock_engine.search.side_effect = fake_search + backend = RemoteBackend(timeout=5, query_engine=mock_engine, poll_interval=5) + + # First (shell) layer resolves to its own answered "allow". + first = backend._poll_for_answer("session1", "permission", prompt_uuid="uuid-shell") + self.assertEqual(first, "allow") + + # Second (target) layer must NOT pick up the shell layer's "allow". + second = backend._poll_for_answer("session1", "permission", prompt_uuid="uuid-target") + self.assertIsNone(second) + + def test_poll_returns_newest_answered_doc(self): + """Defense in depth: resolve against the NEWEST answered doc by _timestamp.""" + from secator.ai.interactivity import RemoteBackend + mock_engine = MagicMock() + mock_engine.search.return_value = [ + {"answer": "stale", "_timestamp": 100.0}, + {"answer": "fresh", "_timestamp": 200.0}, ] + backend = RemoteBackend(timeout=60, query_engine=mock_engine, poll_interval=0.01) + result = backend._poll_for_answer("session1", "permission", prompt_uuid="abc-123") + self.assertEqual(result, "fresh") + + @patch('secator.ai.interactivity.sleep') + def test_ask_user_returns_on_second_poll(self, mock_sleep): + from secator.ai.interactivity import RemoteBackend + mock_engine = MagicMock() + # Query-aware: the follow-up answer poll (ai_type=="follow_up") returns the + # answer on the second call; the interleaved steer poll (ai_type=="steer") + # always returns nothing — so the steer-break never fires here. + answer_calls = {"n": 0} + + def search(query, limit=1): + if query.get("ai_type") == "steer": + return [] + answer_calls["n"] += 1 + return [] if answer_calls["n"] == 1 else [{"answer": "option B"}] + + mock_engine.search.side_effect = search backend = RemoteBackend(timeout=60, query_engine=mock_engine, poll_interval=5) result = backend.ask_user("What next?", [], "session1") @@ -116,6 +230,212 @@ def test_ask_user_returns_on_second_poll(self, mock_sleep): self.assertEqual(result["answer"], "option B") self.assertEqual(mock_sleep.call_count, 1) + @patch('secator.ai.interactivity.sleep') + def test_answer_in_final_window_is_not_lost_to_timeout(self, mock_sleep): + """M10: an answer landing in the last sleep window is returned, not lost. + + The poll loop sees only 'pending' until the loop exits, then the answer + appears. The final post-loop search must pick it up rather than abandon + the turn. + """ + from secator.ai.interactivity import RemoteBackend + answered_doc = [{"answer": "landed late", "_timestamp": 100.0}] + + def fake_search(query, *args, **kwargs): + # Answer only becomes visible AFTER the single poll iteration. + return list(answered_doc) if mock_sleep.call_count >= 1 else [] + + mock_engine = MagicMock() + mock_engine.search.side_effect = fake_search + mock_engine.update.return_value = 0 + backend = RemoteBackend(timeout=5, query_engine=mock_engine, poll_interval=5) + + result = backend.ask_user("What next?", [], "session1", prompt_uuid="abc-123") + + self.assertIsNotNone(result) + self.assertEqual(result["answer"], "landed late") + + @patch('secator.ai.interactivity.sleep') + def test_timeout_noop_flip_rereads_answer(self, mock_sleep): + """M10: if the timeout flip modifies 0 rows, re-read the answer.""" + from secator.ai.interactivity import RemoteBackend + # Empty during the loop AND at the first final search, then the answer + # appears right as we attempt the (no-op) flip. + searches = [[], [], [{"answer": "raced in", "_timestamp": 1.0}]] + mock_engine = MagicMock() + mock_engine.search.side_effect = lambda *a, **k: searches.pop(0) if searches else [] + mock_engine.update.return_value = 0 # nothing pending -> already answered + backend = RemoteBackend(timeout=5, query_engine=mock_engine, poll_interval=5) + + result = backend._poll_for_answer("session1", "permission", prompt_uuid="abc-123") + self.assertEqual(result, "raced in") + + def test_build_pending_prompt_expires_prior_pending(self): + """M10: starting a new prompt marks prior still-pending docs stale.""" + from secator.ai.interactivity import RemoteBackend + mock_engine = MagicMock() + backend = RemoteBackend(timeout=60, query_engine=mock_engine) + + backend.build_pending_prompt( + "Target x requires approval", ["allow", "deny"], "session1", + prompt_type="permission", permission_type="target", value="x", + prompt_uuid="uuid-new", + ) + + # An update flipping this session's pending docs to timed_out must fire. + mock_engine.update.assert_called_once() + flip_query, flip_update = mock_engine.update.call_args[0] + self.assertEqual(flip_query.get("_context.session_id"), "session1") + self.assertEqual(flip_query.get("status"), "pending") + self.assertEqual(flip_update, {"$set": {"status": "timed_out"}}) + + def test_expire_stale_pending_noop_without_engine(self): + """No query engine -> no crash, no update.""" + from secator.ai.interactivity import RemoteBackend + backend = RemoteBackend(timeout=60, query_engine=None) + backend._expire_stale_pending("session1") # must not raise + + def _permission_backend(self, answer): + """RemoteBackend whose poll resolves to `answer` for a shell prompt.""" + from secator.ai.interactivity import RemoteBackend + mock_engine = MagicMock() + mock_engine.search.return_value = [{"answer": answer, "_timestamp": 1.0}] + return RemoteBackend(timeout=60, query_engine=mock_engine, poll_interval=0.01) + + @staticmethod + def _shell_name_allowed(engine, cmd_name): + """True if the engine auto-allows this shell command NAME (no re-prompt). + + Asserted at the command-name layer (``_check_value``) rather than via + check_action() so the test does not depend on the safecmd/shfmt parser, + which is not present in every env. This is the exact layer a persisted + ``shell()`` session rule matches on. + """ + return engine._check_value("shell", cmd_name).decision == "allow" + + def test_allow_all_persists_session_rule_second_action_auto_allowed(self): + """M12: allow_all adds a session-scoped rule; a 2nd matching action needs no prompt.""" + from secator.ai.guardrails import PermissionEngine + engine = PermissionEngine(config={}) # no static rules: unknown cmd -> no auto-allow + backend = self._permission_backend("allow_all") + + # Pre-condition: with no rule, the command name is not pre-allowed. + self.assertFalse(self._shell_name_allowed(engine, "nmap")) + + result = backend.ask_user( + "Shell `nmap -sV` requires approval", ["deny", "allow", "allow_all"], + "session1", prompt_type="permission", engine=engine, + permission_type="shell", value="nmap -sV", prompt_uuid="u1", + ) + self.assertEqual(result["answer"], "allow") + # A session-scoped shell(nmap) pattern rule must now be present. + self.assertTrue( + any(rt == "shell" and "nmap" in patterns for rt, patterns in engine.runtime_allow), + "allow_all must persist a session-scoped shell(nmap) rule", + ) + # A SECOND, DIFFERENT nmap invocation is auto-allowed without a new prompt. + self.assertTrue(self._shell_name_allowed(engine, "nmap")) + + def test_single_allow_does_not_persist_rule_second_action_reprompts(self): + """M12/H9: single allow is one-shot — no rule added, a 2nd match re-prompts.""" + from secator.ai.guardrails import PermissionEngine + engine = PermissionEngine(config={}) + backend = self._permission_backend("allow") + + result = backend.ask_user( + "Shell `nmap -sV` requires approval", ["deny", "allow", "allow_all"], + "session1", prompt_type="permission", engine=engine, + permission_type="shell", value="nmap -sV", prompt_uuid="u1", + ) + self.assertEqual(result["answer"], "allow") + # No session rule was persisted -> a second matching action is not pre-allowed. + self.assertEqual(engine.runtime_allow, []) + self.assertFalse(self._shell_name_allowed(engine, "nmap")) + + def test_deny_unchanged_no_rule(self): + """deny returns deny and never touches runtime_allow.""" + from secator.ai.guardrails import PermissionEngine + engine = PermissionEngine(config={}) + backend = self._permission_backend("deny") + + result = backend.ask_user( + "Shell `nmap` requires approval", ["deny", "allow", "allow_all"], + "session1", prompt_type="permission", engine=engine, + permission_type="shell", value="nmap", prompt_uuid="u1", + ) + self.assertEqual(result["answer"], "deny") + self.assertEqual(engine.runtime_allow, []) + + +class TestRemoteBackendSteer(unittest.TestCase): + """Verify mid-flight steer draining + the blocked-wait break.""" + + def test_poll_steers_returns_and_consumes(self): + """poll_steers returns pending steer content and marks them consumed.""" + from secator.ai.interactivity import RemoteBackend + mock_engine = MagicMock() + mock_engine.search.return_value = [ + {"content": "actually focus on the API", "_timestamp": 2}, + {"content": "and skip port 80", "_timestamp": 1}, + ] + mock_engine.update = MagicMock() + backend = RemoteBackend(timeout=60, query_engine=mock_engine, poll_interval=0.01) + + steers = backend.poll_steers("session1") + + # Oldest-first by _timestamp + self.assertEqual(steers, ["and skip port 80", "actually focus on the API"]) + # Query scoped to pending steer docs for this session + search_query = mock_engine.search.call_args[0][0] + self.assertEqual(search_query.get("ai_type"), "steer") + self.assertEqual(search_query.get("status"), "pending") + self.assertEqual(search_query.get("_context.session_id"), "session1") + # Pending steers flipped to consumed (inject exactly once) + mock_engine.update.assert_called_once() + update_set = mock_engine.update.call_args[0][1] + self.assertEqual(update_set["$set"]["status"], "consumed") + + def test_poll_steers_no_pending_returns_empty(self): + from secator.ai.interactivity import RemoteBackend + mock_engine = MagicMock() + mock_engine.search.return_value = [] + backend = RemoteBackend(timeout=60, query_engine=mock_engine, poll_interval=0.01) + + self.assertEqual(backend.poll_steers("session1"), []) + # Nothing to consume when nothing is pending + mock_engine.update.assert_not_called() + + def test_poll_steers_robust_on_backend_error(self): + """A steer must never crash the run: backend errors return [].""" + from secator.ai.interactivity import RemoteBackend + mock_engine = MagicMock() + mock_engine.search.side_effect = RuntimeError("mongo down") + backend = RemoteBackend(timeout=60, query_engine=mock_engine, poll_interval=0.01) + + self.assertEqual(backend.poll_steers("session1"), []) + + def test_poll_steers_no_query_engine(self): + from secator.ai.interactivity import RemoteBackend + backend = RemoteBackend(timeout=60, query_engine=None, poll_interval=0.01) + self.assertEqual(backend.poll_steers("session1"), []) + + def test_steer_breaks_blocked_follow_up_wait(self): + """A steer arriving during a follow-up wait returns as the answer.""" + from secator.ai.interactivity import RemoteBackend + mock_engine = MagicMock() + # No follow-up answer ever; a steer arrives on the first poll. + mock_engine.search.side_effect = [ + [], # answered? no + [{"content": "change course now", "_timestamp": 1}], # poll_steers -> steer + ] + mock_engine.update = MagicMock() + backend = RemoteBackend(timeout=60, query_engine=mock_engine, poll_interval=0.01) + + result = backend.ask_user("What next?", [], "session1", prompt_uuid="uuid-1") + + # The steer content resolves the blocked wait (returned as the answer). + self.assertEqual(result["answer"], "change course now") + class TestCreateBackend(unittest.TestCase): """Verify create_backend factory.""" @@ -141,5 +461,173 @@ def test_unknown_returns_auto(self): self.assertIsInstance(backend, AutoBackend) +@unittest.skipUnless(HAS_AI, "ai addon required") +class TestRemoteTurnPendingDocCoverage(unittest.TestCase): + """H5: common remote turns (plain-chat reply, max-iter exit) must not poll on + prompt_uuid=None. Plain-chat must persist a proper pending doc and poll on its + real uuid; the max-iter terminal path must not strand a dangling pending doc.""" + + class _FakeHistory: + def add_user(self, *a, **k): + pass + + def count_tokens_by_role(self, model=None): + return {"total": 0} + + def to_messages(self, *a, **k): + return [] + + def test_plain_chat_remote_persists_pending_doc_and_polls_on_uuid(self): + """A plain-chat remote turn persists a pending follow_up doc with a + non-None prompt_uuid and polls scoped to THAT uuid (never None).""" + from secator.tasks.ai import ai as AiTask + from secator.ai.interactivity import RemoteBackend + from secator.output_types import Ai + + mock_engine = MagicMock() + mock_engine.search.return_value = [{"answer": "keep going", "_timestamp": 1.0}] + backend = RemoteBackend(timeout=60, query_engine=mock_engine, poll_interval=0.01) + + persisted = [] + fake_self = MagicMock() + fake_self.backend = backend + fake_self.session_id = "sess-chat" + fake_self.model = "gpt-4o" + fake_self.mode = "chat" + fake_self.encryptor = None + fake_self.max_iterations = 10 + fake_self.history = self._FakeHistory() + fake_self.add_result = lambda item, **kw: persisted.append(item) + + # plain-chat turn: no prompt_uuid passed in (this is the H5 path) + items = AiTask._prompt_and_redetect(fake_self, []) + + pend = [p for p in persisted if isinstance(p, Ai) and p.status == "pending"] + self.assertEqual(len(pend), 1, "exactly one pending doc must be persisted") + uuid_stamped = (pend[0].extra_data or {}).get("prompt_uuid") + self.assertTrue(uuid_stamped, "pending doc must carry a real (non-None) prompt_uuid") + self.assertEqual(pend[0].ai_type, "follow_up") + + # the poll must be scoped to THAT prompt's uuid — never None + search_query = mock_engine.search.call_args[0][0] + self.assertEqual(search_query.get("extra_data.prompt_uuid"), uuid_stamped) + # the answer resolved, so the loop continues (non-None items) rather than exiting + self.assertIsNotNone(items) + + def test_local_plain_chat_does_not_persist_pending_doc(self): + """Local (CLI) plain-chat must NOT create a pending doc — that is a + remote-channel concern only.""" + from secator.tasks.ai import ai as AiTask + from secator.ai.interactivity import CLIBackend + from secator.output_types import Ai + + backend = MagicMock(spec=CLIBackend) + backend.ask_user.return_value = {"answer": "do x"} + + persisted = [] + fake_self = MagicMock() + fake_self.backend = backend + fake_self.session_id = "sess-local" + fake_self.model = "gpt-4o" + fake_self.mode = "chat" + fake_self.encryptor = None + fake_self.max_iterations = 10 + fake_self.history = self._FakeHistory() + fake_self.add_result = lambda item, **kw: persisted.append(item) + + AiTask._prompt_and_redetect(fake_self, []) + + pend = [p for p in persisted if isinstance(p, Ai) and p.status == "pending"] + self.assertEqual(pend, [], "local backend must not persist a pending doc") + + @patch("secator.query.QueryEngine") + @patch("secator.tasks.ai.init_llm") + @patch("secator.tasks.ai.call_llm") + def test_remote_max_iter_does_not_strand_pending_doc(self, mock_call_llm, mock_init, mock_qe_cls): + """At remote max-iter after tool work, the loop ends cleanly: it does not + enter the follow-up poll and does not persist a dangling pending doc.""" + from secator.tasks.ai import ai as AiTask + from secator.ai.interactivity import RemoteBackend + from secator.output_types import Ai, Info + + mock_call_llm.return_value = { + "content": "working", "tool_calls": [object()], "usage": {"tokens": 100, "cost": 0.001}, + } + + persisted = [] + prompt_calls = [] + backend = RemoteBackend(timeout=60, query_engine=MagicMock(), poll_interval=0.01) + + fake_self = MagicMock() + fake_self.backend = backend + fake_self.session_id = "sess-max" + fake_self.model = "gpt-4o" + fake_self.mode = "chat" + fake_self.max_iterations = 1 + fake_self.interactive = "remote" + fake_self.is_subagent = False + fake_self.inputs = [] + fake_self.context = {} + fake_self.scope = "workspace" + fake_self.results = [] + fake_self.max_workers = 3 + fake_self.encryptor = None + fake_self.dry_run = False + fake_self.verbose = False + fake_self._sync = False + fake_self.temp = 0.7 + fake_self.api_base = "" + fake_self.api_key = "" + fake_self.tool_schemas = [] + fake_self.max_tokens_total = 100000 + fake_self.permission_engine = MagicMock() + fake_self.history = self._FakeHistory() + fake_self.add_result = lambda item, **kw: persisted.append(item) + + def _empty_gen(*a, **k): + return + yield # pragma: no cover - make it a generator + + fake_self._summarize_auto = _empty_gen + fake_self._summarize_user = _empty_gen + fake_self._drain_history_usage = lambda: None + fake_self._account_usage = lambda u: None + # _add_assistant_to_history now returns the litellm message dict (the + # caller feeds it to cap_message(...) for persistence); the stub must + # match that contract instead of returning None -- a bare None broke + # cap_message's dict(msg) call and masked this test's real assertions + # behind a swallowed exception. + fake_self._add_assistant_to_history = lambda c, t: {"role": "assistant", "content": c} + fake_self._save_history = lambda: None + + def _fake_process(tool_calls, ctx): + return [{"action": "shell", "tool_call_id": "t", "tool_call_name": "run_shell"}] + yield # pragma: no cover + + fake_self._process_tool_calls = _fake_process + + def _fake_dispatch(actions, ctx): + return {"follow_up_choices": None, "stop_reason": None, "follow_up_prompt_uuid": None} + yield # pragma: no cover + + fake_self._dispatch_and_collect = _fake_dispatch + + def _track_prompt(choices, prompt_uuid=None): + prompt_calls.append((choices, prompt_uuid)) + return [] + + fake_self._prompt_and_redetect = _track_prompt + + items = list(AiTask._run_loop(fake_self)) + + self.assertEqual(prompt_calls, [], "remote max-iter must not enter the follow-up poll") + pend = [p for p in persisted if isinstance(p, Ai) and getattr(p, "status", None) == "pending"] + self.assertEqual(pend, [], "remote max-iter must not persist a dangling pending doc") + self.assertTrue( + any(isinstance(it, Info) and "max iterations" in it.message.lower() for it in items), + "loop must end via the terminal 'reached max iterations' tail", + ) + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/test_ai_loop.py b/tests/unit/test_ai_loop.py index 39e179cfd..195a62c4f 100644 --- a/tests/unit/test_ai_loop.py +++ b/tests/unit/test_ai_loop.py @@ -24,6 +24,21 @@ from secator.output_types import Ai +def _fake_command_runner(output): + """Build a fake `command` runner instance for patching secator.tasks.command.command. + + The AI shell path now runs commands as the `command` task (not subprocess.run), so + these E2E flows patch the class at its source and hand back an instance exposing the + fields _handle_shell reads: `.output` (captured stdout), `.id`, `.status`, a callable + `.run()`, and a settable `.max_timeout`. Mirrors the fake used in TestHandleShell. + """ + fake = MagicMock() + fake.output = output + fake.id = "task_fake" + fake.status = "SUCCESS" + return fake + + def _make_tool_call(name, args, tc_id=None): """Create a mock litellm tool call object.""" tc = MagicMock() @@ -307,6 +322,82 @@ def test_stop_yields_ai_stopped(self): self.assertIn("completed", ai_results[0].content) +# ============================================================================= +# UNIT TESTS: Remote follow-up persistence (status + top-level choices) +# ============================================================================= + +@unittest.skipUnless(HAS_AI, "ai addon required") +class TestRemoteFollowUpPersistence(unittest.TestCase): + """In remote mode, the single persisted follow-up doc must be renderable: + status=="pending" + non-empty top-level `choices` (what the web UI reads).""" + + def _run_dispatch(self, backend): + """Drive the real ai._dispatch_and_collect with a minimal fake self. + + Returns (yielded_items, persisted_items) where persisted_items are what + add_result() received (i.e. what the mongodb on_item hook would persist). + """ + from secator.tasks.ai import ai as AiTask + + choices = ["Fuzz parameters", "Run nuclei", "Deep crawl"] + follow_up = Ai( + content="Presenting actionable next steps", + ai_type="follow_up", + extra_data={"choices": choices}, + _context={"tool_call_id": "tc_fu", "tool_call_name": "follow_up"}, + ) + + persisted = [] + + class _FakeHistory: + def get_action_budget(self, model): + return 10000 + + def add_tool_result(self, *a, **k): + pass + + fake_self = MagicMock() + fake_self.backend = backend + fake_self.session_id = "sess-123" + fake_self.model = "test-model" + fake_self.reports_folder = None + fake_self.history = _FakeHistory() + fake_self.add_result = lambda item, **kw: persisted.append(item) + + ctx = MagicMock() + ctx.results = [] + + def _fake_dispatch_action(action, c): + yield follow_up + + with patch("secator.tasks.ai.safe_dispatch_action", _fake_dispatch_action): + gen = AiTask._dispatch_and_collect(fake_self, [{"tool_call_id": "tc_fu"}], ctx) + yielded = list(gen) + return yielded, persisted, follow_up + + def test_remote_follow_up_persisted_pending_with_choices(self): + backend = RemoteBackend(timeout=60, query_engine=MagicMock()) + yielded, persisted, follow_up = self._run_dispatch(backend) + + # Exactly one follow_up Ai is persisted (no duplicate display + pending docs). + fu_docs = [p for p in persisted if isinstance(p, Ai) and p.ai_type == "follow_up"] + self.assertEqual(len(fu_docs), 1) + doc = fu_docs[0] + self.assertEqual(doc.status, "pending") + self.assertEqual(doc.choices, ["Fuzz parameters", "Run nuclei", "Deep crawl"]) + self.assertEqual(doc.session_id, "sess-123") + # Same object → single doc by _uuid. + self.assertIs(doc, follow_up) + + def test_local_follow_up_not_stamped_pending(self): + """CLI/local mode must NOT stamp status=pending (drives the TUI menu directly).""" + backend = CLIBackend() + yielded, persisted, follow_up = self._run_dispatch(backend) + fu_docs = [p for p in persisted if isinstance(p, Ai) and p.ai_type == "follow_up"] + self.assertEqual(len(fu_docs), 1) + self.assertNotEqual(fu_docs[0].status, "pending") + + # ============================================================================= # UNIT TESTS: Backend and tool schema behavior # ============================================================================= @@ -386,14 +477,19 @@ def test_path_rule_added(self): self.assertEqual(check.decision, "allow") def test_full_remote_permission_flow(self): - """Full flow: RemoteBackend polls, gets 'allow', adds rules.""" + """Full flow: RemoteBackend polls, gets 'allow_all', persists a session rule. + + Per M12/H9 a single 'allow' is a one-shot that adds NO rule (next match + re-prompts); only 'allow_all' persists a session-scoped runtime allow. The + backend still maps both to {"answer": "allow"} for the caller. + """ mock_engine = MagicMock() - mock_engine.search.return_value = [{"answer": "allow"}] + mock_engine.search.return_value = [{"answer": "allow_all"}] engine = PermissionEngine(_make_permission_config(), targets=["10.0.0.1"], workspace="/tmp/ws") backend = RemoteBackend(timeout=60, query_engine=mock_engine, poll_interval=0.01) result = backend.ask_user( - "Allow shell: python3?", ["allow", "deny"], "sess1", + "Allow shell: python3?", ["allow", "allow_all", "deny"], "sess1", prompt_type="permission", permission_type="shell", value="python3 exploit.py", engine=engine ) @@ -441,8 +537,8 @@ def mock_approve(command, reason="", interactive=True): mock_shell_prompt.assert_called_once() # Dispatch the approved action - with patch('secator.ai.actions.subprocess.run') as mock_run: - mock_run.return_value = MagicMock(stdout="exploit output\n", stderr="") + with patch('secator.tasks.command.command') as mock_cmd: + mock_cmd.return_value = _fake_command_runner("exploit output") results1 = list(dispatch_action(shell_action, ctx)) # Verify shell output was produced @@ -522,8 +618,8 @@ def mock_ask(question="", choices=None, session_id="", prompt_type="", **kwargs) self.assertIsNone(denial, f"Expected approval but got: {denial}") # Dispatch the approved action - with patch('secator.ai.actions.subprocess.run') as mock_run: - mock_run.return_value = MagicMock(stdout="exploit output\n", stderr="") + with patch('secator.tasks.command.command') as mock_cmd: + mock_cmd.return_value = _fake_command_runner("exploit output") results1 = list(dispatch_action(shell_action, ctx)) ai_results = [r for r in results1 if isinstance(r, Ai)] @@ -629,8 +725,8 @@ def test_allowed_command_runs_in_auto_mode(self): denial, warnings = check_guardrails(action, ctx) self.assertIsNone(denial) - with patch('secator.ai.actions.subprocess.run') as mock_run: - mock_run.return_value = MagicMock(stdout="response data", stderr="") + with patch('secator.tasks.command.command') as mock_cmd: + mock_cmd.return_value = _fake_command_runner("response data") results = list(dispatch_action(action, ctx)) ai_results = [r for r in results if isinstance(r, Ai)] @@ -743,7 +839,7 @@ def test_multi_turn_local_loop(self): engine = PermissionEngine(_make_permission_config(), targets=["10.0.0.1"], workspace="/tmp/ws") ctx = _make_ctx(interactive="local", engine=engine) history = ChatHistory() - history.add_system("You are a pentester.") + history.set_system("You are a pentester.") history.add_user("Scan 10.0.0.1") # --- Turn 1: LLM returns shell tool call --- @@ -763,8 +859,8 @@ def test_multi_turn_local_loop(self): self.assertIsNone(denial) # Dispatch - with patch('secator.ai.actions.subprocess.run') as mock_subprocess: - mock_subprocess.return_value = MagicMock(stdout="scan results\n", stderr="") + with patch('secator.tasks.command.command') as mock_cmd: + mock_cmd.return_value = _fake_command_runner("scan results") results1 = list(dispatch_action(action, ctx)) self.assertTrue(any(isinstance(r, Ai) and r.ai_type == "shell_output" for r in results1)) @@ -832,7 +928,7 @@ def mock_ask(question="", choices=None, session_id="", prompt_type="", **kwargs) ctx = _make_ctx(interactive="remote", backend=mock_backend, engine=engine, session_id="remote-e2e") history = ChatHistory() - history.add_system("You are a pentester.") + history.set_system("You are a pentester.") history.add_user("Scan 10.0.0.1") # --- Turn 1: Shell command needing remote permission --- @@ -847,8 +943,8 @@ def mock_ask(question="", choices=None, session_id="", prompt_type="", **kwargs) self.assertIsNone(denial, f"Expected remote approval but got: {denial}") # Dispatch (mock subprocess only around dispatch_action) - with patch('secator.ai.actions.subprocess.run') as mock_subprocess: - mock_subprocess.return_value = MagicMock(stdout="scan results\n", stderr="") + with patch('secator.tasks.command.command') as mock_cmd: + mock_cmd.return_value = _fake_command_runner("scan results") results1 = list(dispatch_action(action, ctx)) self.assertTrue(any(isinstance(r, Ai) and r.ai_type == "shell_output" for r in results1)) @@ -891,7 +987,7 @@ def test_multi_turn_auto_loop(self): """ ctx = _make_ctx(interactive="auto") history = ChatHistory() - history.add_system("You are a pentester.") + history.set_system("You are a pentester.") history.add_user("Scan 10.0.0.1") # --- Turn 1: Unknown command → blocked --- @@ -918,8 +1014,8 @@ def test_multi_turn_auto_loop(self): denial2, _ = check_guardrails(action2, ctx) self.assertIsNone(denial2, "Allowed command should pass in auto mode") - with patch('secator.ai.actions.subprocess.run') as mock_subprocess: - mock_subprocess.return_value = MagicMock(stdout="output\n", stderr="") + with patch('secator.tasks.command.command') as mock_cmd: + mock_cmd.return_value = _fake_command_runner("output") results = list(dispatch_action(action2, ctx)) self.assertTrue(any(isinstance(r, Ai) and r.ai_type == "shell_output" for r in results)) @@ -939,5 +1035,201 @@ def test_multi_turn_auto_loop(self): self.assertIsNotNone(stop_reason) +@unittest.skipUnless(HAS_AI, "ai addon required") +class TestLoopResilientToActionErrors(unittest.TestCase): + """A Python error during an action dispatch must NOT kill the main loop. + + It must be caught, turned into an Error item fed back to the LLM as that + tool call's result, and the loop must continue. + """ + + def test_safe_dispatch_catches_exception_and_feeds_back(self): + """safe_dispatch_action converts a raised Exception into an Error item + carrying the action's tool_call_id, instead of propagating.""" + from secator.ai.actions import safe_dispatch_action + from secator.output_types import Error + + ctx = _make_ctx(interactive="auto") + action = { + "action": "shell", + "command": "curl http://10.0.0.1", + "tool_call_id": "tc_err", + "tool_call_name": "run_shell", + } + + # Make the shell handler raise the exact failure from the spec. + def _boom(*a, **k): + raise TypeError("'str' object is not a mapping") + + with patch("secator.ai.actions._handle_shell", _boom): + # Must NOT raise. + results = list(safe_dispatch_action(action, ctx)) + + errors = [r for r in results if isinstance(r, Error)] + self.assertEqual(len(errors), 1, "expected exactly one Error item") + err = errors[0] + # LLM-facing feedback phrasing + the exception type/message. + self.assertIn("Action failed with error", err.message) + self.assertIn("TypeError", err.message) + self.assertIn("'str' object is not a mapping", err.message) + self.assertIn("try again", err.message.lower()) + # Attributed to the failing tool call so it groups into that tool result. + self.assertEqual(err._context.get("tool_call_id"), "tc_err") + self.assertEqual(err._context.get("tool_call_name"), "run_shell") + + def test_does_not_catch_keyboardinterrupt(self): + """Control-flow exceptions (BaseException) must propagate, not be swallowed.""" + from secator.ai.actions import safe_dispatch_action + + ctx = _make_ctx(interactive="auto") + action = {"action": "shell", "command": "x", "tool_call_id": "tc", "tool_call_name": "run_shell"} + + def _interrupt(*a, **k): + raise KeyboardInterrupt() + yield # pragma: no cover - make it a generator + + with patch("secator.ai.actions._handle_shell", _interrupt): + with self.assertRaises(KeyboardInterrupt): + list(safe_dispatch_action(action, ctx)) + + def test_dispatch_and_collect_continues_and_feeds_history(self): + """Drive the real _dispatch_and_collect: a raising action yields an Error, + appends an error result to history (LLM-visible), and does NOT raise.""" + from secator.tasks.ai import ai as AiTask + from secator.output_types import Error + + tool_results = [] # (name, tc_id, content) tuples appended to history + + class _FakeHistory: + def get_action_budget(self, model): + return 10000 + + def add_tool_result(self, name, tc_id, content): + tool_results.append((name, tc_id, content)) + + persisted = [] + fake_self = MagicMock() + fake_self.backend = CLIBackend() + fake_self.session_id = "sess-err" + fake_self.model = "test-model" + fake_self.reports_folder = None + fake_self.encryptor = None + fake_self.history = _FakeHistory() + fake_self.add_result = lambda item, **kw: persisted.append(item) + + ctx = MagicMock() + ctx.results = [] + + action = { + "action": "shell", + "command": "curl http://10.0.0.1", + "tool_call_id": "tc_err", + "tool_call_name": "run_shell", + } + + def _boom(*a, **k): + raise TypeError("'str' object is not a mapping") + yield # pragma: no cover + + with patch("secator.ai.actions._handle_shell", _boom): + # Single action → safe_dispatch_action path. Must not raise. + gen = AiTask._dispatch_and_collect(fake_self, [action], ctx) + yielded = list(gen) + + # An Error item was yielded to the caller (visible in console / persisted). + errors = [r for r in yielded if isinstance(r, Error)] + self.assertEqual(len(errors), 1) + + # The error reached the LLM-visible history as this tool call's result. + self.assertEqual(len(tool_results), 1) + name, tc_id, content = tool_results[0] + self.assertEqual(tc_id, "tc_err") + self.assertIn("error", content.lower()) + self.assertIn("'str' object is not a mapping", content) + + +# ============================================================================= +# Mid-flight steering: _drain_steers injects pending steers into history +# ============================================================================= + +@unittest.skipUnless(HAS_AI, "ai addon required") +class TestDrainSteers(unittest.TestCase): + """The loop's `_drain_steers` drains pending steers and injects them. + + A pending `ai_type:"steer"` doc (user message sent WHILE the agent runs) must + be drained at the loop checkpoint, appended to the LLM history as a + `[User interjected]: …` user message, echoed as a steer Ai item, and marked + consumed — without breaking the loop or the existing follow-up flow. + """ + + def _make_task(self, backend, history=None): + """Build a minimal `ai` task with only what _drain_steers reads.""" + from secator.tasks.ai import ai + task = object.__new__(ai) + task.backend = backend + task.session_id = "steer-sess" + task.encryptor = None + task.history = history or ChatHistory() + task.debug = lambda *a, **k: None + return task + + def test_steer_drained_injected_and_consumed(self): + from secator.ai.interactivity import RemoteBackend + mock_engine = MagicMock() + mock_engine.search.return_value = [ + {"content": "actually focus on the API", "_timestamp": 1}, + ] + mock_engine.update = MagicMock() + backend = RemoteBackend(timeout=60, query_engine=mock_engine, poll_interval=0.01) + task = self._make_task(backend) + + yielded = list(task._drain_steers()) + + # Injected into history as a user "interjected" message. + user_msgs = [m for m in task.history.to_messages() if m["role"] == "user"] + self.assertEqual(len(user_msgs), 1) + self.assertEqual(user_msgs[-1]["content"], "[User interjected]: actually focus on the API") + + # No echo doc is yielded: the API's pending steer doc is itself the + # persisted transcript entry, so a second steer Ai would double-render. + self.assertEqual(yielded, []) + + # Marked consumed so it injects exactly once. + update_set = mock_engine.update.call_args[0][1] + self.assertEqual(update_set["$set"]["status"], "consumed") + + def test_no_steer_is_noop_and_preserves_loop(self): + """No pending steer -> nothing injected, history untouched (loop intact).""" + from secator.ai.interactivity import RemoteBackend + mock_engine = MagicMock() + mock_engine.search.return_value = [] + backend = RemoteBackend(timeout=60, query_engine=mock_engine, poll_interval=0.01) + history = ChatHistory() + history.add_user("original prompt") + task = self._make_task(backend, history=history) + + yielded = list(task._drain_steers()) + + self.assertEqual(yielded, []) + user_msgs = [m for m in task.history.to_messages() if m["role"] == "user"] + self.assertEqual([m["content"] for m in user_msgs], ["original prompt"]) + + def test_non_remote_backend_is_noop(self): + """Local/auto backends have no channel -> drain is a no-op (no crash).""" + task = self._make_task(create_backend("auto")) + self.assertEqual(list(task._drain_steers()), []) + + def test_steer_poll_error_never_crashes_loop(self): + """A backend error during drain is swallowed (run must not crash).""" + from secator.ai.interactivity import RemoteBackend + mock_engine = MagicMock() + mock_engine.search.side_effect = RuntimeError("mongo down") + backend = RemoteBackend(timeout=60, query_engine=mock_engine, poll_interval=0.01) + task = self._make_task(backend) + # Should not raise, yields nothing, history untouched. + self.assertEqual(list(task._drain_steers()), []) + self.assertEqual(task.history.to_messages(), []) + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/test_ai_prompts.py b/tests/unit/test_ai_prompts.py index 166276853..e05513c1f 100644 --- a/tests/unit/test_ai_prompts.py +++ b/tests/unit/test_ai_prompts.py @@ -101,6 +101,26 @@ def test_get_system_prompt_exploit(self): self.assertIn("exploitation verification specialist", prompt) self.assertIn("proof-of-concept", prompt) + def test_get_system_prompt_exploit_no_leftover_placeholders(self): + """D2: exploit renders fully — no unresolved ${include} or template $vars. + + (Literal Mongo operators like $in/$regex and example secrets like $API_KEY + are content, not Template vars, so we check the template names explicitly.) + """ + import re + prompt = get_system_prompt("exploit") + # All ${include} directives resolved (load_prompt) and $var substitutions done. + self.assertEqual(re.findall(r'\$\{\w+\}', prompt), [], "unresolved ${include} in exploit prompt") + template_vars = [ + "library_reference", "discovery", "common", "queries", "findings", + "arsenal", "guardrails", "isolation", "exploitation_report", + "workspace_path", "query_types", "output_types_reference", + ] + leftover = [v for v in template_vars if f"${v}" in prompt] + self.assertEqual(leftover, [], f"unresolved template vars in exploit prompt: {leftover}") + # uses the exploit template, not attack/chat + self.assertIn("exploitation verification specialist", prompt) + def test_get_system_prompt_attack_has_library_reference(self): prompt = get_system_prompt("attack") self.assertIn('', prompt) @@ -281,6 +301,33 @@ def test_common_rules_has_no_shouting(self): self.assertNotIn("NEVER INVENT", COMMON_RULES) self.assertNotIn("ALWAYS provide", COMMON_RULES) + # === Template-drift regression tests (D1) === + + def test_rendered_prompts_have_no_unsubstituted_template_vars(self): + """Rendered prompts must not leak $query_types / $output_types_reference (D1).""" + for mode in ("attack", "chat", "exploit"): + prompt = get_system_prompt(mode) + self.assertNotIn("$query_types", prompt, f"$query_types leaked in {mode!r} prompt") + self.assertNotIn("$output_types_reference", prompt, f"$output_types_reference leaked in {mode!r} prompt") + + def test_rendered_prompts_substitute_query_types_from_registry(self): + """$query_types renders to the real FINDING_TYPES names, not a placeholder.""" + from secator.ai.prompts import build_query_types + expected = build_query_types() + self.assertIn("vulnerability", expected) + for mode in ("attack", "chat", "exploit"): + self.assertIn(expected, get_system_prompt(mode)) + + def test_rendered_prompts_have_no_phantom_run_query_tool(self): + """Examples must call the real query_workspace tool, never a phantom run_query (D1).""" + from secator.ai.tools import TOOL_ACTION_MAP + self.assertEqual(TOOL_ACTION_MAP["query_workspace"], "query") + self.assertNotIn("run_query", TOOL_ACTION_MAP) + for mode in ("attack", "chat", "exploit"): + prompt = get_system_prompt(mode) + self.assertNotIn("run_query", prompt, f"phantom run_query in {mode!r} prompt") + self.assertIn("query_workspace", prompt) + if __name__ == '__main__': unittest.main() diff --git a/tests/unit/test_ai_resilience.py b/tests/unit/test_ai_resilience.py new file mode 100644 index 000000000..06ae7283e --- /dev/null +++ b/tests/unit/test_ai_resilience.py @@ -0,0 +1,304 @@ +"""Resiliency harness: fake the LLM and feed the agent loop every kind of weird +response, asserting the invariant that a malformed LLM response is handled +**turn-locally** and the loop **survives and continues** — never aborts the whole +session via the top-level catch-all, and never raises out of the loop. + +Why "survives and continues" and not just "doesn't crash": `_run_loop` already +wraps each iteration in `try/except Exception -> Error.from_exception; return` +(ai.py). So an unhandled exception won't kill the worker — but it ABORTS the +entire conversation. The resilient path instead rejects the bad tool call as a +clean tool-result error (so the model can retry) and keeps looping. We detect the +difference by counting `call_llm` invocations: a tool-call turn that is handled +turn-locally forces another iteration (the loop asks the model again), so +`call_llm` is called at least twice. An abort stops at one. + +Two layers: + 1. TestWeirdToolCalls — a curated table (regression coverage for the bugs we've + hit + adjacent ones: stringified opts/query, broken JSON, wrong-type args, + unknown tools, missing fields, mixed batches, ...). + 2. TestMalformedArgFuzzer — random malformed tool-call arguments across all + tools; seeded so any failure reproduces. +""" +import contextlib +import json +import random +import types +import unittest +from unittest.mock import patch + +from secator.definitions import ADDONS_ENABLED + +HAS_AI = ADDONS_ENABLED.get('ai', False) + +if HAS_AI: + from secator.tasks.ai import ai + from secator.ai.history import ChatHistory + from secator.ai.interactivity import create_backend + from secator.output_types import Error + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + +def _tc(name, args, call_id="tc1"): + """A litellm-shaped tool_call. ``args`` may be a dict/list (JSON-dumped, the + normal case) or a raw string — pass a string to inject malformed/weird + arguments verbatim, exactly as a misbehaving model would.""" + arguments = json.dumps(args) if isinstance(args, (dict, list)) else args + return types.SimpleNamespace(id=call_id, function=types.SimpleNamespace(name=name, arguments=arguments)) + + +def _resp(content=None, tool_calls=None, usage="default"): + """A call_llm() return dict.""" + if usage == "default": + usage = {"tokens": 100, "cost": 0.001} + return {"content": content, "tool_calls": tool_calls or [], "usage": usage} + + +def _make_loop_task(): + """A bare `ai` task carrying exactly the state `_run_loop` reads, with + dry_run+dangerous so actions neither execute real tools nor need the guardrail + shell parser — this isolates arg-handling / dispatch resilience.""" + task = ai.__new__(ai) + task.context = {"ai_tokens": 0, "ai_prompt_tokens": 0, "ai_completion_tokens": 0, "ai_cost": 0.0} + task.history = ChatHistory() + task.inputs = [] + task.model = "test-model" + task.intent_model = "test-model" + task.temp = 0.7 + task.api_base = None + task.api_key = "key" + task.max_iterations = 6 + task.max_tokens_total = 100000 + task.max_workers = 1 + task.is_subagent = True + task.verbose = False + task.dry_run = True + task.mode = "chat" + task.scope = "workspace" # + empty workspace_id -> query short-circuits, no real search + task.results = [] + task.encryptor = None + task.tool_schemas = [] + task.permission_engine = None + task.dangerous = True + task.interactive = "auto" + task._sync = True + task.session_id = "s" + task.async_tasks = False + task.context_warnings = False + task._reports_folder = None + task.system_prompt = "SYS" + task.debug = lambda *a, **k: None + task.add_result = lambda *a, **k: None + task.print_item = False + task.print_line = False + task.backend = create_backend("auto") + return task + + +@contextlib.contextmanager +def _driven(task, weird_responses): + """Patch call_llm to emit the weird responses, then a terminating content-only + response forever; stub the heavy collaborators; and capture the loop's + top-level abort signal. + + The abort signal is `Error.from_exception`, which `_run_loop` calls at exactly + one site (ai.py) — its `except Exception -> Error.from_exception(e); return` + catch-all. `safe_dispatch_action` (per-action resilience) uses plain `Error(...)`, + not `from_exception`, and `stop`/`follow_up` end the loop without raising — so a + captured `from_exception` means an UNHANDLED exception aborted the session: a + resilience failure, distinct from a clean end.""" + seq = list(weird_responses) + state = {"calls": 0, "aborted_with": None} + + def _next(*a, **k): + state["calls"] += 1 + return seq.pop(0) if seq else _resp(content="__done__", tool_calls=[]) + + real_from_exc = Error.from_exception + + def _capture(exc, *a, **k): + state["aborted_with"] = exc + return real_from_exc(exc, *a, **k) + + with contextlib.ExitStack() as stack: + stack.enter_context(patch('secator.tasks.ai.call_llm', side_effect=_next)) + stack.enter_context(patch('secator.tasks.ai.get_context_window', return_value=8000)) + stack.enter_context(patch('secator.ai.history.get_context_window', return_value=8000)) + stack.enter_context(patch('secator.tasks.ai.save_history')) + stack.enter_context(patch.object(type(task), 'reports_folder', property(lambda self: None))) + stack.enter_context(patch.object(ai, '_summarize_auto', return_value=iter(()))) + stack.enter_context(patch.object(ai, '_summarize_user', return_value=iter(()))) + stack.enter_context(patch('secator.tasks.ai.Error.from_exception', side_effect=_capture)) + # content-only turns exit (no follow-up loop); tool-call turns still continue + stack.enter_context(patch.object(ai, '_prompt_and_redetect', return_value=None)) + yield state + + +def _run(task, weird_responses): + """Drive the real _run_loop; return (yielded_items, call_count, aborted_with).""" + with _driven(task, weird_responses) as state: + items = list(task._run_loop()) + return items, state["calls"], state["aborted_with"] + + +# --------------------------------------------------------------------------- +# 1. Curated weird tool-call responses +# --------------------------------------------------------------------------- + +# Each entry: (label, tool_call). The loop gets ONE turn with this tool call, then +# a terminating content turn. A resilient loop rejects/handles the bad call and +# asks the model again -> call_llm invoked >= 2. +WEIRD_TOOL_CALLS = [ + # --- stringified object/array args (provider quirk; #1273/#1275) --- + ("stringified_opts", _tc("run_task", '{"name":"nmap","targets":["10.0.0.1"],"opts":"{\\"session_name\\":\\"x\\"}"}')), + ("stringified_query", _tc("query_workspace", '{"query":"{\\"_type\\":\\"url\\"}"}')), + ("stringified_targets", _tc("run_task", '{"name":"nmap","targets":"10.0.0.1"}')), + ("stringified_choices", _tc("follow_up", '{"reason":"pick","choices":"[\\"a\\",\\"b\\"]"}')), + # --- valid JSON, wrong shape --- + ("query_as_list", _tc("query_workspace", {"query": ["_type", "url"]})), + ("opts_as_int", _tc("run_task", {"name": "nmap", "targets": ["x"], "opts": 5})), + ("targets_as_number", _tc("run_task", {"name": "nmap", "targets": 12345})), + ("args_top_level_int", _tc("run_task", "12345")), + ("args_top_level_array", _tc("run_task", '["nmap","10.0.0.1"]')), + ("args_top_level_string", _tc("run_task", '"just a string"')), + # --- invalid JSON --- + ("broken_json_unbalanced", _tc("run_shell", '{"command": "curl -s x" ')), + ("broken_json_trailing", _tc("run_task", '{"name":"nmap",}')), + ("empty_string_args", _tc("run_shell", '')), + ("garbage_args", _tc("run_task", 'not json at all')), + # --- missing / empty fields --- + ("empty_object_args", _tc("run_task", {})), + ("missing_name", _tc("run_task", {"targets": ["x"]})), + ("shell_missing_command", _tc("run_shell", {})), + ("null_values", _tc("run_task", {"name": None, "targets": None, "opts": None})), + # --- unknown / nonsense tool --- + ("unknown_tool", _tc("delete_everything", {})), + ("unknown_tool_bad_args", _tc("../../etc/passwd", 'weird')), + # --- add_finding malformations --- + ("add_finding_str_data", _tc("add_finding", {"finding_type": "vulnerability", "data": "not-a-dict"})), + ("add_finding_no_type", _tc("add_finding", {"data": {"name": "x"}})), + # --- deep nesting / large --- + ("deeply_nested_opts", _tc("run_task", {"name": "nmap", "targets": ["x"], "opts": {"a": {"b": {"c": {"d": 1}}}}})), +] + + +@unittest.skipUnless(HAS_AI, 'ai addon required') +class TestWeirdToolCalls(unittest.TestCase): + """A malformed tool call must be handled turn-locally and the loop must + SURVIVE and continue (call_llm invoked again), never abort the session.""" + + def _assert_survives(self, label, tool_call): + task = _make_loop_task() + try: + _items, _n, aborted = _run(task, [_resp(tool_calls=[tool_call])]) + except Exception as e: # noqa: BLE001 - the whole point is nothing escapes + self.fail(f"[{label}] raised out of the loop: {type(e).__name__}: {e}") + self.assertIsNone( + aborted, + f"[{label}] weird tool call aborted the session via the top-level catch-all: " + f"{type(aborted).__name__ if aborted else None}: {aborted}") + + +def _make_weird_test(label, tool_call): + def test(self): + self._assert_survives(label, tool_call) + test.__name__ = f"test_{label}" + return test + + +for _label, _tool_call in WEIRD_TOOL_CALLS: + setattr(TestWeirdToolCalls, f"test_{_label}", _make_weird_test(_label, _tool_call)) + + +# --------------------------------------------------------------------------- +# 2. Non-tool-call weird responses (tailored expectations) +# --------------------------------------------------------------------------- + +@unittest.skipUnless(HAS_AI, 'ai addon required') +class TestWeirdContentResponses(unittest.TestCase): + + def test_single_empty_response_recovers(self): + """One empty response (no content, no tools) -> Warning, then continues.""" + task = _make_loop_task() + items, n, aborted = _run(task, [_resp(content=None, tool_calls=[])]) + self.assertIsNone(aborted) + self.assertGreaterEqual(n, 2) # recovered and asked again + + def test_three_empty_responses_stop_cleanly(self): + """Three consecutive empties stop with a clean Error (intended), no abort.""" + task = _make_loop_task() + items, n, aborted = _run(task, [_resp(content=None, tool_calls=[]) for _ in range(3)]) + self.assertIsNone(aborted) # a deliberate Error(485), not the catch-all + self.assertTrue(any(getattr(i, '_type', '') == 'error' for i in items)) + + def test_missing_usage_does_not_crash(self): + """usage=None must not crash accounting.""" + task = _make_loop_task() + items, n, aborted = _run(task, [_resp(content="hi", tool_calls=[], usage=None)]) + self.assertIsNone(aborted) + self.assertEqual(task.context["ai_tokens"], 0) + + def test_huge_content_does_not_crash(self): + task = _make_loop_task() + items, n, aborted = _run(task, [_resp(content="A" * 500_000, tool_calls=[])]) + self.assertIsNone(aborted) + + +# --------------------------------------------------------------------------- +# 3. Fuzzer: random malformed arguments across every tool +# --------------------------------------------------------------------------- + +_TOOL_NAMES = ["run_task", "run_workflow", "run_shell", "query_workspace", "follow_up", "add_finding", "stop"] + + +def _random_weird_arguments(rng): + """Produce a plausibly-broken `tool_call.arguments` string a model might emit.""" + kind = rng.choice([ + "valid_int", "valid_array", "valid_string", "unbalanced", "trailing_comma", + "empty", "not_json", "stringified_nested", "wrong_types", "null_fields", + ]) + if kind == "valid_int": + return str(rng.randint(0, 10_000)) + if kind == "valid_array": + return json.dumps([rng.choice(["a", 1, None, True]) for _ in range(rng.randint(0, 4))]) + if kind == "valid_string": + return json.dumps("".join(rng.choice("abc {}[]\"") for _ in range(rng.randint(0, 20)))) + if kind == "unbalanced": + return '{"name": "x", "opts": {' + '"k": 1' * rng.randint(0, 2) + if kind == "trailing_comma": + return '{"name": "nmap", "targets": ["x"],}' + if kind == "empty": + return rng.choice(["", " ", "{}"]) + if kind == "not_json": + return rng.choice(["not json", "", "```json\n{}\n```", "\x00\x01"]) + if kind == "stringified_nested": + return json.dumps({"name": "nmap", "targets": '["a","b"]', "opts": '{"x":1}'}) + if kind == "wrong_types": + return json.dumps({"name": rng.choice([1, None, [], {}]), "targets": rng.choice(["s", 5, {}]), + "opts": rng.choice([5, "str", []]), "query": rng.choice([[], "s", 9])}) + # null_fields + return json.dumps({"name": None, "targets": None, "opts": None, "query": None, "command": None}) + + +@unittest.skipUnless(HAS_AI, 'ai addon required') +class TestMalformedArgFuzzer(unittest.TestCase): + """Random malformed arguments across all tools must never abort the loop.""" + + def test_fuzz_arguments_never_abort_loop(self): + rng = random.Random(1337) # deterministic: any failure reproduces + failures = [] + for i in range(200): + name = rng.choice(_TOOL_NAMES) + raw = _random_weird_arguments(rng) + task = _make_loop_task() + tool_call = _tc(name, raw, call_id=f"f{i}") + try: + _items, _n, aborted = _run(task, [_resp(tool_calls=[tool_call])]) + except Exception as e: # noqa: BLE001 + failures.append(f"#{i} {name} args={raw!r} -> RAISED {type(e).__name__}: {e}") + continue + if aborted is not None: + failures.append(f"#{i} {name} args={raw!r} -> ABORTED ({type(aborted).__name__}: {aborted})") + self.assertEqual(failures, [], f"{len(failures)} resilience failures:\n" + "\n".join(failures[:20])) diff --git a/tests/unit/test_ai_session.py b/tests/unit/test_ai_session.py new file mode 100644 index 000000000..a1ff61ab0 --- /dev/null +++ b/tests/unit/test_ai_session.py @@ -0,0 +1,1132 @@ +"""Tests for secator.ai.session restore_history_from_db + remote resume branch.""" +import contextlib +import tempfile +import unittest +from unittest.mock import MagicMock, patch + + +class TestRestoreHistoryFromDB(unittest.TestCase): + """Verify restore_history_from_db rebuilds an equivalent ChatHistory from Mongo docs.""" + + def _docs(self): + # Intentionally out of timestamp order to verify sorting. + return [ + {"_type": "ai", "ai_type": "response", "content": "Hi, how can I help?", "_timestamp": 2}, + {"_type": "ai", "ai_type": "prompt", "content": "Hello", "_timestamp": 1}, + {"_type": "ai", "ai_type": "shell", "content": "nmap -p- host", "_timestamp": 3}, + {"_type": "ai", "ai_type": "prompt", "content": "Scan the target", "_timestamp": 4}, + {"_type": "ai", "ai_type": "follow_up", "content": "What next?", "_timestamp": 5}, + {"_type": "ai", "ai_type": "response", "content": "Found 2 open ports.", "_timestamp": 6}, + ] + + def test_rebuilds_order_roles_and_system(self): + from secator.ai.session import restore_history_from_db + engine = MagicMock() + engine.search.return_value = self._docs() + + history = restore_history_from_db( + "session1", engine, model="gpt-4o", system_prompt="SYSTEM PROMPT") + + # Query was scoped to the session by the auto-stamped `_context.session_id` + # (the same key the remote poll + resume branch use — NOT the top-level field, + # which prompt/response docs don't carry). + engine.search.assert_called_once_with({"_type": "ai", "_context.session_id": "session1"}) + + # System prompt set, conversation turns in timestamp order, non-turn docs skipped + self.assertEqual(history.messages, [ + {"role": "system", "content": "SYSTEM PROMPT"}, + {"role": "user", "content": "Hello"}, + {"role": "assistant", "content": "Hi, how can I help?"}, + {"role": "user", "content": "Scan the target"}, + {"role": "assistant", "content": "Found 2 open ports."}, + ]) + self.assertEqual(history.model, "gpt-4o") + + def test_no_prior_docs_returns_system_only(self): + from secator.ai.session import restore_history_from_db + engine = MagicMock() + engine.search.return_value = [] + + history = restore_history_from_db("s2", engine, system_prompt="SYS") + self.assertEqual(history.messages, [{"role": "system", "content": "SYS"}]) + + def test_no_system_prompt_yields_empty_when_no_docs(self): + from secator.ai.session import restore_history_from_db + engine = MagicMock() + engine.search.return_value = [] + + history = restore_history_from_db("s3", engine) + self.assertEqual(history.messages, []) + + def test_empty_content_docs_skipped(self): + from secator.ai.session import restore_history_from_db + engine = MagicMock() + engine.search.return_value = [ + {"ai_type": "prompt", "content": "", "_timestamp": 1}, + {"ai_type": "response", "content": "Real answer", "_timestamp": 2}, + ] + history = restore_history_from_db("s4", engine) + self.assertEqual(history.messages, [{"role": "assistant", "content": "Real answer"}]) + + def test_search_failure_returns_system_only(self): + from secator.ai.session import restore_history_from_db + engine = MagicMock() + engine.search.side_effect = RuntimeError("backend down") + + history = restore_history_from_db("s5", engine, system_prompt="SYS") + # Failure must not crash; returns just the system prompt + self.assertEqual(history.messages, [{"role": "system", "content": "SYS"}]) + + def test_encryptor_reencrypts_restored_turns(self): + from secator.ai.session import restore_history_from_db + engine = MagicMock() + engine.search.return_value = [ + {"ai_type": "prompt", "content": "scan 10.0.0.1", "_timestamp": 1}, + ] + encryptor = MagicMock() + encryptor.encrypt.side_effect = lambda t: f"ENC({t})" + + history = restore_history_from_db("s6", engine, encryptor=encryptor) + self.assertEqual(history.messages, [{"role": "user", "content": "ENC(scan 10.0.0.1)"}]) + + def test_restore_rebuilds_full_transcript_from_message(self): + """A persisted user->assistant(tool_calls)->tool->assistant round-trip + restores byte-identically, with tool_call_id pairing intact.""" + from secator.ai.session import restore_history_from_db + docs = [ + {"_type": "ai", "ai_type": "prompt", "_timestamp": 1, + "message": {"role": "user", "content": "scan example.com"}, + "_context": {"session_id": "S"}}, + {"_type": "ai", "ai_type": "response", "_timestamp": 2, + "message": {"role": "assistant", "tool_calls": [ + {"id": "c1", "type": "function", "function": {"name": "run_task", "arguments": "{}"}}]}, + "_context": {"session_id": "S"}}, + {"_type": "ai", "ai_type": "tool_result", "_timestamp": 3, + "message": {"role": "tool", "tool_call_id": "c1", "name": "run_task", "content": "80/open"}, + "_context": {"session_id": "S"}}, + {"_type": "ai", "ai_type": "response", "_timestamp": 4, + "message": {"role": "assistant", "content": "port 80 is open"}, + "_context": {"session_id": "S"}}, + ] + + class FakeEngine: + def search(self, q, **k): + return docs + + h = restore_history_from_db("S", FakeEngine(), model="gpt-4o") + roles = [m["role"] for m in h.messages if m["role"] != "system"] + self.assertEqual(roles, ["user", "assistant", "tool", "assistant"]) + self.assertEqual(h.messages[-3]["tool_calls"][0]["id"], "c1") + self.assertEqual(h.messages[-2]["tool_call_id"], "c1") # pairs correctly + # Byte-exact: every field of every persisted message survives verbatim. + self.assertEqual(h.messages, [ + {"role": "user", "content": "scan example.com"}, + {"role": "assistant", "tool_calls": [ + {"id": "c1", "type": "function", "function": {"name": "run_task", "arguments": "{}"}}]}, + {"role": "tool", "tool_call_id": "c1", "name": "run_task", "content": "80/open"}, + {"role": "assistant", "content": "port 80 is open"}, + ]) + self.assertEqual(h.model, "gpt-4o") + + def test_restore_appends_copies_not_shared_doc_references(self): + """Mutating the restored history must not mutate the source doc's message dict.""" + from secator.ai.session import restore_history_from_db + doc_message = {"role": "user", "content": "scan example.com"} + engine = MagicMock() + engine.search.return_value = [ + {"_type": "ai", "ai_type": "prompt", "_timestamp": 1, "message": doc_message}, + ] + history = restore_history_from_db("s7", engine) + history.messages[0]["content"] = "mutated" + self.assertEqual(doc_message["content"], "scan example.com") + + def test_legacy_docs_without_message_field_fall_back_to_text_only(self): + """Pre-upgrade docs (no `message` field) still restore via the old + text-only prompt/response reconstruction.""" + from secator.ai.session import restore_history_from_db + engine = MagicMock() + engine.search.return_value = [ + {"ai_type": "prompt", "content": "legacy scan request", "_timestamp": 1}, + {"ai_type": "response", "content": "legacy scan result", "_timestamp": 2}, + ] + history = restore_history_from_db("s8", engine) + self.assertEqual(history.messages, [ + {"role": "user", "content": "legacy scan request"}, + {"role": "assistant", "content": "legacy scan result"}, + ]) + + def test_message_docs_are_not_reencrypted_on_restore(self): + """Persisted `message.content` is already encrypted at persist time — restore + must append it verbatim and NOT pass it through the encryptor again, whereas + the legacy text-only fallback still encrypts (its content was never encrypted + at persist time).""" + from secator.ai.session import restore_history_from_db + from secator.ai.encryption import SensitiveDataEncryptor + + encryptor = SensitiveDataEncryptor() + plaintext = "scan admin@example.com now" + already_encrypted = encryptor.encrypt(plaintext) + self.assertNotEqual(already_encrypted, plaintext) # sanity: encryption actually changed it + + engine = MagicMock() + engine.search.return_value = [ + {"_type": "ai", "ai_type": "prompt", "_timestamp": 1, + "message": {"role": "user", "content": already_encrypted}}, + ] + history = restore_history_from_db("s9", engine, encryptor=encryptor) + + # Verbatim: restored content matches the already-encrypted string exactly + # (re-encrypting placeholders would change/garble it further). + self.assertEqual(history.messages[0]["content"], already_encrypted) + # And it decrypts back to the original plaintext downstream, proving the + # round-trip survived restore intact. + self.assertEqual(encryptor.decrypt(history.messages[0]["content"]), plaintext) + + +class TestRestoreScopingIsolation(unittest.TestCase): + """Task 6 Step 1: restore_history_from_db must be scoped to a single + conversation even when the underlying store holds docs from several. The + scoping key is `_context.session_id` (see restore_history_from_db's query), + so this uses a fake engine that ACTUALLY applies that filter to a MIXED doc + list of two conversations -- not a MagicMock stub that always returns a + fixed set regardless of the query -- to prove zero cross-bleed end to end.""" + + class _FilteringEngine: + """Fake engine standing in for a real backend: applies the caller's + `_context.session_id` filter (and `_type` filter) to `docs`, exactly the + way JsonBackend/MongoDBBackend would.""" + + def __init__(self, docs): + self.docs = docs + self.calls = [] + + def search(self, query, **kwargs): + self.calls.append(query) + wanted_session = query.get('_context.session_id') + wanted_type = query.get('_type') + return [ + d for d in self.docs + if (wanted_type is None or d.get('_type') == wanted_type) + and (wanted_session is None or (d.get('_context') or {}).get('session_id') == wanted_session) + ] + + def _mixed_docs(self): + """Two interleaved conversations, A and B, in one combined doc list.""" + return [ + {"_type": "ai", "ai_type": "prompt", "_timestamp": 1, + "message": {"role": "user", "content": "A: scan host1"}, + "_context": {"session_id": "A"}}, + {"_type": "ai", "ai_type": "prompt", "_timestamp": 1, + "message": {"role": "user", "content": "B: scan host2"}, + "_context": {"session_id": "B"}}, + {"_type": "ai", "ai_type": "response", "_timestamp": 2, + "message": {"role": "assistant", "content": "A: host1 has port 80 open"}, + "_context": {"session_id": "A"}}, + {"_type": "ai", "ai_type": "response", "_timestamp": 2, + "message": {"role": "assistant", "content": "B: host2 has port 443 open"}, + "_context": {"session_id": "B"}}, + {"_type": "ai", "ai_type": "prompt", "_timestamp": 3, + "message": {"role": "user", "content": "A: anything else?"}, + "_context": {"session_id": "A"}}, + ] + + def test_restore_isolates_session_a_from_mixed_docs(self): + from secator.ai.session import restore_history_from_db + engine = self._FilteringEngine(self._mixed_docs()) + + history = restore_history_from_db("A", engine, model="gpt-4o") + + self.assertEqual(history.messages, [ + {"role": "user", "content": "A: scan host1"}, + {"role": "assistant", "content": "A: host1 has port 80 open"}, + {"role": "user", "content": "A: anything else?"}, + ]) + # No content from B leaked into A's transcript. + for msg in history.messages: + self.assertNotIn("B:", msg["content"]) + # The engine was actually called with the session-scoping filter. + self.assertEqual(engine.calls, [{"_type": "ai", "_context.session_id": "A"}]) + + def test_restore_isolates_session_b_from_mixed_docs(self): + from secator.ai.session import restore_history_from_db + engine = self._FilteringEngine(self._mixed_docs()) + + history = restore_history_from_db("B", engine, model="gpt-4o") + + self.assertEqual(history.messages, [ + {"role": "user", "content": "B: scan host2"}, + {"role": "assistant", "content": "B: host2 has port 443 open"}, + ]) + for msg in history.messages: + self.assertNotIn("A:", msg["content"]) + self.assertEqual(engine.calls, [{"_type": "ai", "_context.session_id": "B"}]) + + +class TestRestoreCrossBackend(unittest.TestCase): + """Task 6 Step 2: restore_history_from_db must round-trip identically + whether the docs come from a mocked engine or a REAL local backend + (JsonBackend via QueryEngine reading an on-disk report.json). Proves the + restore logic itself -- not just the mock plumbing -- is backend-agnostic.""" + + def test_restore_from_real_json_backend_matches_fake_engine_path(self): + import json as _json + import tempfile as _tempfile + from pathlib import Path + from secator.query import QueryEngine + from secator.query.json import JsonBackend + from secator.ai.session import restore_history_from_db + + tmp = _tempfile.mkdtemp(prefix="secator-test-reports-") + # No hyphens/spaces: sanitize_folder_name would otherwise rewrite the + # workspace directory name and the test would look in the wrong place. + workspace_name = "wscrossbackend" + task_dir = Path(tmp) / workspace_name / "tasks" / "task1" + task_dir.mkdir(parents=True) + + session_id = "CROSS-SESSION-1" + ai_items = [ + {"_type": "ai", "ai_type": "prompt", "_timestamp": 1, + "message": {"role": "user", "content": "scan example.com"}, + "_context": {"session_id": session_id}}, + {"_type": "ai", "ai_type": "response", "_timestamp": 2, + "message": {"role": "assistant", "content": "port 80 is open"}, + "_context": {"session_id": session_id}}, + ] + report = {"results": {"ai": ai_items}} + (task_dir / "report.json").write_text(_json.dumps(report)) + + # Resolve a real QueryEngine down to the real JsonBackend, then apply the + # backend's documented reports_dir override (JsonBackend.__init__ accepts + # config={'reports_dir': ...}; QueryEngine doesn't forward a `config` kwarg + # of its own, so the override is applied directly on the resolved backend + # instance -- same effect, same attribute the constructor would have set). + engine = QueryEngine( + workspace_id=workspace_name, + context={"drivers": ["local"], "workspace_name": workspace_name}, + ) + self.assertIsInstance(engine.backend, JsonBackend) + engine.backend.reports_dir = Path(tmp) + + history = restore_history_from_db(session_id, engine, model="gpt-4o") + + expected = [ + {"role": "user", "content": "scan example.com"}, + {"role": "assistant", "content": "port 80 is open"}, + ] + self.assertEqual(history.messages, expected) + + # Same round-trip via a fake engine over the identical doc list -> proves + # parity between the real local backend and the mock/fake path. + class FakeEngine: + def search(self, query, **kwargs): + return ai_items + + fake_history = restore_history_from_db(session_id, FakeEngine(), model="gpt-4o") + self.assertEqual(history.messages, fake_history.messages) + + +class TestSizeBackstopEndToEnd(unittest.TestCase): + """Task 6 Step 3: an oversized assistant `content` is capped by cap_message + before persist (mirrors ai.py:526's `message=cap_message(assistant_msg)`), + and the capped message still restores into a valid transcript -- proving + the persist-time backstop and the restore path compose correctly.""" + + def test_oversized_assistant_content_capped_and_restores_cleanly(self): + from secator.ai.history import cap_message, MAX_PERSISTED_MESSAGE_CHARS + from secator.ai.session import restore_history_from_db + + oversized = "y" * (MAX_PERSISTED_MESSAGE_CHARS * 3) + assistant_msg = {"role": "assistant", "content": oversized} + capped = cap_message(assistant_msg) + + # Backstop applied at persist time: capped, marker present, well under budget. + slack = 20 # room for the '…[capped]' marker itself + self.assertLessEqual(len(capped["content"]), MAX_PERSISTED_MESSAGE_CHARS + slack) + self.assertIn("[capped]", capped["content"]) + self.assertLess(len(capped["content"]), len(oversized)) + + docs = [ + {"_type": "ai", "ai_type": "prompt", "_timestamp": 1, + "message": {"role": "user", "content": "scan and summarize everything"}, + "_context": {"session_id": "SZ"}}, + {"_type": "ai", "ai_type": "response", "_timestamp": 2, + "message": capped, + "_context": {"session_id": "SZ"}}, + ] + + class FakeEngine: + def search(self, query, **kwargs): + return docs + + history = restore_history_from_db("SZ", FakeEngine(), model="gpt-4o") + + # Valid transcript: correct roles/order, restored verbatim (including cap). + self.assertEqual([m["role"] for m in history.messages], ["user", "assistant"]) + self.assertEqual(history.messages[-1]["content"], capped["content"]) + self.assertLessEqual(len(history.messages[-1]["content"]), MAX_PERSISTED_MESSAGE_CHARS + slack) + self.assertIn("[capped]", history.messages[-1]["content"]) + + +class TestRemoteResumeBranch(unittest.TestCase): + """Verify the yielder remote-resume branch picks Mongo restore vs fresh start.""" + + def _make_task(self, prior_docs, backend_name="mongodb"): + from secator.tasks.ai import ai + + task = ai.__new__(ai) + # Minimal attributes the branch touches + task.interactive = "remote" + task.session_id = "sess-123" + task.session_name = "" + task.mode = "chat" + task.model = "gpt-4o" + task.encryptor = None + task.context = {"workspace_id": "ws1", "drivers": ["mongodb"]} + task.run_opts = {"prompt": "Tell me about this workspace"} + # An existing dir short-circuits the reports_folder property (no dir creation) + task._reports_folder = tempfile.mkdtemp(prefix="secator-test-") + task.backend = MagicMock() + task.debug = MagicMock() + task.history = MagicMock() + + # Stub query engine + engine = MagicMock() + engine.backend = MagicMock() + engine.backend.name = backend_name + + def _search(query, limit=0): + # The resume branch scopes by `_context.session_id` (not the top-level field). + if query.get("_type") == "ai" and any("session_id" in k for k in query): + return prior_docs + return [] + engine.search.side_effect = _search + task._get_query_engine = MagicMock(return_value=engine) + return task, engine + + def test_fresh_when_no_prior_docs(self): + task, engine = self._make_task(prior_docs=[]) + # Generator return value is the StopIteration value. + gen = task._maybe_resume_remote() + restored = None + try: + while True: + next(gen) + except StopIteration as e: + restored = e.value + self.assertFalse(restored) + + @patch("secator.tasks.ai.restore_history_from_db") + @patch("secator.tasks.ai.get_system_prompt", return_value="SYS") + def test_restores_when_prior_docs(self, mock_sys, mock_restore): + mock_history = MagicMock() + mock_history.messages = [{"role": "system", "content": "SYS"}] + mock_restore.return_value = mock_history + + task, engine = self._make_task(prior_docs=[{"ai_type": "prompt", "content": "hi"}]) + # Stub the heavy methods the branch calls + task._detect_mode = MagicMock() + task._run_loop = MagicMock(return_value=iter([])) + + gen = task._maybe_resume_remote() + restored = None + try: + while True: + next(gen) + except StopIteration as e: + restored = e.value + + self.assertTrue(restored) + mock_restore.assert_called_once() + # Restored from Mongo via the resolved query engine + _, kwargs = mock_restore.call_args + self.assertEqual(mock_restore.call_args[0][0], "sess-123") + task._run_loop.assert_called_once() + + @patch("secator.tasks.ai.restore_history_from_db") + @patch("secator.tasks.ai.get_system_prompt", return_value="SYS") + def test_warns_on_non_mongo_backend(self, mock_sys, mock_restore): + from secator.output_types import Warning as WarningType + mock_restore.return_value = MagicMock(messages=[]) + + task, engine = self._make_task( + prior_docs=[{"ai_type": "prompt", "content": "hi"}], backend_name="local") + task._detect_mode = MagicMock() + task._run_loop = MagicMock(return_value=iter([])) + + items = [] + gen = task._maybe_resume_remote() + try: + while True: + items.append(next(gen)) + except StopIteration: + pass + + warnings = [i for i in items if isinstance(i, WarningType)] + self.assertTrue(any("remote" in w.message for w in warnings)) + + +class TestTurnIdempotency(unittest.TestCase): + """C3: an acks_late redelivery of an already-completed turn must NOT replay + _run_loop (no re-run tool actions, no re-billed tokens); a genuinely + incomplete turn must still resume and run.""" + + def _make_task(self, marker_docs, prior_docs, celery_id="turn-abc"): + from secator.tasks.ai import ai + + task = ai.__new__(ai) + task.interactive = "remote" + task.session_id = "sess-123" + task.session_name = "" + task.mode = "chat" + task.model = "gpt-4o" + task.encryptor = None + task.context = {"workspace_id": "ws1", "drivers": ["mongodb"], "celery_id": celery_id} + task.context["ai_tokens"] = 0 + task.run_opts = {"prompt": "continue"} + task._reports_folder = tempfile.mkdtemp(prefix="secator-test-") + task.backend = MagicMock() + task.debug = MagicMock() + task.history = MagicMock() + + engine = MagicMock() + engine.backend = MagicMock() + engine.backend.name = "mongodb" + + def _search(query, limit=0): + # The idempotency marker query is the only one keyed by turn_completed. + if query.get("ai_type") == "turn_completed": + return marker_docs + if query.get("_type") == "ai": + return prior_docs + return [] + engine.search.side_effect = _search + task._get_query_engine = MagicMock(return_value=engine) + return task, engine + + def _drive(self, task): + gen = task._maybe_resume_remote() + restored = None + try: + while True: + next(gen) + except StopIteration as e: + restored = e.value + return restored + + def test_completed_turn_short_circuits_without_replay(self): + """A redelivery whose turn already has a completion marker short-circuits: + _run_loop is not called and no tokens are billed.""" + task, engine = self._make_task( + marker_docs=[{"ai_type": "turn_completed", "extra_data": {"turn_uuid": "turn-abc"}}], + prior_docs=[{"ai_type": "prompt", "content": "hi"}], + ) + task._run_loop = MagicMock(return_value=iter([])) + + with patch("secator.tasks.ai.restore_history_from_db") as mock_restore: + restored = self._drive(task) + + self.assertTrue(restored) # turn handled (as a no-op) + task._run_loop.assert_not_called() # no tool actions replayed + mock_restore.assert_not_called() # didn't even rebuild/append + self.assertEqual(task.context["ai_tokens"], 0) # nothing re-billed + + @patch("secator.tasks.ai.restore_history_from_db") + @patch("secator.tasks.ai.get_system_prompt", return_value="SYS") + def test_incomplete_turn_still_resumes(self, mock_sys, mock_restore): + """No marker (a real mid-turn crash) → the turn resumes and runs _run_loop.""" + mock_restore.return_value = MagicMock(messages=[{"role": "system", "content": "SYS"}]) + task, engine = self._make_task( + marker_docs=[], + prior_docs=[{"ai_type": "prompt", "content": "hi"}], + ) + task._detect_mode = MagicMock() + task._run_loop = MagicMock(return_value=iter([])) + task._mark_turn_completed = MagicMock() + + restored = self._drive(task) + + self.assertTrue(restored) + mock_restore.assert_called_once() + task._run_loop.assert_called_once() + + def test_mark_turn_completed_persists_marker(self): + """_mark_turn_completed persists exactly one turn_completed Ai stamped with + the celery_id turn_uuid; it is a no-op off the remote channel.""" + from secator.output_types import Ai + + task, engine = self._make_task(marker_docs=[], prior_docs=[]) + persisted = [] + task.add_result = lambda item, **kw: persisted.append(item) + + task._mark_turn_completed() + self.assertEqual(len(persisted), 1) + marker = persisted[0] + self.assertIsInstance(marker, Ai) + self.assertEqual(marker.ai_type, "turn_completed") + self.assertEqual(marker.extra_data.get("turn_uuid"), "turn-abc") + # The marker carries no top-level session_id; its conversation id is stamped + # onto `_context.session_id` by the runner persist pipeline (from self.context), + # which _turn_completed_marker queries by. That stamping is out of scope here. + + # Local channel: no marker persisted (idempotency is a remote concern). + persisted.clear() + task.interactive = "local" + task._mark_turn_completed() + self.assertEqual(persisted, []) + + +class TestFastDetectMode(unittest.TestCase): + """D4: the deterministic mode fast-path skips the intent LLM round-trip for + unambiguous prompts, while ambiguous ones still fall back to the LLM.""" + + def test_fast_detect_mode_pure(self): + from secator.tasks.ai import fast_detect_mode + self.assertEqual(fast_detect_mode("scan the target"), "attack") + self.assertEqual(fast_detect_mode("summarize the findings"), "chat") + self.assertEqual(fast_detect_mode(""), "chat") + # exploit-ish → defer to LLM (no behavior change for those) + self.assertIsNone(fast_detect_mode("write an exploit for this CVE-2024-1234")) + # conflicting cues → ambiguous → defer to LLM + self.assertIsNone(fast_detect_mode("scan and explain the results")) + # no cues → ambiguous → defer to LLM + self.assertIsNone(fast_detect_mode("please handle the situation")) + + def _make_task(self, prompt, mode=""): + from secator.tasks.ai import ai + task = ai.__new__(ai) + task.mode = mode + task.prompt = prompt + task.intent_model = "intent-model" + task.model = "main-model" + task.api_base = None + task.api_key = None + task.max_iterations = 10 + task.is_subagent = False + task.backend = MagicMock() + task._reports_folder = tempfile.mkdtemp(prefix="secator-test-") + task._account_usage = MagicMock() + return task + + def _patches(self): + return ( + patch("secator.tasks.ai.get_system_prompt", return_value="SYS"), + patch("secator.tasks.ai.build_tool_schemas", return_value=[]), + patch("secator.tasks.ai.get_mode_config", return_value={"max_iterations": 5}), + ) + + def test_fast_path_resolves_without_llm(self): + """Unambiguous prompt → mode set deterministically, call_llm untouched.""" + task = self._make_task("scan the target") + p_sys, p_tools, p_cfg = self._patches() + with p_sys, p_tools, p_cfg, patch("secator.tasks.ai.call_llm") as mock_llm: + task._detect_mode() + mock_llm.assert_not_called() + self.assertEqual(task.mode, "attack") + + def test_ambiguous_falls_back_to_llm(self): + """Conflicting cues → the LLM classifier still runs and decides.""" + task = self._make_task("scan and explain the results") + p_sys, p_tools, p_cfg = self._patches() + with p_sys, p_tools, p_cfg, \ + patch("secator.tasks.ai.load_prompt", return_value="SELECT"), \ + patch("secator.tasks.ai.call_llm", return_value={"content": "chat", "usage": {}}) as mock_llm: + task._detect_mode() + mock_llm.assert_called_once() + self.assertEqual(mock_llm.call_args[0][1], "intent-model") # uses intent_model + self.assertEqual(task.mode, "chat") + + def test_force_redetects_over_explicit_mode(self): + """force=True re-detects even when mode was explicitly set (fast-path applies).""" + task = self._make_task("scan the target", mode="chat") + p_sys, p_tools, p_cfg = self._patches() + # Without force, explicit mode short-circuits (no detection, no LLM). + with p_sys, p_tools, p_cfg, patch("secator.tasks.ai.call_llm") as mock_llm: + task._detect_mode() + self.assertEqual(task.mode, "chat") + mock_llm.assert_not_called() + # With force, detection runs again → fast-path flips to attack. + p_sys, p_tools, p_cfg = self._patches() + with p_sys, p_tools, p_cfg, patch("secator.tasks.ai.call_llm") as mock_llm: + task._detect_mode(force=True) + self.assertEqual(task.mode, "attack") + mock_llm.assert_not_called() + + +class TestSessionIdStampedOnContext(unittest.TestCase): + """_init_options writes the resolved session_id back onto self.context. + + Every persisted item copies self.context into its `_context` (Runner._process_item), + so this is what makes `prompt`/`response` docs queryable by `_context.session_id` + (restore_history_from_db + the remote poll both key on it). Without the stamp a + locally-resolved session_id (str(self.id)/session_name) leaves the transcript turns + unqueryable and a remote resume restores an empty history. + """ + + def _drive_init(self, context, run_opts=None): + from secator.tasks.ai import ai + task = ai.__new__(ai) + task.context = context + task.run_opts = run_opts or {} + task.results = [] + task.inputs = [] + task._reports_folder = None + task.sync = True + opt_values = { + "resume": False, "subagent": False, "model": "m", "intent_model": "im", + "api_base": None, "api_key": "k", "sensitive": False, "mode": "chat", + "max_tokens_total": 100000, "max_workers": 1, "max_iterations": 10, + "temperature": 0.7, "context_warnings": True, "async_tasks": False, + "dangerous": False, "interactive": "remote", + } + task.get_opt_value = lambda key: opt_values.get(key) + with contextlib.ExitStack() as stack: + stack.enter_context(patch('secator.tasks.ai.PermissionEngine')) + stack.enter_context(patch('secator.tasks.ai.create_backend')) + stack.enter_context(patch('secator.tasks.ai.SensitiveDataEncryptor')) + stack.enter_context(patch.object(ai, '_auto_approve_workspace_targets')) + stack.enter_context(patch.object(type(task), 'reports_folder', property(lambda self: None))) + stack.enter_context(patch.object(type(task), 'id', 'runner-id-42', create=True)) + task._init_options() + return task + + def test_stamped_when_locally_derived(self): + """No session_id anywhere -> falls back to str(self.id) AND is written to context.""" + task = self._drive_init(context={"workspace_id": "ws1"}) + self.assertEqual(task.session_id, "runner-id-42") + self.assertEqual(task.context["session_id"], "runner-id-42") + + def test_platform_supplied_session_id_preserved(self): + """A dispatcher-supplied context session_id is kept and remains the stamped value.""" + task = self._drive_init(context={"workspace_id": "ws1", "session_id": "ui-sess-abc"}) + self.assertEqual(task.session_id, "ui-sess-abc") + self.assertEqual(task.context["session_id"], "ui-sess-abc") + + def test_stamp_matches_restore_query_key(self): + """The stamped context key is exactly what restore/poll query (`_context.session_id`).""" + task = self._drive_init(context={}) + # Simulate the generic per-item context copy (Runner._process_item does self.context.copy()). + item_context = dict(task.context) + self.assertEqual(item_context.get("session_id"), task.session_id) + + +class TestLocalResumeAdoptsSessionId(unittest.TestCase): + """Task 5: a resumed LOCAL run (yielder's `self.resume` branch) adopts the + picked session's session_id and rebuilds history via the unified + restore_history_from_db over the local query engine -- not the bespoke + replay_session, which mints a fresh str(self.id) and can't correlate a + re-resumed conversation's appended docs with the prior ones.""" + + def _make_task(self, model="gpt-4o"): + from secator.tasks.ai import ai + + task = ai.__new__(ai) + task.inputs = [] + task.results = [] + task.run_opts = {"resume": True} + task.sync = True + task._reports_folder = tempfile.mkdtemp(prefix="secator-test-") + task.context = {"workspace_id": "ws1"} + task.debug = MagicMock() + + opt_values = { + "resume": True, "subagent": False, "model": model, "intent_model": "im", + "api_base": None, "api_key": "k", "sensitive": False, "mode": "", + "max_tokens_total": 100000, "max_workers": 1, "max_iterations": 10, + "temperature": 0.7, "context_warnings": True, "async_tasks": False, + "dangerous": False, "interactive": "local", + } + task.get_opt_value = lambda key: opt_values.get(key) + + # Stub the local (JSON) query engine _get_query_engine() would build. + engine = MagicMock() + task._get_query_engine = MagicMock(return_value=engine) + + # Short-circuit right after the resume block adopts session_id + restores + # history: declining the interactive "what's next?" prompt (as if the user + # cancelled) ends the run cleanly, so the full _run_loop is out of scope. + task._prompt_and_redetect = MagicMock(return_value=None) + task._save_history = MagicMock() + + return task, engine + + def _patches(self): + from secator.tasks.ai import ai + return ( + patch('secator.tasks.ai.PermissionEngine'), + patch('secator.tasks.ai.create_backend'), + patch('secator.tasks.ai.SensitiveDataEncryptor'), + patch.object(ai, '_auto_approve_workspace_targets'), + patch('secator.tasks.ai.get_system_prompt', return_value='SYS'), + patch('secator.tasks.ai.build_tool_schemas', return_value=[]), + ) + + @patch('secator.tasks.ai.show_session_picker') + @patch('secator.tasks.ai.restore_history_from_db') + def test_resume_adopts_prior_session_id_and_uses_unified_restore(self, mock_restore, mock_picker): + prior_folder = tempfile.mkdtemp(prefix="secator-test-prior-") + mock_picker.return_value = {"name": "prior chat", "folder": prior_folder, "session_id": "PRIOR-SESSION"} + mock_history = MagicMock() + mock_restore.return_value = mock_history + + task, engine = self._make_task() + + with contextlib.ExitStack() as stack: + for p in self._patches(): + stack.enter_context(p) + list(task.yielder()) + + # Adopted the picked session's id (not a freshly-minted str(self.id)), + # and stamped it back onto the context (every persisted item copies + # self.context into its `_context`, so this is what makes appended docs + # queryable by `_context.session_id` under the SAME id going forward). + self.assertEqual(task.session_id, "PRIOR-SESSION") + self.assertEqual(task.context["session_id"], "PRIOR-SESSION") + + # Restored via the unified restore over the LOCAL query engine, keyed by + # the adopted session_id -- not replay_session's bespoke rebuild. + mock_restore.assert_called_once() + args, kwargs = mock_restore.call_args + self.assertEqual(args[0], "PRIOR-SESSION") + self.assertIs(args[1], engine) + self.assertEqual(kwargs.get("model"), "gpt-4o") + self.assertIs(task.history, mock_history) + + @patch('secator.tasks.ai.print_session_results') + @patch('secator.tasks.ai.show_session_picker') + @patch('secator.tasks.ai.restore_history_from_db') + def test_new_format_resume_prints_prior_conversation(self, mock_restore, mock_picker, mock_print): + """New-format resume rebuilds history in-memory via the unified restore, which + (unlike the legacy replay_session) does NOT print anything — so the branch must + call print_session_results(session) to keep the prior conversation visible on + the console.""" + prior_folder = tempfile.mkdtemp(prefix="secator-test-print-") + mock_picker.return_value = {"name": "prior chat", "folder": prior_folder, "session_id": "PRIOR"} + mock_restore.return_value = MagicMock() + + task, engine = self._make_task() + with contextlib.ExitStack() as stack: + for p in self._patches(): + stack.enter_context(p) + list(task.yielder()) + + mock_print.assert_called_once_with(mock_picker.return_value) + + @patch('secator.tasks.ai.show_session_picker') + @patch('secator.tasks.ai.replay_session') + @patch('secator.tasks.ai.restore_history_from_db') + def test_legacy_session_without_session_id_falls_back_to_replay(self, mock_restore, mock_replay, mock_picker): + """A picked LEGACY session (pre session_id-stamping: session_id absent/'') + must resume via replay_session (reads history.json directly), NOT the + unified restore -- whose nested `_context.session_id` filter would exclude + the legacy docs (they carry no session_id) and rebuild an EMPTY history.""" + prior_folder = tempfile.mkdtemp(prefix="secator-test-legacy-") + # session_id absent entirely (older list_sessions) — same as '' in behavior. + mock_picker.return_value = {"name": "legacy chat", "folder": prior_folder} + mock_history = MagicMock() + mock_replay.return_value = mock_history + + task, engine = self._make_task() + + with contextlib.ExitStack() as stack: + for p in self._patches(): + stack.enter_context(p) + list(task.yielder()) + + # Legacy path: replay_session used to rebuild history, unified restore NOT + # called (its session_id filter would exclude the legacy docs → empty). + mock_replay.assert_called_once_with(mock_picker.return_value) + mock_restore.assert_not_called() + self.assertIs(task.history, mock_history) + # No picked session_id, so context was NOT stamped with an adopted id — it + # keeps whatever _init_options locally resolved. + self.assertEqual(task.context.get("session_id"), task.session_id) + + @patch('secator.tasks.ai.show_session_picker') + @patch('secator.tasks.ai.replay_session') + @patch('secator.tasks.ai.restore_history_from_db') + def test_empty_string_session_id_also_falls_back_to_replay(self, mock_restore, mock_replay, mock_picker): + """An explicit empty-string session_id (what list_sessions returns for a + legacy session) is falsy too → same replay fallback.""" + prior_folder = tempfile.mkdtemp(prefix="secator-test-legacy2-") + mock_picker.return_value = {"name": "legacy chat", "folder": prior_folder, "session_id": ""} + mock_replay.return_value = MagicMock() + + task, engine = self._make_task() + with contextlib.ExitStack() as stack: + for p in self._patches(): + stack.enter_context(p) + list(task.yielder()) + + mock_replay.assert_called_once() + mock_restore.assert_not_called() + + +class TestListSessionsSurfacesSessionId(unittest.TestCase): + """list_sessions() must surface each session's session_id (read from its + report.json ai docs' `_context.session_id`, first non-empty) so the local + resume branch has something to adopt (Task 5).""" + + def _write_session(self, tmp_root, ai_items, info=None): + import json as _json + from pathlib import Path + + task_dir = Path(tmp_root) / 'ws1' / 'tasks' / 'task1' + task_dir.mkdir(parents=True) + (task_dir / 'history.json').write_text('[]') + report = {"info": info or {}, "results": {"ai": ai_items}} + (task_dir / 'report.json').write_text(_json.dumps(report)) + return str(task_dir / 'history.json') + + @patch('secator.ai.session.glob.glob') + def test_list_sessions_includes_session_id(self, mock_glob): + from secator.ai.session import list_sessions + + tmp_root = tempfile.mkdtemp(prefix="secator-test-reports-") + history_path = self._write_session(tmp_root, [ + {"ai_type": "prompt", "content": "hello", "_context": {"session_name": "hi", "session_id": "SESSION-XYZ"}}, + {"ai_type": "response", "content": "hi there", "_context": {"session_id": "SESSION-XYZ"}}, + ]) + mock_glob.return_value = [history_path] + + sessions = list_sessions() + + self.assertEqual(len(sessions), 1) + self.assertEqual(sessions[0]["session_id"], "SESSION-XYZ") + + @patch('secator.ai.session.glob.glob') + def test_list_sessions_session_id_falls_back_to_first_non_empty(self, mock_glob): + """The prompt doc itself may carry no session_id (pre-stamp docs); scan + ALL ai docs and take the first non-empty one, not just the prompt doc.""" + from secator.ai.session import list_sessions + + tmp_root = tempfile.mkdtemp(prefix="secator-test-reports-") + history_path = self._write_session(tmp_root, [ + {"ai_type": "prompt", "content": "hello", "_context": {}}, + {"ai_type": "response", "content": "hi there", "_context": {"session_id": "SESSION-ABC"}}, + ]) + mock_glob.return_value = [history_path] + + sessions = list_sessions() + + self.assertEqual(sessions[0]["session_id"], "SESSION-ABC") + + @patch('secator.ai.session.glob.glob') + def test_list_sessions_session_id_empty_when_absent(self, mock_glob): + from secator.ai.session import list_sessions + + tmp_root = tempfile.mkdtemp(prefix="secator-test-reports-") + history_path = self._write_session(tmp_root, [ + {"ai_type": "prompt", "content": "hello", "_context": {}}, + ]) + mock_glob.return_value = [history_path] + + sessions = list_sessions() + + self.assertEqual(sessions[0]["session_id"], '') + + +class TestAddAssistantToHistory(unittest.TestCase): + """_add_assistant_to_history must build + append the litellm message to chat + history AND return that exact dict, so the caller (the response emission in + _run_loop) can persist it verbatim via Ai.message — including tool-call-only + turns that carry no text content (today's `if content:` gate silently drops + those turns; Task 2 fixes that by always emitting on this returned message). + """ + + def _make_task(self, encryptor=None): + from secator.tasks.ai import ai + from secator.ai.history import ChatHistory + + task = ai.__new__(ai) + task.encryptor = encryptor + task.history = ChatHistory() + return task + + def test_add_assistant_returns_message_with_tool_calls(self): + task = self._make_task() + + class TC: + id = "call_1" + + class function: + name = "run_task" + arguments = '{"name":"nmap"}' + + msg = task._add_assistant_to_history(None, [TC]) + self.assertEqual(msg["role"], "assistant") + self.assertEqual(msg["tool_calls"][0]["id"], "call_1") + self.assertEqual(msg["tool_calls"][0]["function"]["name"], "run_task") + self.assertTrue("content" not in msg or msg["content"] is None) + # The returned dict is the exact one appended to history (same content). + self.assertEqual(task.history.messages[-1], msg) + + def test_add_assistant_returns_message_text_only(self): + task = self._make_task() + msg = task._add_assistant_to_history("hello", []) + self.assertEqual(msg, {"role": "assistant", "content": "hello"}) + self.assertEqual(task.history.messages[-1], msg) + + def test_cap_message_applies_to_returned_message(self): + from secator.ai.history import cap_message, MAX_PERSISTED_MESSAGE_CHARS + task = self._make_task() + long_content = "x" * (MAX_PERSISTED_MESSAGE_CHARS + 500) + msg = task._add_assistant_to_history(long_content, []) + capped = cap_message(msg) + self.assertLess(len(capped["content"]), len(long_content)) + self.assertTrue(capped["content"].endswith('…[capped]')) + + +class TestDispatchAndCollectPersistsToolResult(unittest.TestCase): + """Task 3: _dispatch_and_collect must emit a tool_result Ai (in addition to + appending to self.history) at every add_tool_result site, so the tool turn + round-trips through session restore. Drives the real _dispatch_and_collect + with a minimal fake self, patching the shell handler to control the single + collected result (mirrors the existing _dispatch_and_collect harness in + test_ai_loop.py's TestLoopResilientToActionErrors).""" + + def _make_fake_self(self, context=None): + from secator.ai.interactivity import CLIBackend + + class _FakeHistory: + def __init__(self): + self.tool_results = [] + + def get_action_budget(self, model): + return 10000 + + def add_tool_result(self, name, tc_id, content): + self.tool_results.append((name, tc_id, content)) + + fake_self = MagicMock() + fake_self.backend = CLIBackend() + fake_self.session_id = "sess-tool-result" + fake_self.model = "test-model" + fake_self.reports_folder = None + fake_self.encryptor = None + fake_self.context = context if context is not None else {} + fake_self.history = _FakeHistory() + persisted = [] + fake_self.add_result = lambda item, **kw: persisted.append(item) + return fake_self, persisted + + def _drive(self, fake_self, ctx, action): + from secator.tasks.ai import ai as AiTask + from secator.output_types import Ai as _Ai + + def _shell_ok(*a, **k): + # ai_type="shell_output" is required for the result to reach `collected` + # (see _dispatch_and_collect: Info/Stat/Progress/State are skipped + # entirely, and Ai results only collect for ai_type in + # ("shell_output", "response")). + yield _Ai(content="ok", ai_type="shell_output", _context={ + "tool_call_id": action["tool_call_id"], + "tool_call_name": action["tool_call_name"], + }) + + with patch("secator.ai.actions._handle_shell", _shell_ok): + gen = AiTask._dispatch_and_collect(fake_self, [action], ctx) + yielded = list(gen) + return yielded + + def test_tool_result_ai_emitted_matching_tool_call_id_and_content(self): + from secator.output_types import Ai + + ctx = MagicMock() + ctx.results = [] + action = { + "action": "shell", + "command": "echo hi", + "tool_call_id": "tc_ok", + "tool_call_name": "run_shell", + } + fake_self, _persisted = self._make_fake_self() + yielded = self._drive(fake_self, ctx, action) + + # The tool result was appended to LLM-visible history exactly once. + self.assertEqual(len(fake_self.history.tool_results), 1) + name, tc_id, tool_result_str = fake_self.history.tool_results[0] + self.assertEqual(tc_id, "tc_ok") + + # A matching tool_result Ai was yielded alongside the history append. + tool_result_ais = [r for r in yielded if isinstance(r, Ai) and r.ai_type == "tool_result"] + self.assertEqual(len(tool_result_ais), 1) + doc = tool_result_ais[0] + self.assertEqual(doc.message["role"], "tool") + self.assertEqual(doc.message["tool_call_id"], tc_id) + self.assertEqual(doc.message["name"], name) + # Byte-exact: same (truncated+encrypted) string that went to history, no re-processing. + self.assertEqual(doc.message["content"], tool_result_str) + + def test_tool_result_runner_id_from_group_result_context(self): + """runner_id is pulled from the collected result's own _context + (task_id/workflow_id/scan_id), stamped by the persistence hooks - + NOT from self.context.""" + from secator.output_types import Ai + + ctx = MagicMock() + ctx.results = [] + action = { + "action": "shell", + "command": "echo hi", + "tool_call_id": "tc_runner", + "tool_call_name": "run_shell", + } + fake_self, _persisted = self._make_fake_self() + + def _shell_with_task_id(*a, **k): + yield Ai(content="ok", ai_type="shell_output", _context={ + "tool_call_id": action["tool_call_id"], + "tool_call_name": action["tool_call_name"], + "task_id": "task-xyz", + }) + + with patch("secator.ai.actions._handle_shell", _shell_with_task_id): + from secator.tasks.ai import ai as AiTask + yielded = list(AiTask._dispatch_and_collect(fake_self, [action], ctx)) + + tool_result_ais = [r for r in yielded if isinstance(r, Ai) and r.ai_type == "tool_result"] + self.assertEqual(len(tool_result_ais), 1) + self.assertEqual(tool_result_ais[0].extra_data.get("runner_id"), "task-xyz") + + def test_tool_result_ai_emitted_on_error_paths(self): + """Malformed-JSON / unknown-tool / guardrail-denial error paths (in + _process_tool_calls) also emit a tool_result Ai, not just the success + path in _dispatch_and_collect.""" + from secator.tasks.ai import ai as AiTask + from secator.output_types import Ai + + class _FakeHistory: + def __init__(self): + self.tool_results = [] + + def add_tool_result(self, name, tc_id, content): + self.tool_results.append((name, tc_id, content)) + + fake_self = MagicMock() + fake_self.encryptor = None + fake_self.context = {} + fake_self.history = _FakeHistory() + fake_self.debug = MagicMock() + fake_self.dangerous = True # skip guardrails entirely; force the malformed-JSON path + + tc = MagicMock() + tc.id = "tc_bad_json" + tc.function.name = "run_shell" + tc.function.arguments = "{not json" + + ctx = MagicMock() + items = list(AiTask._process_tool_calls(fake_self, [tc], ctx)) + + self.assertEqual(len(fake_self.history.tool_results), 1) + _, tc_id, error_content = fake_self.history.tool_results[0] + + tool_result_ais = [i for i in items if isinstance(i, Ai) and i.ai_type == "tool_result"] + self.assertEqual(len(tool_result_ais), 1) + doc = tool_result_ais[0] + self.assertEqual(doc.message["role"], "tool") + self.assertEqual(doc.message["tool_call_id"], tc_id) + self.assertEqual(doc.message["content"], error_content) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/test_ai_task_opts.py b/tests/unit/test_ai_task_opts.py index 518dec89b..401c0a0ec 100644 --- a/tests/unit/test_ai_task_opts.py +++ b/tests/unit/test_ai_task_opts.py @@ -1,5 +1,6 @@ """Tests for AI task subagent opts.""" import unittest +from unittest.mock import MagicMock, patch from secator.definitions import ADDONS_ENABLED @@ -25,6 +26,59 @@ def test_max_workers_opt_exists(self): self.assertEqual(ai.opts["max_workers"].get("default"), 3) self.assertTrue(ai.opts["max_workers"].get("internal", False)) + def test_mode_opt_help_lists_all_modes(self): + """D2: the mode opt help documents every real mode (derived from MODES).""" + from secator.tasks.ai import ai + from secator.ai.prompts import MODES + help_text = ai.opts["mode"]["help"] + for mode in MODES: + self.assertIn(mode, help_text) + self.assertIn("exploit", help_text) # the previously-omitted one + + +@unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') +class TestDetectMode(unittest.TestCase): + """D2: _detect_mode must honor an LLM 'exploit' classification (was discarded).""" + + def _make_task(self, prompt): + """Bare ai instance with just the attributes _detect_mode reads.""" + from secator.tasks.ai import ai + t = ai.__new__(ai) + t.mode = "" # no explicit mode -> detection runs + t.prompt = prompt + t.intent_model = "test-intent-model" + t.api_base = None + t.api_key = None + t.backend = MagicMock() + t.is_subagent = False + t.max_iterations = 10 + t._account_usage = MagicMock() + return t + + def _run_detect(self, prompt, llm_word): + """Force the LLM branch (ambiguous prompt) and stub call_llm's verdict.""" + from secator.tasks.ai import ai + t = self._make_task(prompt) + with patch("secator.tasks.ai.call_llm", return_value={"content": llm_word, "usage": {}}), \ + patch("secator.tasks.ai.get_system_prompt", return_value="sys"), \ + patch("secator.tasks.ai.build_tool_schemas", return_value=[]), \ + patch.object(ai, "reports_folder", "/tmp/ws"): + t._detect_mode() + return t.mode + + def test_llm_exploit_classification_is_honored(self): + # ambiguous prompt -> defers to LLM; LLM says exploit -> mode is exploit (was 'chat') + self.assertEqual(self._run_detect("take a look at this thing", "exploit"), "exploit") + + def test_llm_attack_classification_unchanged(self): + self.assertEqual(self._run_detect("take a look at this thing", "attack"), "attack") + + def test_llm_chat_classification_unchanged(self): + self.assertEqual(self._run_detect("take a look at this thing", "chat"), "chat") + + def test_llm_unknown_classification_falls_back_to_chat(self): + self.assertEqual(self._run_detect("take a look at this thing", "banana"), "chat") + if __name__ == '__main__': unittest.main() diff --git a/tests/unit/test_ai_tokens.py b/tests/unit/test_ai_tokens.py new file mode 100644 index 000000000..72d311226 --- /dev/null +++ b/tests/unit/test_ai_tokens.py @@ -0,0 +1,491 @@ +"""Tests for per-run billed AI token accounting. + +The `ai` task accumulates billed tokens from every LLM call it makes into +`context.ai_tokens` (and cost into `context.ai_cost`). The platform billing +chore reads `context.ai_tokens` — the AI analog of `context.scan_hours`. + +These tests verify: +- N calls with known token counts sum onto `context.ai_tokens`. +- Missing/None usage counts as 0 and never crashes the run. +- History summarization usage is rolled in exactly once. +""" +import contextlib +import types +import unittest +from unittest.mock import patch + +from secator.definitions import ADDONS_ENABLED + +HAS_AI = ADDONS_ENABLED.get('ai', False) + +if HAS_AI: + from secator.tasks.ai import ai + from secator.ai.history import ChatHistory + + +def _fake_tool_call(name="noop", call_id="t1"): + """A minimal litellm-shaped tool_call object (has .id and .function.*).""" + return types.SimpleNamespace( + id=call_id, + function=types.SimpleNamespace(name=name, arguments="{}"), + ) + + +def _make_task(): + """Construct a bare `ai` task instance with a context dict, bypassing __init__. + + We avoid the full runner construction (which needs a workspace, backend, etc.) + since the accounting helpers only touch `self.context` and `self.history`. + """ + task = ai.__new__(ai) + task.context = {} + task.history = ChatHistory() + # Mirror what _init_options seeds. + task.context.setdefault("ai_tokens", 0) + task.context.setdefault("ai_prompt_tokens", 0) + task.context.setdefault("ai_completion_tokens", 0) + task.context.setdefault("ai_cost", 0.0) + return task + + +@unittest.skipUnless(HAS_AI, 'ai addon required') +class TestAiTokenAccounting(unittest.TestCase): + + def test_sum_over_n_calls(self): + """N call_llm usages sum onto context.ai_tokens (and ai_cost).""" + task = _make_task() + usages = [ + {"tokens": 100, "cost": 0.001}, + {"tokens": 250, "cost": 0.002}, + {"tokens": 50, "cost": 0.0005}, + ] + for u in usages: + task._account_usage(u) + self.assertEqual(task.context["ai_tokens"], 400) + self.assertAlmostEqual(task.context["ai_cost"], 0.0035) + + def test_missing_usage_counts_as_zero(self): + """None / empty / missing-key usage never crashes and adds 0.""" + task = _make_task() + task._account_usage(None) + task._account_usage({}) + task._account_usage({"tokens": None, "cost": None}) + task._account_usage({"cost": 0.5}) # no tokens key + self.assertEqual(task.context["ai_tokens"], 0) + + def test_malformed_usage_does_not_crash(self): + """Non-numeric token/cost values are ignored, not raised.""" + task = _make_task() + task._account_usage({"tokens": "abc", "cost": "xyz"}) + task._account_usage({"tokens": 42, "cost": 0.01}) + self.assertEqual(task.context["ai_tokens"], 42) + + def test_field_persisted_on_context(self): + """The platform reads context.ai_tokens — confirm that exact key.""" + task = _make_task() + task._account_usage({"tokens": 123, "cost": 0.0}) + self.assertIn("ai_tokens", task.context) + self.assertEqual(task.context["ai_tokens"], 123) + self.assertIsInstance(task.context["ai_tokens"], int) + + def test_prompt_completion_split_accumulated(self): + """prompt_tokens/completion_tokens accumulate into their own context keys.""" + task = _make_task() + task._account_usage({"tokens": 300, "prompt_tokens": 200, "completion_tokens": 100, "cost": 0.0}) + task._account_usage({"tokens": 60, "prompt_tokens": 40, "completion_tokens": 20, "cost": 0.0}) + self.assertEqual(task.context["ai_tokens"], 360) + self.assertEqual(task.context["ai_prompt_tokens"], 240) + self.assertEqual(task.context["ai_completion_tokens"], 120) + + def test_prompt_completion_split_missing_is_zero(self): + """Usage with only total tokens leaves the split at 0 (no crash).""" + task = _make_task() + task._account_usage({"tokens": 100, "cost": 0.0}) + self.assertEqual(task.context["ai_tokens"], 100) + self.assertEqual(task.context["ai_prompt_tokens"], 0) + self.assertEqual(task.context["ai_completion_tokens"], 0) + + def test_history_summarization_usage_drained_once(self): + """Billed tokens accrued by history compaction roll in exactly once.""" + task = _make_task() + # Simulate ChatHistory.compact stashing summarization usage. + task.history.billed_tokens = 500 + task.history.billed_cost = 0.004 + task._drain_history_usage() + self.assertEqual(task.context["ai_tokens"], 500) + self.assertAlmostEqual(task.context["ai_cost"], 0.004) + # Draining again must not double-count. + task._drain_history_usage() + self.assertEqual(task.context["ai_tokens"], 500) + + def test_history_compact_records_billed_usage(self): + """ChatHistory.compact accrues the summarization call's billed tokens.""" + history = ChatHistory(model="test-model") + history.set_system("system") + history.add_user("u1") + history.add_assistant("a1") + history.add_user("u2") + history.add_assistant("a2") + history.add_user("u3") + history.add_assistant("a3") + + fake = {"content": "summary", "usage": {"tokens": 321, "cost": 0.003}} + with patch('secator.ai.utils.call_llm', return_value=fake): + with patch('secator.ai.history.get_context_window', return_value=8000): + history.compact("test-model") + + self.assertEqual(history.billed_tokens, 321) + self.assertAlmostEqual(history.billed_cost, 0.003) + + def test_history_compact_missing_usage_is_zero(self): + """compact() with no usage on the response adds 0 billed tokens.""" + history = ChatHistory(model="test-model") + history.set_system("system") + history.add_user("u1") + history.add_assistant("a1") + history.add_user("u2") + history.add_assistant("a2") + history.add_user("u3") + history.add_assistant("a3") + + fake = {"content": "summary", "usage": None} + with patch('secator.ai.utils.call_llm', return_value=fake): + with patch('secator.ai.history.get_context_window', return_value=8000): + history.compact("test-model") + + self.assertEqual(history.billed_tokens, 0) + + +@unittest.skipUnless(HAS_AI, 'ai addon required') +class TestAiModelRecording(unittest.TestCase): + """The resolved run model is recorded on context.ai_model for the metering chore.""" + + def _run_init_options(self, model): + """Drive _init_options with the heavy collaborators stubbed out. + + Only the bits _init_options touches are stubbed; we assert the + context.ai_model recording, which sits next to the ai_tokens seeding. + """ + task = ai.__new__(ai) + task.context = {} + task.run_opts = {} + task.results = [] + task.inputs = [] + task._reports_folder = None + task.sync = True + + opt_values = { + "resume": False, + "subagent": False, + "model": model, + "intent_model": "intent-model", + "api_base": None, + "api_key": "key", + "sensitive": False, + "mode": "chat", + "max_tokens_total": 100000, + "max_workers": 1, + "max_iterations": 10, + "temperature": 0.7, + "context_warnings": True, + "async_tasks": False, + "dangerous": False, + "interactive": "auto", + } + task.get_opt_value = lambda key: opt_values.get(key) + + with contextlib.ExitStack() as stack: + stack.enter_context(patch('secator.tasks.ai.PermissionEngine')) + stack.enter_context(patch('secator.tasks.ai.create_backend')) + stack.enter_context(patch('secator.tasks.ai.SensitiveDataEncryptor')) + stack.enter_context(patch.object(ai, '_auto_approve_workspace_targets')) + stack.enter_context(patch.object(type(task), 'reports_folder', property(lambda self: None))) + stack.enter_context(patch.object(type(task), 'id', 'task-id', create=True)) + task._init_options() + return task + + def test_ai_model_recorded_on_context(self): + """context.ai_model == the resolved run model (the chore prices against it).""" + task = self._run_init_options("openrouter/anthropic/claude-sonnet-4.6") + self.assertEqual(task.context["ai_model"], "openrouter/anthropic/claude-sonnet-4.6") + + def test_ai_model_recorded_alongside_token_seeds(self): + """ai_model is seeded next to the ai_tokens accounting keys.""" + task = self._run_init_options("openrouter/google/gemma-4-26b-a4b-it:free") + self.assertEqual(task.context["ai_model"], "openrouter/google/gemma-4-26b-a4b-it:free") + self.assertEqual(task.context["ai_tokens"], 0) + self.assertIn("ai_prompt_tokens", task.context) + + +@contextlib.contextmanager +def _loop_patches(task, responses): + """Patch the heavy collaborators _run_loop touches so we can drive it bare. + + Leaves call_llm token accounting intact (that is what we are testing). + """ + with contextlib.ExitStack() as stack: + stack.enter_context(patch('secator.tasks.ai.call_llm', side_effect=responses)) + stack.enter_context(patch('secator.ai.history.get_context_window', return_value=8000)) + stack.enter_context(patch('secator.tasks.ai.get_context_window', return_value=8000)) + stack.enter_context(patch('secator.tasks.ai.save_history')) + stack.enter_context(patch.object(type(task), 'reports_folder', property(lambda self: None))) + stack.enter_context(patch.object(ai, '_summarize_auto', return_value=iter(()))) + stack.enter_context(patch.object(ai, '_summarize_user', return_value=iter(()))) + yield stack + + +@unittest.skipUnless(HAS_AI, 'ai addon required') +class TestAiTokenAccountingEndToEnd(unittest.TestCase): + """Drive the real _run_loop with mocked call_llm and assert the sum lands.""" + + def _make_loop_task(self): + task = _make_task() + # Minimal state _run_loop reads. + task.inputs = [] + task.model = "test-model" + task.intent_model = "test-model" + task.temp = 0.7 + task.api_base = None + task.api_key = "key" + task.max_iterations = 3 + task.max_tokens_total = 100000 + task.max_workers = 1 + task.is_subagent = True + task.verbose = False + task.dry_run = False + task.mode = "chat" + task.scope = "workspace" + task.results = [] + task.encryptor = None + task.tool_schemas = [] + task.permission_engine = None + task.dangerous = True + task.interactive = "auto" + task._sync = True + task.session_id = "s" + task._reports_folder = None + task.debug = lambda *a, **k: None + task.add_result = lambda *a, **k: None + from secator.ai.interactivity import create_backend + task.backend = create_backend("auto") + return task + + def test_loop_sums_token_usage(self): + """Three content responses with known tokens sum onto context.ai_tokens.""" + task = self._make_loop_task() + responses = [ + {"content": "r1", "tool_calls": [], "usage": {"tokens": 100, "cost": 0.001}}, + {"content": "r2", "tool_calls": [], "usage": {"tokens": 200, "cost": 0.002}}, + {"content": "r3", "tool_calls": [], "usage": {"tokens": 300, "cost": 0.003}}, + ] + # auto backend returns None on follow-up prompt -> loop exits after first + # content-only response. Force it to keep going by mocking the prompt to + # add a user turn for the first two, then exit. + prompt_calls = {"n": 0} + + def fake_prompt(choices, **kwargs): # loop passes prompt_uuid= on the content-only path + prompt_calls["n"] += 1 + if prompt_calls["n"] >= 3: + return None # exit + task.history.add_user("continue") + return [] + + with _loop_patches(task, responses): + with patch.object(ai, '_prompt_and_redetect', side_effect=fake_prompt): + list(task._run_loop()) + + self.assertEqual(task.context["ai_tokens"], 600) + self.assertAlmostEqual(task.context["ai_cost"], 0.006) + + def test_loop_with_no_usage_is_zero(self): + """Responses without usage leave context.ai_tokens at 0 (no crash).""" + task = self._make_loop_task() + responses = [ + {"content": "r1", "tool_calls": [], "usage": None}, + ] + with _loop_patches(task, responses): + with patch.object(ai, '_prompt_and_redetect', return_value=None): + list(task._run_loop()) + + self.assertEqual(task.context["ai_tokens"], 0) + + def test_tool_only_turn_is_counted(self): + """A main-loop turn that only calls tools (no display content) is billed. + + The old `Σ ai_type=="response"` approach missed these turns entirely + (the `response` Ai is gated on `if content:`). The accumulator must count + the call's tokens regardless of whether it produced content. + """ + task = self._make_loop_task() + # Turn 1: tool-only (no content). Turn 2: content -> exits. + responses = [ + {"content": "", "tool_calls": [_fake_tool_call()], "usage": {"tokens": 100, "cost": 0.001}}, + {"content": "done", "tool_calls": [], "usage": {"tokens": 50, "cost": 0.0005}}, + ] + with _loop_patches(task, responses): + # Tool call is consumed, returns no follow-up actions -> loop continues. + with patch.object(ai, '_process_tool_calls', return_value=iter(())): + with patch.object(ai, '_prompt_and_redetect', return_value=None): + list(task._run_loop()) + + # Both turns counted, including the content-less tool-only turn. + self.assertEqual(task.context["ai_tokens"], 150) + self.assertAlmostEqual(task.context["ai_cost"], 0.0015) + + def test_combined_main_intent_compaction_sum(self): + """Main-loop + intent-detection + compaction usages all sum onto context. + + Proves the three distinct billed call sites the audit flagged + (tool-only main turn, _detect_mode intent call, history compaction) + are aggregated into a single context.ai_tokens total. + """ + task = self._make_loop_task() + # (b) intent-detection call (as _detect_mode does it). + task._account_usage({"tokens": 30, "cost": 0.0003}) + # (c) compaction call (as ChatHistory.compact stashes, then drained). + task.history.billed_tokens = 70 + task.history.billed_cost = 0.0007 + task._drain_history_usage() + # (a) tool-only main-loop turn driven through the real loop. + responses = [ + {"content": "", "tool_calls": [_fake_tool_call()], "usage": {"tokens": 100, "cost": 0.001}}, + {"content": "done", "tool_calls": [], "usage": {"tokens": 50, "cost": 0.0005}}, + ] + with _loop_patches(task, responses): + with patch.object(ai, '_process_tool_calls', return_value=iter(())): + with patch.object(ai, '_prompt_and_redetect', return_value=None): + list(task._run_loop()) + + # 30 (intent) + 70 (compaction) + 100 (tool-only) + 50 (content) = 250 + self.assertEqual(task.context["ai_tokens"], 250) + self.assertAlmostEqual(task.context["ai_cost"], 0.0025) + + +@unittest.skipUnless(HAS_AI, 'ai addon required') +class TestAiRateLimitTermination(unittest.TestCase): + """A persistent 429 must terminate the loop after a bounded number of failures (H1).""" + + def _make_loop_task(self, max_iterations): + task = _make_task() + task.inputs = [] + task.model = "test-model" + task.intent_model = "test-model" + task.temp = 0.7 + task.api_base = None + task.api_key = "key" + task.max_iterations = max_iterations + task.max_tokens_total = 100000 + task.max_workers = 1 + task.is_subagent = True + task.verbose = False + task.dry_run = False + task.mode = "chat" + task.scope = "workspace" + task.results = [] + task.encryptor = None + task.tool_schemas = [] + task.permission_engine = None + task.dangerous = True + task.interactive = "auto" + task._sync = True + task.session_id = "s" + task._reports_folder = None + task.debug = lambda *a, **k: None + task.add_result = lambda *a, **k: None + from secator.ai.interactivity import create_backend + task.backend = create_backend("auto") + return task + + def test_persistent_rate_limit_aborts_bounded(self): + """A 429 on every call_llm aborts after 4 attempts, regardless of max_iterations.""" + import litellm + from secator.output_types import Error + + task = self._make_loop_task(max_iterations=50) + calls = {"n": 0} + + def always_rate_limited(*args, **kwargs): + calls["n"] += 1 + raise litellm.RateLimitError("rate limited", "openai", "test-model") + + with _loop_patches(task, always_rate_limited): + results = list(task._run_loop()) + + # bounded by the 4-consecutive-429 cap, not the 50-iteration budget + self.assertEqual(calls["n"], 4) + errors = [r for r in results if isinstance(r, Error)] + self.assertTrue(errors, "expected an Error to be yielded on abort") + self.assertIn("Rate limit", errors[-1].message) + + +@unittest.skipUnless(HAS_AI, 'ai addon required') +class TestAiToolPairTrim(unittest.TestCase): + """Trim/compaction must not leave a leading orphan tool_result (H2). + + litellm trim_messages and the blind keep_last tail cut drop the OLDEST + messages with no tool-pairing awareness, so the kept window can START with a + tool_result whose assistant(tool_calls) parent was dropped — which + Anthropic/OpenAI reject. The fix strips those leading orphans. + """ + + def test_strip_leading_orphan_tools_keeps_system(self): + from secator.ai.utils import _strip_leading_orphan_tools + msgs = [ + {"role": "system", "content": "s"}, + {"role": "tool", "tool_call_id": "t1", "content": "{}"}, + {"role": "tool", "tool_call_id": "t2", "content": "{}"}, + {"role": "user", "content": "u"}, + ] + removed = _strip_leading_orphan_tools(msgs) + self.assertEqual(removed, 2) + self.assertEqual([m["role"] for m in msgs], ["system", "user"]) + + def test_repair_handles_leading_orphan_tool(self): + from secator.ai.utils import _repair_orphan_tool_uses + msgs = [ + {"role": "tool", "tool_call_id": "t1", "content": "{}"}, + {"role": "user", "content": "u"}, + ] + n = _repair_orphan_tool_uses(msgs) + self.assertEqual(n, 1) + self.assertEqual(msgs[0]["role"], "user") + + def test_trim_strips_leading_orphan_tool(self): + """After litellm drops the assistant parent, trim() removes the orphan tool.""" + history = ChatHistory(model="test-model") + history.set_system("sys") + history.add_assistant_with_tool_calls(None, [{"id": "t1", "function": {"name": "noop", "arguments": "{}"}}]) + history.add_tool_result("noop", "t1", "{}") + history.add_user("u1") + history.add_assistant("a1") + # Simulate litellm dropping the oldest (assistant parent) but keeping its tool_result. + simulated = [history.messages[0], history.messages[2], history.messages[3], history.messages[4]] + with patch('litellm.utils.trim_messages', return_value=simulated): + out = history.trim(100) + nonsys = [m for m in out if m["role"] != "system"] + self.assertEqual(nonsys[0]["role"], "user") + self.assertFalse(any(m["role"] == "tool" for m in out)) + + def test_compact_strips_leading_orphan_tool_in_kept_tail(self): + """keep_last (fixed at 4) tail cut that starts on a tool_result is repaired.""" + history = ChatHistory(model="test-model") + history.set_system("sys") + history.add_user("u1") + history.add_assistant_with_tool_calls(None, [{"id": "t1", "function": {"name": "noop", "arguments": "{}"}}]) + history.add_tool_result("noop", "t1", "{}") + history.add_user("u_filler") + history.add_assistant("a2") + history.add_user("u2") + + fake = {"content": "summary", "usage": None} + with patch('secator.ai.utils.call_llm', return_value=fake): + with patch('secator.ai.history.get_context_window', return_value=8000): + history.compact("test-model") + + nonsys = [m for m in history.messages if m["role"] != "system"] + self.assertIn(nonsys[0]["role"], ("user", "assistant")) + self.assertFalse(any(m["role"] == "tool" for m in history.messages)) + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/unit/test_ai_tools.py b/tests/unit/test_ai_tools.py index 8a491fe3e..efcb19285 100644 --- a/tests/unit/test_ai_tools.py +++ b/tests/unit/test_ai_tools.py @@ -198,6 +198,51 @@ def test_unknown_tool_returns_none(self): result = tool_call_to_action("nonexistent_tool", {"foo": "bar"}) self.assertIsNone(result) + def test_non_dict_arguments_rejected(self): + """Non-object arguments (a bare JSON int/array/string) reject cleanly to None + instead of raising AttributeError on .items() and aborting the loop.""" + from secator.ai.tools import tool_call_to_action + for bad in (12345, ["nmap", "10.0.0.1"], "just a string"): + self.assertIsNone(tool_call_to_action("run_task", bad)) + + +@unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') +class TestCoerceStringifiedArgs(unittest.TestCase): + """Models sometimes serialize object/array params as JSON strings even though + the schema says object/array — coerce them back at the tool-call boundary.""" + + def test_stringified_opts_and_targets_coerced(self): + from secator.ai.tools import coerce_stringified_args + args = coerce_stringified_args("run_task", { + "name": "nmap", + "targets": '["10.0.0.1", "10.0.0.2"]', # array sent as string + "opts": '{"session_name": "scan-x", "top_ports": 100}', # object sent as string + }) + self.assertEqual(args["targets"], ["10.0.0.1", "10.0.0.2"]) + self.assertEqual(args["opts"], {"session_name": "scan-x", "top_ports": 100}) + + def test_stringified_query_coerced(self): + from secator.ai.tools import coerce_stringified_args + args = coerce_stringified_args("query_workspace", {"query": '{"_type": "url"}'}) + self.assertEqual(args["query"], {"_type": "url"}) + + def test_already_typed_args_untouched(self): + from secator.ai.tools import coerce_stringified_args + args = coerce_stringified_args("run_task", {"name": "nmap", "targets": ["a"], "opts": {"x": 1}}) + self.assertEqual(args["targets"], ["a"]) + self.assertEqual(args["opts"], {"x": 1}) + + def test_malformed_json_left_as_is(self): + from secator.ai.tools import coerce_stringified_args + args = coerce_stringified_args("run_task", {"name": "nmap", "opts": "not json"}) + self.assertEqual(args["opts"], "not json") # left for the handler to reject cleanly + + def test_scalar_string_params_not_coerced(self): + """A string-typed param (e.g. run_shell.command) must stay a string.""" + from secator.ai.tools import coerce_stringified_args + args = coerce_stringified_args("run_shell", {"command": '{"looks": "like json"}'}) + self.assertEqual(args["command"], '{"looks": "like json"}') + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/test_ai_utils.py b/tests/unit/test_ai_utils.py index 4268b0b44..15c59688e 100644 --- a/tests/unit/test_ai_utils.py +++ b/tests/unit/test_ai_utils.py @@ -82,15 +82,18 @@ def test_call_llm_basic(self, mock_cost, mock_completion): self.assertEqual(result["tool_calls"], []) mock_completion.assert_called_once() + @patch('litellm.token_counter') @patch('litellm.completion') - def test_call_llm_no_usage(self, mock_completion): - """Response without usage data.""" + def test_call_llm_no_usage_estimates_tokens(self, mock_completion, mock_token_counter): + """M5: response without usage still yields a non-zero estimated token count.""" mock_response = MagicMock() mock_response.choices = [MagicMock()] mock_response.choices[0].message.content = "Response" mock_response.choices[0].message.tool_calls = None mock_response.usage = None mock_completion.return_value = mock_response + # prompt (messages=...) then completion (text=...) + mock_token_counter.side_effect = [42, 8] from secator.ai.utils import call_llm result = call_llm( @@ -99,7 +102,11 @@ def test_call_llm_no_usage(self, mock_completion): ) self.assertEqual(result["content"], "Response") - self.assertIsNone(result["usage"]) + self.assertIsNotNone(result["usage"]) + self.assertEqual(result["usage"]["tokens"], 50) + self.assertEqual(result["usage"]["prompt_tokens"], 42) + self.assertEqual(result["usage"]["completion_tokens"], 8) + self.assertIsNone(result["usage"]["cost"]) self.assertEqual(result["tool_calls"], []) @patch('litellm.completion') @@ -234,6 +241,126 @@ def test_call_llm_tool_call_with_malformed_json(self, mock_completion): self.assertEqual(tc.function.name, "broken_tool") self.assertEqual(tc.function.arguments, "{not valid json") + @patch('time.sleep') + @patch('litellm.completion') + def test_call_llm_non_orphan_400_fails_fast(self, mock_completion, mock_sleep): + """M4: a plain (non-orphan) 400 is raised immediately, NOT retried 3x.""" + import litellm + from secator.ai.utils import call_llm + + err = litellm.BadRequestError( + message="litellm.BadRequestError: context_length_exceeded", + model="test-model", llm_provider="anthropic", + ) + mock_completion.side_effect = err + + with self.assertRaises(litellm.BadRequestError): + call_llm([{"role": "user", "content": "hi"}], "test-model") + + self.assertEqual(mock_completion.call_count, 1) # no 3x spin + mock_sleep.assert_not_called() + + @patch('time.sleep') + @patch('litellm.completion') + def test_call_llm_orphan_400_repairs_and_retries(self, mock_completion, mock_sleep): + """M4: the orphan tool_use 400 still triggers repair-and-retry (no fail-fast).""" + import litellm + from secator.ai.utils import call_llm + + ok_response = MagicMock() + ok_response.choices = [MagicMock(message=MagicMock(content="ok", tool_calls=None))] + ok_response.usage = None + + err = litellm.BadRequestError( + message="AnthropicException - tool_use ids were found without tool_result blocks", + model="claude", llm_provider="anthropic", + ) + + calls = [] + + def side_effect(**kwargs): + if not calls: # first call: inject orphan, then raise the orphan 400 + kwargs["messages"].insert(0, { + "role": "assistant", "content": None, + "tool_calls": [{"id": "toolu_late", "type": "function", + "function": {"name": "f", "arguments": "{}"}}], + }) + calls.append(1) + raise err + return ok_response + + mock_completion.side_effect = side_effect + result = call_llm([{"role": "user", "content": "hi"}], "claude") + + self.assertEqual(result["content"], "ok") + self.assertEqual(mock_completion.call_count, 2) # repaired then succeeded + mock_sleep.assert_not_called() # repair skips the backoff + + @patch('time.sleep') + @patch('litellm.completion') + def test_call_llm_duplicate_tool_result_400_repairs_and_retries(self, mock_completion, mock_sleep): + """A 'multiple tool_result blocks with id' 400 is now deduped and retried, + instead of failing fast as non-retryable.""" + import litellm + from secator.ai.utils import call_llm + + ok_response = MagicMock() + ok_response.choices = [MagicMock(message=MagicMock(content="ok", tool_calls=None))] + ok_response.usage = None + + err = litellm.BadRequestError( + message=("messages.24.content.3: each tool_use must have a single result. " + "Found multiple tool_result blocks with id: toolu_dup"), + model="claude", llm_provider="anthropic", + ) + calls = [] + + def side_effect(**kwargs): + if not calls: # first call: inject an assistant + duplicate tool_results, then raise + kwargs["messages"][:0] = [ + {"role": "assistant", "content": None, + "tool_calls": [{"id": "toolu_dup", "type": "function", + "function": {"name": "f", "arguments": "{}"}}]}, + {"role": "tool", "tool_call_id": "toolu_dup", "name": "f", "content": "r1"}, + {"role": "tool", "tool_call_id": "toolu_dup", "name": "f", "content": "r2"}, + ] + calls.append(1) + raise err + return ok_response + + mock_completion.side_effect = side_effect + result = call_llm([{"role": "user", "content": "hi"}], "claude") + + self.assertEqual(result["content"], "ok") + self.assertEqual(mock_completion.call_count, 2) # deduped then succeeded + mock_sleep.assert_not_called() + + @patch('time.sleep') + @patch('litellm.completion') + @patch('litellm.completion_cost') + def test_call_llm_transient_error_still_retries(self, mock_cost, mock_completion, mock_sleep): + """M4: genuinely-transient errors (429/500) still retry then succeed.""" + import litellm + from secator.ai.utils import call_llm + + ok_response = MagicMock() + ok_response.choices = [MagicMock()] + ok_response.choices[0].message.content = "ok" + ok_response.choices[0].message.tool_calls = None + ok_response.usage.total_tokens = 10 + mock_cost.return_value = 0.0 + + err = litellm.RateLimitError( + message="rate limited", model="test-model", llm_provider="anthropic", + ) + mock_completion.side_effect = [err, ok_response] + + result = call_llm([{"role": "user", "content": "hi"}], "test-model") + + self.assertEqual(result["content"], "ok") + self.assertEqual(mock_completion.call_count, 2) # transient retry honored + mock_sleep.assert_called_once() + @unittest.skipUnless(ADDONS_ENABLED['ai'], 'ai addon not installed') class TestPromptUserAllChoices(unittest.TestCase): @@ -251,7 +378,7 @@ def test_all_choices_not_shown_with_single_choice(self, mock_menu_class): mock_menu_class.return_value = mock_menu history = ChatHistory() - history.add_system("system") + history.set_system("system") prompt_user(history, choices=["Single choice"]) @@ -274,7 +401,7 @@ def test_all_choices_shown_with_multiple_choices(self, mock_menu_class): mock_menu_class.return_value = mock_menu history = ChatHistory() - history.add_system("system") + history.set_system("system") prompt_user(history, choices=["Choice A", "Choice B"]) @@ -296,7 +423,7 @@ def test_all_choices_position_after_llm_choices(self, mock_menu_class): mock_menu_class.return_value = mock_menu history = ChatHistory() - history.add_system("system") + history.set_system("system") prompt_user(history, choices=["Choice A", "Choice B", "Choice C"]) @@ -323,7 +450,7 @@ def test_all_choices_formats_message_correctly(self, mock_menu_class): mock_menu_class.return_value = mock_menu history = ChatHistory() - history.add_system("system") + history.set_system("system") result = prompt_user(history, choices=choices, max_iterations=10) @@ -344,7 +471,7 @@ def test_all_choices_with_extra_instructions(self, mock_menu_class): mock_menu_class.return_value = mock_menu history = ChatHistory() - history.add_system("system") + history.set_system("system") result = prompt_user(history, choices=choices) @@ -466,11 +593,69 @@ def completion_side_effect(**kwargs): with patch('litellm.completion', side_effect=completion_side_effect), \ patch('time.sleep') as mock_sleep: - result = call_llm(messages, "claude", max_retries=3) + result = call_llm(messages, "claude") self.assertEqual(result["content"], "ok") mock_sleep.assert_not_called() # repair branch should skip the backoff sleep +class TestDedupeToolResults(unittest.TestCase): + """Providers reject >1 tool_result per tool_use id ('multiple tool_result blocks + with id X') — a non-retryable 400. Duplicates arise from batch results grouped + out of order or history trim/compaction; drop the extras, keep the first.""" + + def test_drops_consecutive_duplicate_same_id(self): + from secator.ai.utils import _dedupe_tool_results + messages = [ + {"role": "assistant", "content": None, + "tool_calls": [{"id": "x", "type": "function", "function": {"name": "f", "arguments": "{}"}}]}, + {"role": "tool", "tool_call_id": "x", "name": "f", "content": "first"}, + {"role": "tool", "tool_call_id": "y", "name": "f", "content": "other"}, + {"role": "tool", "tool_call_id": "x", "name": "f", "content": "DUP"}, + {"role": "user", "content": "next"}, + ] + removed = _dedupe_tool_results(messages) + self.assertEqual(removed, 1) + tool_ids = [m["tool_call_id"] for m in messages if m.get("role") == "tool"] + self.assertEqual(tool_ids, ["x", "y"]) # first x kept, dup dropped, y intact + # the kept x is the FIRST result + self.assertEqual(next(m for m in messages if m.get("tool_call_id") == "x")["content"], "first") + + def test_no_op_when_unique(self): + from secator.ai.utils import _dedupe_tool_results + messages = [ + {"role": "tool", "tool_call_id": "a", "content": "1"}, + {"role": "tool", "tool_call_id": "b", "content": "2"}, + ] + before = [dict(m) for m in messages] + self.assertEqual(_dedupe_tool_results(messages), 0) + self.assertEqual(messages, before) + + def test_dedupe_scoped_per_consecutive_run(self): + """The same id in two SEPARATE tool runs (own assistant each) is not a dup.""" + from secator.ai.utils import _dedupe_tool_results + messages = [ + {"role": "tool", "tool_call_id": "x", "content": "r1"}, + {"role": "assistant", "content": "thinking"}, + {"role": "tool", "tool_call_id": "x", "content": "r2"}, + ] + self.assertEqual(_dedupe_tool_results(messages), 0) # separated by a non-tool msg + + def test_repair_dedupes_duplicate_tool_results(self): + """_repair_orphan_tool_uses now removes duplicates as part of its pass.""" + from secator.ai.utils import _repair_orphan_tool_uses + messages = [ + {"role": "assistant", "content": None, + "tool_calls": [{"id": "x", "type": "function", "function": {"name": "f", "arguments": "{}"}}]}, + {"role": "tool", "tool_call_id": "x", "name": "f", "content": "first"}, + {"role": "tool", "tool_call_id": "x", "name": "f", "content": "DUP"}, + {"role": "user", "content": "next"}, + ] + changed = _repair_orphan_tool_uses(messages) + self.assertGreaterEqual(changed, 1) + tool_ids = [m["tool_call_id"] for m in messages if m.get("role") == "tool"] + self.assertEqual(tool_ids, ["x"]) # exactly one result for x + + if __name__ == '__main__': unittest.main() diff --git a/tests/unit/test_api_hook_transport.py b/tests/unit/test_api_hook_transport.py new file mode 100644 index 000000000..61ac501fa --- /dev/null +++ b/tests/unit/test_api_hook_transport.py @@ -0,0 +1,54 @@ +"""API hook transport-security tests. + +The ``api`` driver POSTs runner/finding data — including raw targets, accumulated +command output and the Bearer API key — to the configured ``addons.api.url``. +``_make_request`` must refuse to do that over cleartext HTTP to a remote host, +regardless of ``force_ssl`` (which only governs TLS certificate *verification*). +Loopback http:// stays allowed for local dev. These tests pin that guard. +""" + +import unittest +from unittest import mock + +from secator.hooks import api as api_hook + + +class TestApiHookTransportGuard(unittest.TestCase): + def _check(self, url): + with mock.patch.object(api_hook, 'API_URL', url): + # Stop before any real network call: the guard runs first, so if it + # allows the request, requests.request is reached and we short-circuit. + with mock.patch.object(api_hook.requests, 'request') as req: + req.side_effect = RuntimeError('reached_network') + api_hook._make_request('GET', 'workspaces') + + def test_remote_http_rejected(self): + with self.assertRaises(Exception) as ctx: + self._check('http://app.secator.cloud/api') + self.assertIn('cleartext', str(ctx.exception).lower()) + + def test_remote_http_rejected_even_with_force_ssl_false(self): + with mock.patch.object(api_hook, 'FORCE_SSL', False): + with self.assertRaises(Exception) as ctx: + self._check('http://10.0.0.5:8081/api') + self.assertIn('cleartext', str(ctx.exception).lower()) + + def test_https_remote_allowed(self): + # Allowed by the guard -> proceeds to the network call (which we stub). + with self.assertRaises(RuntimeError) as ctx: + self._check('https://app.secator.cloud/api') + self.assertEqual(str(ctx.exception), 'reached_network') + + def test_http_localhost_allowed(self): + with self.assertRaises(RuntimeError) as ctx: + self._check('http://localhost:8081/api') + self.assertEqual(str(ctx.exception), 'reached_network') + + def test_http_loopback_ip_allowed(self): + with self.assertRaises(RuntimeError) as ctx: + self._check('http://127.0.0.1:8081/api') + self.assertEqual(str(ctx.exception), 'reached_network') + + +if __name__ == '__main__': + unittest.main() diff --git a/tests/unit/test_command_task.py b/tests/unit/test_command_task.py new file mode 100644 index 000000000..d6dc7eada --- /dev/null +++ b/tests/unit/test_command_task.py @@ -0,0 +1,176 @@ +import os +import unittest +from unittest import mock + +from secator.runners import Command +from secator.tasks.command import command + + +class TestCommandTask(unittest.TestCase): + """The generic `command` task runs an arbitrary command line verbatim in shell mode.""" + + def test_runs_verbatim_and_captures_stdout(self): + """A trivial echo runs through the real Command yielder and reaches SUCCESS with stdout captured.""" + runner = command(inputs=["echo secator-pr3"], run_opts={"sync": True, "print_line": False, "print_item": False}) + runner.run() + self.assertEqual(runner.status, "SUCCESS") + self.assertIn("secator-pr3", runner.output) + self.assertEqual(runner.return_code, 0) + + def test_shell_metacharacters_are_interpreted(self): + """Shell operators (&&) must be interpreted, not passed as literal echo args. + + Under shell=False the whole string is shlex-split and `&&` becomes a literal echo + argument, so only the first echo runs and 'world' never appears. This proves the + task actually runs in shell mode end-to-end. + """ + opts = {"sync": True, "print_line": False, "print_item": False} + runner = command(inputs=["echo hello && echo world"], run_opts=opts) + runner.run() + self.assertEqual(runner.status, "SUCCESS") + self.assertIn("hello", runner.output) + self.assertIn("world", runner.output) + # Under shell=False the whole thing is one echo, so '&&' is echoed literally. + # Interpreted correctly, '&&' is an operator and never appears in stdout. + self.assertNotIn("&&", runner.output) + self.assertTrue(runner.shell) + + def test_bare_single_word_command_is_not_stripped(self): + """A bare single-word command (no space) must run, not get type-filtered away. + + `input_types` must be [] — a non-empty input_types makes the base _validate_inputs() + run autodetect_type() and DROP inputs whose detected type isn't listed. "true" is + autodetected as 'slug', so [STRING] would strip it -> empty inputs -> empty cmd -> + FAILURE. This is the exact shape that broke most ordinary commands. + """ + runner = command(inputs=["true"], run_opts={"sync": True, "print_line": False, "print_item": False}) + runner.run() + self.assertEqual(runner.cmd, "true") + self.assertEqual(runner.status, "SUCCESS") + + def test_empty_inputs_does_not_crash(self): + """With no inputs, _build_cmd must not crash (cmd stays empty rather than indexing inputs[0]).""" + runner = command(inputs=[], run_opts={"sync": True, "print_line": False, "print_item": False}) + runner._build_cmd() + self.assertEqual(runner.cmd, "") + + def test_is_a_command_subclass(self): + """Sanity check on the inheritance the rest of the PR relies on.""" + self.assertTrue(issubclass(command, Command)) + + def test_env_run_opt_is_honored(self): + """A custom `env` run_opt overrides the process env for the subprocess. + + The AI shell handler relies on this to pass a SANITIZED env (LLM key / cloud + creds stripped) so an AI-run `env`/`printenv` can't leak them. PATH is included + so /bin/sh can still resolve the shell builtin. Opts are spread as kwargs + (command takes **run_opts) so `env` actually reaches self.run_opts. + """ + custom_env = {'FOO': 'bar', 'PATH': os.environ.get('PATH', '')} + runner = command( + ['echo $FOO'], + sync=True, print_line=False, print_item=False, env=custom_env, + ) + runner.run() + self.assertEqual(runner.status, 'SUCCESS') + self.assertIn('bar', runner.output) + + def test_no_env_run_opt_uses_process_env(self): + """Control: with no `env` run_opt, the subprocess inherits the process env + (default behavior unchanged).""" + with mock.patch.dict(os.environ, {'SECATOR_ENV_PROBE': 'present'}): + runner = command( + ['echo $SECATOR_ENV_PROBE'], + sync=True, print_line=False, print_item=False, + ) + runner.run() + self.assertEqual(runner.status, 'SUCCESS') + self.assertIn('present', runner.output) + + def test_empty_env_run_opt_is_honored(self): + """An explicit empty `env={}` must be honored (deliberate empty environment), + NOT silently fall back to the full process env — otherwise a caller asking for a + locked-down env would leak every process var. Guards the `.get('env', os.environ)` + (vs truthiness `or os.environ`) semantics. + """ + with mock.patch.dict(os.environ, {'ANTHROPIC_API_KEY': 'sk-leakme'}): + runner = command( + ['echo "[$ANTHROPIC_API_KEY]"'], + sync=True, print_line=False, print_item=False, env={}, + ) + runner.run() + self.assertEqual(runner.status, 'SUCCESS') + # With an empty env the var is unset, so the shell expands it to nothing. + self.assertNotIn('sk-leakme', runner.output) + self.assertIn('[]', runner.output) + + def test_env_run_opt_is_not_persisted(self): + """The `env` run_opt must be runtime-only, NOT persisted into the runner doc. + + `resolved_opts` -> `toDict()['run_opts']` is written to Mongo by the mongodb driver + and exposed via the runner-list API. Persisting `env` would write the whole (denylist- + sanitized) process environment — DB URLs, internal hostnames, etc. that the denylist + misses — into every AI-shell runner's history, a broader disclosure than the command's + own output. `resolved_opts` excludes `env` so the subprocess still reads it from + `run_opts` (command.py) but it never reaches the persisted doc. + """ + custom_env = {'FOO': 'bar', 'DATABASE_URL': 'postgres://u:pw@host', 'PATH': os.environ.get('PATH', '')} + runner = command(['echo hi'], sync=True, print_line=False, print_item=False, env=custom_env) + persisted_opts = runner.toDict()['run_opts'] + self.assertNotIn('env', persisted_opts) + # sanity: the runner still holds env at runtime for the subprocess to use + self.assertEqual(runner.run_opts.get('env'), custom_env) + + +class TestCommandFromResult(unittest.TestCase): + """`command.from_result` imports an already-run command's result into a runner doc, + without executing anything (the forward-looking seam for importing externally-run + commands into Secator Cloud). + """ + + @mock.patch("subprocess.Popen") + def test_from_result_populates_runner_without_executing(self, mock_popen): + """A successful imported result populates output/status and never spawns a subprocess.""" + runner = command.from_result("nmap -p80 x", "PORT 80 open", 0) + + self.assertEqual(runner.output, "PORT 80 open") + self.assertEqual(runner.status, "SUCCESS") + mock_popen.assert_not_called() + + data = runner.toDict() + self.assertEqual(data["cmd"], "nmap -p80 x") + self.assertEqual(data["output"], "PORT 80 open") + self.assertEqual(data["status"], "SUCCESS") + self.assertEqual(data["return_code"], 0) + + @mock.patch("subprocess.Popen") + def test_from_result_bare_command_success(self, mock_popen): + """A bare single-word command imports as SUCCESS with its cmd + output intact. + + Regression for the input-type-stripping bug: before the fix, "whoami" was + autodetected as 'slug' and dropped, so cmd came back '' and status FAILURE (the + spurious empty-input Error). The passed-in output is a fixed literal so the + assertion is deterministic (not machine-dependent). + """ + runner = command.from_result("whoami", "someoutput", 0) + + self.assertEqual(runner.status, "SUCCESS") + self.assertEqual(runner.cmd, "whoami") + self.assertEqual(runner.output, "someoutput") + mock_popen.assert_not_called() + + @mock.patch("subprocess.Popen") + def test_from_result_failure_preserves_output_verbatim(self, mock_popen): + """A non-zero return code yields FAILURE, and the caller's output is preserved verbatim. + + Regression for the output-corruption bug: the synthetic Error added on the FAILURE + path must NOT be appended onto self.output (add_result must be called output=False). + """ + runner = command.from_result("somecmd", "the real stdout", 1) + + self.assertEqual(runner.status, "FAILURE") + self.assertEqual(runner.toDict()["status"], "FAILURE") + # Exact match — no ANSI Error repr appended. + self.assertEqual(runner.output, "the real stdout") + self.assertEqual(runner.toDict()["output"], "the real stdout") + mock_popen.assert_not_called() diff --git a/tests/unit/test_eviction.py b/tests/unit/test_eviction.py new file mode 100644 index 000000000..194d85d09 --- /dev/null +++ b/tests/unit/test_eviction.py @@ -0,0 +1,82 @@ +import queue +import threading +import types +import unittest +import unittest.mock +from pathlib import Path +import shutil +import tempfile + +import secator.celery_signals as cs +from secator.celery_signals import ( + clear_shutdown_flag, + is_worker_shutting_down, + worker_shutting_down_handler, +) +from secator.runners import command as command_mod + + +class TestEvictionSelfFinalize(unittest.TestCase): + """Worker-eviction self-finalize: on shutdown (e.g. a K8s pod SIGTERM eviction) the in-flight + task's monitor stops it early and returns partial results, so the surrounding chord proceeds + instead of hanging until the broker visibility timeout redelivers the task.""" + + def setUp(self): + # Isolate the shutdown flag to a per-test temp path so tests can't interfere with each + # other (or with a real worker) through the shared global flag file. + self._tmpdir = tempfile.mkdtemp() + self._patcher = unittest.mock.patch.object(cs, 'SHUTDOWN_FLAG', Path(self._tmpdir) / 'worker_shutdown') + self._patcher.start() + clear_shutdown_flag() + + def tearDown(self): + clear_shutdown_flag() + self._patcher.stop() + shutil.rmtree(self._tmpdir, ignore_errors=True) + + def _bare_command(self, process): + """A Command shell with only the attributes _monitor_process touches (no real subprocess).""" + cmd = command_mod.Command.__new__(command_mod.Command) + cmd.process = process + cmd.monitor_stop_event = threading.Event() + cmd.monitor_queue = queue.Queue() + cmd.debug = lambda *a, **k: None + return cmd + + def test_shutdown_flag_lifecycle(self): + """worker_shutting_down_handler raises the flag; clear_shutdown_flag drops it.""" + self.assertFalse(is_worker_shutting_down()) + worker_shutting_down_handler() + self.assertTrue(is_worker_shutting_down()) + clear_shutdown_flag() + self.assertFalse(is_worker_shutting_down()) + + def test_monitor_stops_process_when_flag_set(self): + """When a shutdown is raised during the run, the monitor stops the process (exit_ok=True, so + the task returns partial results) and emits the eviction Warning — letting the chord proceed.""" + cmd = self._bare_command(types.SimpleNamespace(pid=999999)) + stopped = {} + cmd.stop_process = lambda **kw: (stopped.update(kw), cmd.monitor_stop_event.set()) + with unittest.mock.patch('secator.celery_signals.is_worker_shutting_down', return_value=True): + cmd._monitor_process() + self.assertTrue(stopped.get('exit_ok'), 'monitor did not stop the process on the shutdown flag') + queued = [] + while not cmd.monitor_queue.empty(): + queued.append(cmd.monitor_queue.get()) + self.assertTrue( + any('shutting down' in str(getattr(i, 'message', '')).lower() for i in queued), + 'monitor did not emit the eviction warning', + ) + + def test_monitor_clears_stale_flag_at_start(self): + """The monitor clears any pre-existing (stale) flag at start, so a flag left by a previous + worker sharing the state dir does not stop a fresh task. Regression for the integration leak: + a leaked flag had been self-aborting every later task.""" + cmd = self._bare_command(None) # process=None -> loop breaks right after the clear + with unittest.mock.patch('secator.celery_signals.clear_shutdown_flag') as mock_clear: + cmd._monitor_process() + mock_clear.assert_called_once() + + +if __name__ == '__main__': + unittest.main()