diff --git a/secator/ai/actions.py b/secator/ai/actions.py index 1fd07a099..c2c1def39 100644 --- a/secator/ai/actions.py +++ b/secator/ai/actions.py @@ -28,6 +28,8 @@ class ActionContext: """ targets: List[str] model: str + api_key: str = "" + api_base: str = "" encryptor: Any = None dry_run: bool = False verbose: bool = False @@ -525,6 +527,50 @@ def _sanitize_child_opts(opts: Any) -> Dict: 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 _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 _run_runner(action: Dict, ctx: ActionContext, runner_type: str) -> Generator: """Execute a secator task or workflow. @@ -548,6 +594,20 @@ def _run_runner(action: Dict, ctx: ActionContext, runner_type: str) -> Generator return opts["subagent"] = True opts["interactive"] = False + # Inherit the parent's resolved LLM config so the subagent can actually run. + # Without this it falls back to CONFIG.addons.ai.default_model, which may be a + # different provider than the parent (e.g. anthropic-direct vs openrouter) with + # no key set -> AuthenticationError before the subagent does anything. 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 diff --git a/secator/tasks/ai.py b/secator/tasks/ai.py index 1a0b2c4fd..cdbc87c28 100644 --- a/secator/tasks/ai.py +++ b/secator/tasks/ai.py @@ -405,6 +405,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, diff --git a/tests/unit/test_ai_actions.py b/tests/unit/test_ai_actions.py index 84cf0821f..ca5fac58a 100644 --- a/tests/unit/test_ai_actions.py +++ b/tests/unit/test_ai_actions.py @@ -557,6 +557,62 @@ def test_run_runner_preserves_existing_session_id(self, mock_build_hooks, mock_t _, 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): @@ -1372,5 +1428,59 @@ def test_child_context_strips_parent_identity_keeps_session(self): 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()