jarvis-ai-assistant 0.7.0__py3-none-any.whl → 0.7.6__py3-none-any.whl

jarvis/jarvis_sec/agents.py

@@ -0,0 +1,143 @@
+# -*- coding: utf-8 -*-
+"""Agent创建和订阅模块"""
+
+from typing import Dict, Optional
+from jarvis.jarvis_agent import Agent
+from jarvis.jarvis_tools.registry import ToolRegistry
+from jarvis.jarvis_sec.prompts import (
+    build_summary_prompt,
+    get_review_system_prompt,
+    get_review_summary_prompt,
+    get_cluster_system_prompt,
+    get_cluster_summary_prompt,
+)
+
+
+def subscribe_summary_event(agent: Agent) -> Dict[str, str]:
+    """订阅Agent摘要事件"""
+    summary_container: Dict[str, str] = {"text": ""}
+    try:
+        from jarvis.jarvis_agent.events import AFTER_SUMMARY as _AFTER_SUMMARY
+    except Exception:
+        _AFTER_SUMMARY = None
+    
+    if _AFTER_SUMMARY:
+        def _on_after_summary(**kwargs):
+            try:
+                summary_container["text"] = str(kwargs.get("summary", "") or "")
+            except Exception:
+                summary_container["text"] = ""
+        try:
+            agent.event_bus.subscribe(_AFTER_SUMMARY, _on_after_summary)
+        except Exception:
+            pass
+    return summary_container
+
+
+def create_analysis_agent(task_id: str, llm_group: Optional[str], force_save_memory: bool = False) -> Agent:
+    """创建分析Agent"""
+    system_prompt = """
+# 单Agent安全分析约束
+- 你的核心任务是评估代码的安全问题，目标：针对本候选问题进行证据核实、风险评估与修复建议补充，查找漏洞触发路径，确认在某些条件下会触发；以此来判断是否是漏洞。
+- **必须进行调用路径推导**：
+  - 对于每个候选问题，必须明确推导从可控输入到缺陷代码的完整调用路径。
+  - 调用路径推导必须包括：
+    1. 识别可控输入的来源（例如：用户输入、网络数据、文件读取、命令行参数等）
+    2. 追踪数据流：从输入源开始，逐步追踪数据如何传递到缺陷代码位置
+    3. 识别调用链：明确列出从入口函数到缺陷代码的所有函数调用序列（例如：main() -> parse_input() -> process_data() -> vulnerable_function()）
+    4. 分析每个调用点的数据校验情况：检查每个函数是否对输入进行了校验、边界检查或安全检查
+    5. 确认触发条件：明确说明在什么条件下，未校验或恶意输入能够到达缺陷代码位置
+  - 如果无法推导出完整的调用路径，或者所有调用路径都有充分的保护措施，则应该判定为误报。
+  - 调用路径推导必须在分析过程中明确展示，不能省略或假设。
+- 工具优先：使用 read_code 读取目标文件附近源码（行号前后各 ~50 行），必要时用 execute_script 辅助检索。
+- **调用路径追溯要求**：
+  - 必须向上追溯所有可能的调用者，查看完整的调用路径，以确认风险是否真实存在。
+  - 使用 read_code 和 execute_script 工具查找函数的调用者（例如：使用 grep 搜索函数名，查找所有调用该函数的位置）。
+  - 对于每个调用者，必须检查其是否对输入进行了校验。
+  - 如果发现任何调用路径未做校验，必须明确记录该路径。
+  - 例如：一个函数存在空指针解引用风险，必须检查所有调用者。如果所有调用者均能确保传入的指针非空，则该风险在当前代码库中可能不会实际触发；但如果存在任何调用者未做校验，则风险真实存在。
+- 若多条告警位于同一文件且行号相距不远，可一次性读取共享上下文，对这些相邻告警进行联合分析与判断；但仍需避免无关扩展与大范围遍历。
+- 禁止修改任何文件或执行写操作命令（rm/mv/cp/echo >、sed -i、git、patch、chmod、chown 等）；仅进行只读分析与读取。
+- 每次仅执行一个操作；等待工具结果后再进行下一步。
+- **记忆使用**：
+  - 在分析过程中，充分利用 retrieve_memory 工具检索已有的记忆，特别是与当前分析函数相关的记忆。
+  - 如果有必要，使用 save_memory 工具保存每个函数的分析要点，使用函数名作为 tag（例如：函数名、文件名等）。
+  - 记忆内容示例：某个函数的指针已经判空、某个函数已有输入校验、某个函数的调用路径分析结果等。
+  - 这样可以避免重复分析，提高效率，并保持分析的一致性。
+- 完成对本批次候选问题的判断后，主输出仅打印结束符 <!!!COMPLETE!!!> ，不需要汇总结果。
+""".strip()
+    
+    agent_kwargs: Dict = dict(
+        system_prompt=system_prompt,
+        name=task_id,
+        auto_complete=True,
+        need_summary=True,
+        summary_prompt=build_summary_prompt(),
+        non_interactive=True,
+        in_multi_agent=False,
+        use_methodology=False,
+        use_analysis=False,
+        output_handler=[ToolRegistry()],
+        force_save_memory=force_save_memory,
+        use_tools=["read_code", "execute_script", "save_memory", "retrieve_memory"],
+    )
+    if llm_group:
+        agent_kwargs["model_group"] = llm_group
+    return Agent(**agent_kwargs)
+
+
+def create_review_agent(
+    current_review_num: int,
+    llm_group: Optional[str],
+) -> Agent:
+    """创建复核Agent"""
+    review_system_prompt = get_review_system_prompt()
+    review_summary_prompt = get_review_summary_prompt()
+    
+    review_task_id = f"JARVIS-SEC-Review-Batch-{current_review_num}"
+    review_agent_kwargs: Dict = dict(
+        system_prompt=review_system_prompt,
+        name=review_task_id,
+        auto_complete=True,
+        need_summary=True,
+        summary_prompt=review_summary_prompt,
+        non_interactive=True,
+        in_multi_agent=False,
+        use_methodology=False,
+        use_analysis=False,
+        output_handler=[ToolRegistry()],
+        use_tools=["read_code", "execute_script", "retrieve_memory", "save_memory"],
+    )
+    if llm_group:
+        review_agent_kwargs["model_group"] = llm_group
+    return Agent(**review_agent_kwargs)
+
+
+def create_cluster_agent(
+    file: str,
+    chunk_idx: int,
+    llm_group: Optional[str],
+    force_save_memory: bool = False,
+) -> Agent:
+    """创建聚类Agent"""
+    cluster_system_prompt = get_cluster_system_prompt()
+    cluster_summary_prompt = get_cluster_summary_prompt()
+    
+    agent_kwargs_cluster: Dict = dict(
+        system_prompt=cluster_system_prompt,
+        name=f"JARVIS-SEC-Cluster::{file}::batch{chunk_idx}",
+        auto_complete=True,
+        need_summary=True,
+        summary_prompt=cluster_summary_prompt,
+        non_interactive=True,
+        in_multi_agent=False,
+        use_methodology=False,
+        use_analysis=False,
+        output_handler=[ToolRegistry()],
+        force_save_memory=force_save_memory,
+        use_tools=["read_code", "execute_script", "save_memory", "retrieve_memory"],
+    )
+    if llm_group:
+        agent_kwargs_cluster["model_group"] = llm_group
+    return Agent(**agent_kwargs_cluster)
+

