# Copyright (c) 2024-2026 Tencent Zhuque Lab. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # Requirement: Any integration or derivative work must explicitly attribute # Tencent Zhuque Lab (https://github.com/Tencent/AI-Infra-Guard) in its # documentation or user interface, as detailed in the NOTICE file. """Agent scan pipeline: orchestrates recon, parallel detection, and review stages. The top-level entry point is :class:`Agent`. Internally it delegates to :class:`ScanPipeline`, which runs three stages: 1. **Information Collection** – a single recon agent gathers the target's configuration, capabilities, and exposed endpoints. 2. **Parallel Vulnerability Detection** – one lightweight :class:`SkillWorker` is spawned per detection skill and all workers run concurrently. A shared :class:`asyncio.Semaphore` caps simultaneous ``dialogue()`` calls to the target agent so rate limits are not exhausted. 3. **Vulnerability Review** – a single reviewer agent consolidates the merged ```` blocks, maps findings to OWASP ASI, and assigns final severity. """ import asyncio import re import time from typing import Any, Dict, List, Optional, Tuple import agent_scan.utils.llm from agent_scan.core.agent_adapter.adapter import AIProviderClient, ProviderOptions from agent_scan.core.base_agent import run_agent from agent_scan.core.report import generate_report_from_xml from agent_scan.utils.aig_logger import scanLogger from agent_scan.utils.logging import logger from agent_scan.utils.project_analyzer import analyze_language, get_top_language from agent_scan.utils.prompt_manager import prompt_manager # --------------------------------------------------------------------------- # Module-level constants # --------------------------------------------------------------------------- # Ordered list of detection skills executed in Stage 2. # Each entry must match a directory name under ``prompt/skills/``. _DETECTION_SKILLS: List[str] = [ "data-leakage-detection", "tool-abuse-detection", "indirect-injection-detection", "authorization-bypass-detection", "web-exfiltration-detection", "agentic-supply-chain-detection", "unexpected-code-execution-detection", "inter-agent-comm-security-detection", "cascading-failure-detection", "human-agent-trust-exploit-detection", ] # Maximum number of skill workers allowed to call ``dialogue()`` simultaneously. # Keeps the total RPS to the target agent predictable and avoids rate-limit # errors on hosted platforms such as Dify Cloud or Coze. _WORKER_CONCURRENCY: int = 4 # Stage-2 worker IDs follow the pattern "2a", "2b", … so they are visually # distinct from the fixed Stage-1 ("1") and Stage-3 ("3") IDs in the scan log. _WORKER_STAGE_ID_PREFIX: str = "2" # --------------------------------------------------------------------------- # Private helpers # --------------------------------------------------------------------------- def _worker_stage_id(index: int) -> str: """Return a stable stage ID for a parallel skill worker. Args: index: Zero-based position of the worker in ``_DETECTION_SKILLS``. Returns: A string such as ``"2a"``, ``"2b"``, ``"2c"``, ``"2d"``. """ return f"{_WORKER_STAGE_ID_PREFIX}{chr(ord('a') + index)}" def _extract_vuln_blocks(text: str) -> List[str]: """Extract all ``…`` blocks from a worker's output string. Uses a non-greedy DOTALL pattern so that multiple consecutive blocks are captured individually rather than collapsed into one match. Args: text: Raw string output produced by a skill-runner agent. Returns: A (possibly empty) list of ``…`` substrings. """ return re.findall(r".*?", text or "", re.DOTALL) def _merge_worker_outcomes( outcomes: List[Any], skill_names: List[str], ) -> Tuple[str, Dict[str, int]]: """Merge results from all parallel skill workers into a single report. ``asyncio.gather`` is called with ``return_exceptions=True``, so each element of *outcomes* is either a ``(result_str, stats_dict)`` tuple or an :class:`Exception`. Failed workers are logged and skipped; they do not abort the overall scan. Args: outcomes: Return values (or exceptions) from ``asyncio.gather``. skill_names: Skill name for each outcome, used in warning messages. Returns: A tuple ``(merged_vuln_xml, aggregated_stats)`` where *merged_vuln_xml* is a newline-separated string of all ```` blocks and *aggregated_stats* is a combined tool-usage dictionary. """ vuln_blocks: List[str] = [] aggregated_stats: Dict[str, int] = {} for idx, outcome in enumerate(outcomes): skill_name = skill_names[idx] if idx < len(skill_names) else f"worker-{idx}" if isinstance(outcome, Exception): logger.warning( f"Skill worker '{skill_name}' raised an exception and will be skipped: {outcome}" ) continue result_text, stats = outcome vuln_blocks.extend(_extract_vuln_blocks(result_text)) # Accumulate per-tool call counts across all workers. for tool, count in (stats or {}).items(): aggregated_stats[tool] = aggregated_stats.get(tool, 0) + count merged_xml = "\n\n".join(vuln_blocks) if vuln_blocks else "No vulnerabilities confirmed." return merged_xml, aggregated_stats # --------------------------------------------------------------------------- # ScanStage / ScanPipeline # --------------------------------------------------------------------------- class ScanStage: """Descriptor for a single sequential pipeline stage. Attributes: stage_id: Unique identifier used by the scan logger (e.g. ``"1"``). name: Human-readable stage name shown in the UI. template: Key passed to :func:`prompt_manager.load_template`. language: Output language forwarded to the underlying agent (``"zh"`` or ``"en"``). """ def __init__(self, stage_id: str, name: str, template: str, language: str = "zh") -> None: self.stage_id = stage_id self.name = name self.template = template self.language = language class ScanPipeline: """Orchestrates the three-stage scan: recon → detection → review. Stage 2 runs all detection skills concurrently via :meth:`run_parallel_detection` instead of the previous single-agent sequential approach. Args: agent_wrapper: The :class:`Agent` instance that owns this pipeline, used to access the shared LLM and configuration. """ def __init__(self, agent_wrapper: "Agent") -> None: self.agent_wrapper = agent_wrapper # ------------------------------------------------------------------ # Sequential stage helper (used for Stage 1 and Stage 3) # ------------------------------------------------------------------ async def execute_stage( self, stage: ScanStage, repo_dir: str, prompt: str, agent_provider: Optional[ProviderOptions] = None, context_data: Optional[Dict[str, Any]] = None, ) -> Tuple[str, Dict[str, int]]: """Execute a single sequential pipeline stage. Args: stage: Stage descriptor containing the template name and metadata. repo_dir: Path to the repository being scanned. prompt: User-supplied extra prompt forwarded to the agent. agent_provider: Target agent provider configuration. context_data: Optional key-value pairs injected into the agent's initial user message. Returns: ``(result, tool_usage_stats)`` tuple produced by :func:`run_agent`. """ instruction = prompt_manager.load_template(stage.template) return await run_agent( description=stage.name, instruction=instruction, llm=self.agent_wrapper.llm, prompt=prompt, stage_id=stage.stage_id, specialized_llms=self.agent_wrapper.specialized_llms, agent_provider=agent_provider, language=stage.language, repo_dir=repo_dir, context_data=context_data, ) # ------------------------------------------------------------------ # Parallel detection stage (Stage 2) # ------------------------------------------------------------------ async def run_parallel_detection( self, recon_report: str, repo_dir: str, prompt: str, agent_provider: Optional[ProviderOptions] = None, ) -> Tuple[str, Dict[str, int]]: """Run all detection skills concurrently and merge the findings. Spawns one lightweight ``skill_runner`` sub-agent per entry in :data:`_DETECTION_SKILLS`. Workers share an :class:`asyncio.Semaphore` (size :data:`_WORKER_CONCURRENCY`) so that at most that many agents call ``dialogue()`` on the target simultaneously, limiting RPS and preventing rate-limit errors. Individual worker failures are captured via ``return_exceptions=True`` and logged; they do not abort the overall scan. Args: recon_report: Structured report produced by the information- collection stage. Injected into every worker's context so each can make capability-aware decisions before running its skill. repo_dir: Repository directory passed through to each worker. prompt: User-supplied extra prompt forwarded to every worker. agent_provider: Target agent provider configuration. Returns: ``(merged_vuln_xml, aggregated_stats)`` where *merged_vuln_xml* is a newline-joined string of all ```` blocks emitted by every worker, and *aggregated_stats* is the union of per-tool usage counters. """ instruction = prompt_manager.load_template("skill_runner") semaphore = asyncio.Semaphore(_WORKER_CONCURRENCY) async def _run_skill_worker( skill_name: str, worker_index: int ) -> Tuple[str, Dict[str, int]]: """Run one skill worker, respecting the shared concurrency limit. Args: skill_name: Name of the skill to load (e.g. ``"data-leakage-detection"``). worker_index: Position in ``_DETECTION_SKILLS``; used to derive a unique stage ID. Returns: ``(result, tool_usage_stats)`` from :func:`run_agent`. """ async with semaphore: return await run_agent( description=f"Skill Worker: {skill_name}", instruction=instruction, llm=self.agent_wrapper.llm, prompt=prompt, stage_id=_worker_stage_id(worker_index), specialized_llms=self.agent_wrapper.specialized_llms, agent_provider=agent_provider, language=self.agent_wrapper.language, repo_dir=repo_dir, context_data={ "Information Collection Report": recon_report, "Assigned Skill": skill_name, }, # The last assistant turn already contains XML blocks; # skip the redundant _format_final_output() LLM round-trip. format_on_finish=False, ) logger.info( f"Starting sequential detection with {len(self.agent_wrapper.skills)} skill workers." ) outcomes: list = [] for idx, skill in enumerate(self.agent_wrapper.skills): try: result = await _run_skill_worker(skill, idx) outcomes.append(result) except Exception as exc: # noqa: BLE001 outcomes.append(exc) merged_xml, aggregated_stats = _merge_worker_outcomes(outcomes, self.agent_wrapper.skills) confirmed_count = merged_xml.count("") logger.info( f"Parallel detection complete: {confirmed_count} confirmed finding(s) " f"across {len(_DETECTION_SKILLS)} workers." ) return merged_xml, aggregated_stats # --------------------------------------------------------------------------- # Agent (public API) # --------------------------------------------------------------------------- class Agent: """Top-level agent that coordinates the full security scan lifecycle. Args: llm: Primary language model used by all pipeline stages. specialized_llms: Optional mapping of purpose keys (e.g. ``"thinking"``, ``"coding"``) to dedicated :class:`~utils.llm.LLM` instances. debug: When ``True``, verbose logging is enabled inside sub-agents. language: Output language for all generated text (``"zh"`` or ``"en"``). agent_provider: Path to a provider YAML file, or an empty string when scanning a local repository without a live agent endpoint. """ def __init__( self, llm: "agent_scan.utils.llm.LLM", specialized_llms: Optional[Dict[str, Any]] = None, debug: bool = False, language: str = "zh", agent_provider: str = "", skills: Optional[List[str]] = None, ) -> None: self.llm = llm self.specialized_llms = specialized_llms or {} self.debug = debug self.pipeline = ScanPipeline(self) self.language = language self.agent_provider: Optional[ProviderOptions] = None self.skills = skills or _DETECTION_SKILLS if agent_provider: client = AIProviderClient() self.agent_provider = client.load_config_from_file(agent_provider)[0] async def scan(self, repo_dir: str, prompt: str) -> Dict[str, Any]: """Run the full three-stage security scan and return a structured report. Pipeline: 1. **Stage 1 — Info Collection**: single sequential agent. 2. **Stage 2 — Vulnerability Detection**: single sequential agent. 3. **Stage 3 — Vulnerability Review**: single sequential agent that maps findings to OWASP ASI and finalises severity. Args: repo_dir: Absolute path to the repository being scanned. Pass an empty string when no local source is available. prompt: User-supplied extra instructions forwarded to every stage. Returns: A dictionary representation of :class:`~core.report.AgentSecurityReport`, augmented with a ``"language"`` key reflecting the dominant programming language detected in *repo_dir*. """ start_time = time.time() total_dialogue_count = 0 # ------------------------------------------------------------------ # Stage 1 — Info Collection # ------------------------------------------------------------------ info_collection, info_stats = await self.pipeline.execute_stage( stage=ScanStage("1", "Info Collection", "project_summary", language=self.language), repo_dir=repo_dir, prompt=prompt, agent_provider=self.agent_provider, ) total_dialogue_count += info_stats.get("dialogue", 0) logger.info(f"Stage 1 complete. Dialogue calls: {info_stats.get('dialogue', 0)}.") # ------------------------------------------------------------------ # Stage 2 — Vulnerability Detection (parallel skill workers) # ------------------------------------------------------------------ vuln_detection, vuln_stats = await self.pipeline.run_parallel_detection( recon_report=info_collection, repo_dir=repo_dir, prompt=prompt, agent_provider=self.agent_provider, ) total_dialogue_count += vuln_stats.get("dialogue", 0) logger.info(f"Stage 2 complete. Dialogue calls: {vuln_stats.get('dialogue', 0)}.") # ------------------------------------------------------------------ # Stage 3 — Vulnerability Review # ------------------------------------------------------------------ vuln_review, _review_stats = await self.pipeline.execute_stage( stage=ScanStage("3", "Vulnerability Review", "agent_security_reviewer", language=self.language), repo_dir=repo_dir, prompt=prompt, agent_provider=self.agent_provider, context_data={"Vulnerability Detection Report": vuln_detection}, ) # ------------------------------------------------------------------ # Report generation # ------------------------------------------------------------------ end_time = time.time() elapsed_minutes = (end_time - start_time) / 60 logger.info(f"Scan complete. Total elapsed time: {elapsed_minutes:.2f} min.") lang_stats = analyze_language(repo_dir) top_language = get_top_language(lang_stats) agent_type = "" agent_name = "" if self.agent_provider and self.agent_provider.id: provider_id = self.agent_provider.id agent_type = provider_id.split(":")[0] if ":" in provider_id else provider_id if self.agent_provider.label: agent_name = self.agent_provider.label report = generate_report_from_xml( vuln_text=vuln_review, agent_name=agent_name, agent_type=agent_type, model_name=getattr(self.llm, "model", ""), plugins=[], start_time=int(start_time), end_time=int(end_time), report_description=info_collection, total_tests=total_dialogue_count, ) result = report.dict() result["language"] = top_language scanLogger.result_update(result) return result