jarvis/jarvis_sec/analysis.py

@@ -0,0 +1,276 @@
+# -*- coding: utf-8 -*-
+"""分析相关模块"""
+
+from typing import Dict, List, Optional
+import typer
+
+from jarvis.jarvis_agent import Agent
+from jarvis.jarvis_sec.prompts import build_summary_prompt
+from jarvis.jarvis_sec.parsers import try_parse_summary_report
+from jarvis.jarvis_sec.utils import git_restore_if_dirty
+
+
+def valid_items(items: Optional[List]) -> bool:
+    """验证分析结果项的格式"""
+    if not isinstance(items, list):
+        return False
+    for it in items:
+        if not isinstance(it, dict):
+            return False
+        has_gid = "gid" in it
+        has_gids = "gids" in it
+        if not has_gid and not has_gids:
+            return False
+        if has_gid and has_gids:
+            return False
+        if has_gid:
+            try:
+                if int(it["gid"]) < 1:
+                    return False
+            except Exception:
+                return False
+        elif has_gids:
+            if not isinstance(it["gids"], list) or len(it["gids"]) == 0:
+                return False
+            for gid_val in it["gids"]:
+                try:
+                    if int(gid_val) < 1:
+                        return False
+                except Exception:
+                    return False
+        if "has_risk" not in it or not isinstance(it["has_risk"], bool):
+            return False
+        if it.get("has_risk"):
+            for key in ["preconditions", "trigger_path", "consequences", "suggestions"]:
+                if key not in it:
+                    return False
+                if not isinstance(it[key], str) or not it[key].strip():
+                    return False
+    return True
+
+
+def build_analysis_task_context(batch: List[Dict], entry_path: str, langs: List[str]) -> str:
+    """构建分析任务上下文"""
+    import json as _json2
+    batch_ctx: List[Dict] = list(batch)
+    cluster_verify = str(batch_ctx[0].get("verify") if batch_ctx else "")
+    cluster_gids_ctx = [it.get("gid") for it in batch_ctx]
+    return f"""
+# 安全子任务批次
+上下文参数：
+- entry_path: {entry_path}
+- languages: {langs}
+- cluster_verification: {cluster_verify}
+
+- cluster_gids: {cluster_gids_ctx}
+- note: 每个候选含 gid/verify 字段，模型仅需输出 gid 统一给出验证/判断结论（全局编号）；无需使用局部 id
+
+批次候选(JSON数组):
+{_json2.dumps(batch_ctx, ensure_ascii=False, indent=2)}
+""".strip()
+
+
+def build_validation_error_guidance(
+    parse_error_analysis: Optional[str],
+    prev_parsed_items: Optional[List],
+) -> str:
+    """构建验证错误指导信息"""
+    if parse_error_analysis:
+        return f"\n\n**格式错误详情（请根据以下错误修复输出格式）：**\n- JSON解析失败: {parse_error_analysis}\n\n请确保输出的JSON格式正确，包括正确的引号、逗号、大括号等。仅输出一个 <REPORT> 块，块内直接包含 JSON 数组（不需要额外的标签）。支持jsonnet语法（如尾随逗号、注释、||| 或 ``` 分隔符多行字符串等）。"
+    elif prev_parsed_items is None:
+        return "\n\n**格式错误详情（请根据以下错误修复输出格式）：**\n- 无法从摘要中解析出有效的 JSON 数组"
+    elif not valid_items(prev_parsed_items):
+        validation_errors = []
+        if not isinstance(prev_parsed_items, list):
+            validation_errors.append("结果不是数组")
+        else:
+            for idx, it in enumerate(prev_parsed_items):
+                if not isinstance(it, dict):
+                    validation_errors.append(f"元素{idx}不是字典")
+                    break
+                has_gid = "gid" in it
+                has_gids = "gids" in it
+                if not has_gid and not has_gids:
+                    validation_errors.append(f"元素{idx}缺少必填字段 gid 或 gids")
+                    break
+                if has_gid and has_gids:
+                    validation_errors.append(f"元素{idx}不能同时包含 gid 和 gids")
+                    break
+                if has_gid:
+                    try:
+                        if int(it.get("gid", 0)) < 1:
+                            validation_errors.append(f"元素{idx}的 gid 必须 >= 1")
+                            break
+                    except Exception:
+                        validation_errors.append(f"元素{idx}的 gid 格式错误（必须是整数）")
+                        break
+                elif has_gids:
+                    if not isinstance(it.get("gids"), list) or len(it.get("gids", [])) == 0:
+                        validation_errors.append(f"元素{idx}的 gids 必须是非空数组")
+                        break
+                    try:
+                        for gid_idx, gid_val in enumerate(it.get("gids", [])):
+                            if int(gid_val) < 1:
+                                validation_errors.append(f"元素{idx}的 gids[{gid_idx}] 必须 >= 1")
+                                break
+                        if validation_errors:
+                            break
+                    except Exception:
+                        validation_errors.append(f"元素{idx}的 gids 格式错误（必须是整数数组）")
+                        break
+                if "has_risk" not in it or not isinstance(it.get("has_risk"), bool):
+                    validation_errors.append(f"元素{idx}缺少必填字段 has_risk（必须是布尔值）")
+                    break
+                if it.get("has_risk"):
+                    for key in ["preconditions", "trigger_path", "consequences", "suggestions"]:
+                        if key not in it:
+                            validation_errors.append(f"元素{idx}的 has_risk 为 true，但缺少必填字段 {key}")
+                            break
+                        if not isinstance(it[key], str) or not it[key].strip():
+                            validation_errors.append(f"元素{idx}的 {key} 字段不能为空")
+                            break
+                    if validation_errors:
+                        break
+        if validation_errors:
+            return "\n\n**格式错误详情（请根据以下错误修复输出格式）：**\n" + "\n".join(f"- {err}" for err in validation_errors)
+    return ""
+
+
+def run_analysis_agent_with_retry(
+    agent: Agent,
+    per_task: str,
+    summary_container: Dict[str, str],
+    entry_path: str,
+    task_id: str,
+    bidx: int,
+    meta_records: List[Dict],
+) -> tuple[Optional[List[Dict]], Optional[Dict]]:
+    """运行分析Agent并重试直到成功"""
+    summary_items: Optional[List[Dict]] = None
+    workspace_restore_info: Optional[Dict] = None
+    use_direct_model_analysis = False
+    prev_parsed_items: Optional[List] = None
+    parse_error_analysis: Optional[str] = None
+    attempt = 0
+    
+    while True:
+        attempt += 1
+        summary_container["text"] = ""
+        
+        if use_direct_model_analysis:
+            summary_prompt_text = build_summary_prompt()
+            error_guidance = build_validation_error_guidance(parse_error_analysis, prev_parsed_items)
+            full_prompt = f"{per_task}{error_guidance}\n\n{summary_prompt_text}"
+            try:
+                response = agent.model.chat_until_success(full_prompt)  # type: ignore
+                summary_container["text"] = response
+            except Exception as e:
+                try:
+                    typer.secho(f"[jarvis-sec] 直接模型调用失败: {e}，回退到 run()", fg=typer.colors.YELLOW)
+                except Exception:
+                    pass
+                agent.run(per_task)
+        else:
+            agent.run(per_task)
+
+        # 工作区保护
+        try:
+            _changed = git_restore_if_dirty(entry_path)
+            workspace_restore_info = {
+                "performed": bool(_changed),
+                "changed_files_count": int(_changed or 0),
+                "action": "git checkout -- .",
+            }
+            meta_records.append({
+                "task_id": task_id,
+                "batch_index": bidx,
+                "workspace_restore": workspace_restore_info,
+                "attempt": attempt + 1,
+            })
+            if _changed:
+                try:
+                    typer.secho(f"[jarvis-sec] 工作区已恢复 ({_changed} 个文件），操作: git checkout -- .", fg=typer.colors.BLUE)
+                except Exception:
+                    pass
+        except Exception:
+            pass
+
+        # 解析摘要中的 <REPORT>（JSON）
+        summary_text = summary_container.get("text", "")
+        parsed_items: Optional[List] = None
+        parse_error_analysis = None
+        if summary_text:
+            rep, parse_error_analysis = try_parse_summary_report(summary_text)
+            if parse_error_analysis:
+                try:
+                    typer.secho(f"[jarvis-sec] 分析结果JSON解析失败: {parse_error_analysis}", fg=typer.colors.YELLOW)
+                except Exception:
+                    pass
+            elif isinstance(rep, list):
+                parsed_items = rep
+            elif isinstance(rep, dict):
+                items = rep.get("issues")
+                if isinstance(items, list):
+                    parsed_items = items
+
+        # 关键字段校验
+        # 空数组 [] 是有效的（表示没有发现问题），需要单独处理
+        if parsed_items is not None:
+            if len(parsed_items) == 0:
+                # 空数组表示没有发现问题，这是有效的格式
+                summary_items = parsed_items
+                break
+            elif valid_items(parsed_items):
+                # 非空数组需要验证格式
+                summary_items = parsed_items
+                break
+        
+        # 格式校验失败，后续重试使用直接模型调用
+        use_direct_model_analysis = True
+        prev_parsed_items = parsed_items
+        if parse_error_analysis:
+            try:
+                typer.secho(f"[jarvis-sec] 分析结果JSON解析失败 -> 重试第 {attempt} 次 (批次={bidx}，使用直接模型调用，将反馈解析错误)", fg=typer.colors.YELLOW)
+            except Exception:
+                pass
+        else:
+            try:
+                typer.secho(f"[jarvis-sec] 分析结果格式无效 -> 重试第 {attempt} 次 (批次={bidx}，使用直接模型调用)", fg=typer.colors.YELLOW)
+            except Exception:
+                pass
+    
+    return summary_items, workspace_restore_info
+
+
+def expand_and_filter_analysis_results(summary_items: List[Dict]) -> tuple[List[Dict], List[Dict]]:
+    """展开gids格式为单个gid格式，并过滤出有风险的项目"""
+    items_with_risk: List[Dict] = []
+    items_without_risk: List[Dict] = []
+    merged_items: List[Dict] = []
+    
+    for it in summary_items:
+        has_risk = it.get("has_risk") is True
+        if "gids" in it and isinstance(it.get("gids"), list):
+            for gid_val in it.get("gids", []):
+                try:
+                    gid_int = int(gid_val)
+                    if gid_int >= 1:
+                        item = {
+                            **{k: v for k, v in it.items() if k != "gids"},
+                            "gid": gid_int,
+                        }
+                        if has_risk:
+                            merged_items.append(item)
+                            items_with_risk.append(item)
+                        else:
+                            items_without_risk.append(item)
+                except Exception:
+                    pass
+        elif "gid" in it:
+            if has_risk:
+                merged_items.append(it)
+                items_with_risk.append(it)
+            else:
+                items_without_risk.append(it)
+    
+    return items_with_risk, items_without_risk

jarvis/jarvis_sec/cli.py

@@ -5,10 +5,12 @@ Jarvis 安全演进套件 —— 命令行入口（Typer 版本）
 用法示例：
 - Agent模式（单Agent，逐条子任务分析）
   python -m jarvis.jarvis_sec.cli agent --path ./target_project
+  python -m jarvis.jarvis_sec.cli agent  # 默认分析当前目录
 
 可选参数：
 
-- --output: 最终Markdown报告输出路径（默认 ./report.md）
+- --path, -p: 待分析的根目录（默认当前目录）
+- --output, -o: 最终报告输出路径（默认 ./report.md）。如果后缀为 .csv，则输出 CSV 格式；否则输出 Markdown 格式
 """
 
 from __future__ import annotations
@@ -20,7 +22,8 @@ from typing import Optional
 import typer
 from jarvis.jarvis_utils.utils import init_env
 # removed: set_config import（避免全局覆盖模型组配置）
-from jarvis.jarvis_sec.workflow import run_with_agent, direct_scan, format_markdown_report
+from jarvis.jarvis_sec.workflow import run_with_agent, direct_scan, format_markdown_report as format_markdown_report_workflow
+from jarvis.jarvis_sec.report import format_csv_report, aggregate_issues
 
 app = typer.Typer(
     add_completion=False,
@@ -33,22 +36,28 @@ app = typer.Typer(
 
 @app.command("agent", help="Agent模式（单Agent逐条子任务分析）")
 def agent(
-    path: str = typer.Option(..., "--path", "-p", help="待分析的根目录"),
+    path: str = typer.Option(".", "--path", "-p", help="待分析的根目录（默认当前目录）"),
 
     llm_group: Optional[str] = typer.Option(
         None, "--llm-group", "-g", help="使用的模型组（仅对本次运行生效，不修改全局配置）"
     ),
     output: Optional[str] = typer.Option(
-        "report.md", "--output", "-o", help="最终Markdown报告输出路径（默认 ./report.md）"
+        "report.md", "--output", "-o", help="最终报告输出路径（默认 ./report.md）。如果后缀为 .csv，则输出 CSV 格式；否则输出 Markdown 格式"
     ),
 
     cluster_limit: int = typer.Option(
         50, "--cluster-limit", "-c", help="聚类每批最多处理的告警数（按文件分批聚类，默认50）"
     ),
+    enable_verification: bool = typer.Option(
+        True, "--enable-verification/--no-verification", help="是否启用二次验证（默认开启）"
+    ),
+    force_save_memory: bool = typer.Option(
+        False, "--force-save-memory/--no-force-save-memory", help="强制使用记忆（默认关闭）"
+    ),
 ) -> None:
     # 初始化环境，确保平台/模型等全局配置就绪（避免 NoneType 平台）
     try:
-        init_env("欢迎使用 Jarvis 安全套件！", None)
+        init_env()
     except Exception:
         # 环境初始化失败不应阻塞CLI基础功能，继续后续流程
         pass
@@ -62,6 +71,9 @@ def agent(
             path,
             llm_group=llm_group,
             cluster_limit=cluster_limit,
+            enable_verification=enable_verification,
+            force_save_memory=force_save_memory,
+            output_file=output,
         )
     except Exception as e:
         try:
@@ -76,7 +88,18 @@ def agent(
         except Exception:
             pass
         result = direct_scan(path)
-        text = format_markdown_report(result)
+        # 根据输出文件后缀选择格式
+        if output and output.lower().endswith('.csv'):
+            # 使用 report.py 中的函数来格式化 CSV
+            report_json = aggregate_issues(
+                result.get("issues", []),
+                scanned_root=result.get("summary", {}).get("scanned_root"),
+                scanned_files=result.get("summary", {}).get("scanned_files"),
+            )
+            text = format_csv_report(report_json)
+        else:
+            # 使用 workflow.py 中的 format_markdown_report（与 direct_scan 返回结构匹配）
+            text = format_markdown_report_workflow(result)
 
     if output:
         try: