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

jarvis/jarvis_sec/utils.py

@@ -0,0 +1,499 @@
+# -*- coding: utf-8 -*-
+"""工具函数模块"""
+
+from typing import Dict, List, Optional
+from pathlib import Path
+import json
+import typer
+
+from jarvis.jarvis_sec.workflow import direct_scan
+
+
+def git_restore_if_dirty(repo_root: str) -> int:
+    """
+    若 repo_root 为 git 仓库：检测工作区是否有变更；如有则使用 'git checkout -- .' 恢复。
+    返回估算的变更文件数（基于 git status --porcelain 的行数）。
+    """
+    try:
+        import subprocess as _sub
+        root = Path(repo_root)
+        if not (root / ".git").exists():
+            return 0
+        proc = _sub.run(["git", "status", "--porcelain"], cwd=str(root), capture_output=True, text=True)
+        if proc.returncode != 0:
+            return 0
+        lines = [line for line in proc.stdout.splitlines() if line.strip()]
+        if lines:
+            _sub.run(["git", "checkout", "--", "."], cwd=str(root), capture_output=True, text=True)
+            return len(lines)
+    except Exception:
+        pass
+    return 0
+
+
+def get_sec_dir(base_path: str) -> Path:
+    """获取 .jarvis/sec 目录路径，支持 base_path 是项目根目录或已经是 .jarvis/sec 目录"""
+    base = Path(base_path)
+    # 检查 base_path 是否已经是 .jarvis/sec 目录
+    if base.name == "sec" and base.parent.name == ".jarvis":
+        return base
+    # 否则，假设 base_path 是项目根目录
+    return base / ".jarvis" / "sec"
+
+
+def initialize_analysis_context(
+    entry_path: str,
+    status_mgr,
+) -> tuple:
+    """
+    初始化分析上下文，包括状态管理、进度文件、目录等。
+    
+    返回: (sec_dir, progress_path, _progress_append)
+    """
+    # 获取 .jarvis/sec 目录
+    sec_dir = get_sec_dir(entry_path)
+    progress_path = None  # 不再使用 progress.jsonl
+    
+    # 进度追加函数（空函数，不再记录）
+    def _progress_append(rec: Dict) -> None:
+        pass  # 不再记录进度日志
+    
+    return sec_dir, progress_path, _progress_append
+
+
+def load_or_run_heuristic_scan(
+    entry_path: str,
+    langs: List[str],
+    exclude_dirs: Optional[List[str]],
+    sec_dir: Path,
+    status_mgr,
+    _progress_append,
+) -> tuple[List[Dict], Dict]:
+    """
+    加载或运行启发式扫描。
+    
+    优先从新的 candidates.jsonl 文件加载，如果不存在则回退到旧的 heuristic_issues.jsonl。
+    
+    返回: (candidates, summary)
+    """
+    candidates: List[Dict] = []
+    summary: Dict = {}
+    
+    # 优先使用新的 candidates.jsonl 文件
+    from jarvis.jarvis_sec.file_manager import load_candidates, get_candidates_file
+    candidates = load_candidates(sec_dir)
+    
+    if candidates:
+        try:
+            typer.secho(f"[jarvis-sec] 从 {get_candidates_file(sec_dir)} 恢复启发式扫描", fg=typer.colors.BLUE)
+            _progress_append({
+                "event": "pre_scan_resumed",
+                "path": str(get_candidates_file(sec_dir)),
+                "issues_found": len(candidates)
+            })
+        except Exception:
+            pass
+    else:
+        # 回退到旧的 heuristic_issues.jsonl 文件（向后兼容）
+        _heuristic_path = sec_dir / "heuristic_issues.jsonl"
+        if _heuristic_path.exists():
+            try:
+                typer.secho(f"[jarvis-sec] 从 {_heuristic_path} 恢复启发式扫描（旧格式）", fg=typer.colors.BLUE)
+                with _heuristic_path.open("r", encoding="utf-8") as f:
+                    for line in f:
+                        if line.strip():
+                            candidates.append(json.loads(line))
+                _progress_append({
+                    "event": "pre_scan_resumed",
+                    "path": str(_heuristic_path),
+                    "issues_found": len(candidates)
+                })
+            except Exception as e:
+                typer.secho(f"[jarvis-sec] 恢复启发式扫描失败，执行完整扫描: {e}", fg=typer.colors.YELLOW)
+                candidates = []  # 重置以便执行完整扫描
+    
+    if not candidates:
+        _progress_append({"event": "pre_scan_start", "entry_path": entry_path, "languages": langs})
+        status_mgr.update_pre_scan(message="开始启发式扫描...")
+        pre_scan = direct_scan(entry_path, languages=langs, exclude_dirs=exclude_dirs)
+        candidates = pre_scan.get("issues", [])
+        summary = pre_scan.get("summary", {})
+        scanned_files = summary.get("scanned_files", 0)
+        status_mgr.update_pre_scan(
+            current_files=scanned_files,
+            total_files=scanned_files,
+            issues_found=len(candidates),
+            message=f"启发式扫描完成，发现 {len(candidates)} 个候选问题"
+        )
+        _progress_append({
+            "event": "pre_scan_done",
+            "entry_path": entry_path,
+            "languages": langs,
+            "scanned_files": scanned_files,
+            "issues_found": len(candidates)
+        })
+        # 持久化
+        try:
+            _heuristic_path.parent.mkdir(parents=True, exist_ok=True)
+            with _heuristic_path.open("w", encoding="utf-8") as f:
+                for item in candidates:
+                    f.write(json.dumps(item, ensure_ascii=False) + "\n")
+            _progress_append({
+                "event": "heuristic_report_written",
+                "path": str(_heuristic_path),
+                "issues_count": len(candidates),
+            })
+            typer.secho(f"[jarvis-sec] 已将 {len(candidates)} 个启发式扫描问题写入 {_heuristic_path}", fg=typer.colors.GREEN)
+        except Exception:
+            pass
+    else:
+        # 从断点恢复启发式扫描结果
+        status_mgr.update_pre_scan(
+            issues_found=len(candidates),
+            message=f"从断点恢复，已发现 {len(candidates)} 个候选问题"
+        )
+    
+    return candidates, summary
+
+
+def compact_candidate(it: Dict) -> Dict:
+    """精简候选问题，只保留必要字段"""
+    result = {
+        "language": it.get("language"),
+        "category": it.get("category"),
+        "pattern": it.get("pattern"),
+        "file": it.get("file"),
+        "line": it.get("line"),
+        "evidence": it.get("evidence"),
+        "confidence": it.get("confidence"),
+        "severity": it.get("severity", "medium"),
+    }
+    # 如果候选已经有gid，保留它（用于断点恢复）
+    if "gid" in it:
+        try:
+            gid_val = int(it.get("gid", 0))
+            if gid_val >= 1:
+                result["gid"] = gid_val
+        except Exception:
+            pass
+    return result
+
+
+def prepare_candidates(candidates: List[Dict]) -> List[Dict]:
+    """
+    将候选问题精简为子任务清单，控制上下文长度，并分配全局唯一ID。
+    
+    返回: compact_candidates (已分配gid的候选列表)
+    """
+    compact_candidates = [compact_candidate(it) for it in candidates]
+    
+    # 检查是否所有候选都已经有gid（从heuristic_issues.jsonl恢复时）
+    all_have_gid = all("gid" in it and isinstance(it.get("gid"), int) and it.get("gid", 0) >= 1 for it in compact_candidates)
+    
+    if not all_have_gid:
+        # 如果有候选没有gid，需要分配
+        # 优先保留已有的gid，为没有gid的候选分配新的gid
+        existing_gids = set()
+        for it in compact_candidates:
+            try:
+                gid_val = int(it.get("gid", 0))
+                if gid_val >= 1:
+                    existing_gids.add(gid_val)
+            except Exception:
+                pass
+        
+        # 为没有gid的候选分配新的gid
+        next_gid = 1
+        for it in compact_candidates:
+            if "gid" not in it or not isinstance(it.get("gid"), int) or it.get("gid", 0) < 1:
+                # 找到一个未使用的gid
+                while next_gid in existing_gids:
+                    next_gid += 1
+                try:
+                    it["gid"] = next_gid
+                    existing_gids.add(next_gid)
+                    next_gid += 1
+                except Exception:
+                    pass
+    
+    return compact_candidates
+
+
+def group_candidates_by_file(candidates: List[Dict]) -> Dict[str, List[Dict]]:
+    """按文件分组候选问题"""
+    from collections import defaultdict
+    groups: Dict[str, List[Dict]] = defaultdict(list)
+    for it in candidates:
+        groups[str(it.get("file") or "")].append(it)
+    return groups
+
+
+def create_report_writer(sec_dir: Path, report_file: Optional[str]):
+    """创建报告写入函数"""
+    from jarvis.jarvis_sec.file_manager import save_analysis_result, load_clusters
+    
+    def _append_report(items, source: str, task_id: str, cand: Dict):
+        """
+        将当前子任务的检测结果追加写入 analysis.jsonl 文件。
+        
+        参数:
+        - items: 验证通过的问题列表（has_risk: true）
+        - source: 来源（"analysis_only" 或 "verified"）
+        - task_id: 任务ID（如 "JARVIS-SEC-Batch-1"）
+        - cand: 候选信息，包含 batch 和 candidates
+        """
+        if not items:
+            return
+        
+        try:
+            # 从批次中提取信息
+            batch = cand.get("batch", False)
+            candidates = cand.get("candidates", [])
+            
+            if not batch or not candidates:
+                # 如果没有批次信息，回退到旧格式（向后兼容）
+                path = Path(report_file) if report_file else sec_dir / "agent_issues.jsonl"
+                path.parent.mkdir(parents=True, exist_ok=True)
+                with path.open("a", encoding="utf-8") as f:
+                    for item in items:
+                        line = json.dumps(item, ensure_ascii=False)
+                        f.write(line + "\n")
+                try:
+                    typer.secho(f"[jarvis-sec] 已将 {len(items)} 个问题写入 {path}（旧格式）", fg=typer.colors.GREEN)
+                except Exception:
+                    pass
+                return
+            
+            # 从批次中提取 file 和 gids
+            batch_file = candidates[0].get("file") if candidates else ""
+            batch_gids = []
+            for c in candidates:
+                try:
+                    gid = int(c.get("gid", 0))
+                    if gid >= 1:
+                        batch_gids.append(gid)
+                except Exception:
+                    pass
+            
+            # 从 clusters.jsonl 中查找对应的 cluster_id
+            clusters = load_clusters(sec_dir)
+            cluster_id = None
+            batch_index = None
+            cluster_index = None
+            
+            # 尝试从 task_id 中提取 batch_index（格式：JARVIS-SEC-Batch-1）
+            try:
+                if "Batch-" in task_id:
+                    batch_index = int(task_id.split("Batch-")[1])
+            except Exception:
+                pass
+            
+            # 查找匹配的聚类（通过 file 和 gids）
+            for cluster in clusters:
+                cluster_file = str(cluster.get("file", ""))
+                cluster_gids = cluster.get("gids", [])
+                
+                if cluster_file == batch_file and set(cluster_gids) == set(batch_gids):
+                    cluster_id = cluster.get("cluster_id", "")
+                    if not cluster_id:
+                        # 如果没有 cluster_id，生成一个
+                        cluster_id = f"{cluster_file}|{cluster.get('batch_index', batch_index or 0)}|{cluster.get('cluster_index', 0)}"
+                    batch_index = cluster.get("batch_index", batch_index or 0)
+                    cluster_index = cluster.get("cluster_index", 0)
+                    break
+            
+            # 如果找不到匹配的聚类，生成一个临时的 cluster_id
+            if not cluster_id:
+                cluster_id = f"{batch_file}|{batch_index or 0}|0"
+                batch_index = batch_index or 0
+                cluster_index = 0
+            
+            # 分离验证为问题的gid和误报的gid
+            verified_gids = []
+            false_positive_gids = []
+            issues = []
+            
+            # 从 items 中提取已验证的问题
+            for item in items:
+                try:
+                    gid = int(item.get("gid", 0))
+                    if gid >= 1:
+                        has_risk = item.get("has_risk", False)
+                        if has_risk:
+                            verified_gids.append(gid)
+                            issues.append(item)
+                        else:
+                            false_positive_gids.append(gid)
+                except Exception:
+                    pass
+            
+            # 从 candidates 中提取所有未在 items 中的 gid（这些可能是误报）
+            for c in candidates:
+                try:
+                    gid = int(c.get("gid", 0))
+                    if gid >= 1 and gid not in verified_gids and gid not in false_positive_gids:
+                        # 如果这个 gid 不在已验证的问题中，可能是误报
+                        false_positive_gids.append(gid)
+                except Exception:
+                    pass
+            
+            # 构建分析结果记录
+            analysis_result = {
+                "cluster_id": cluster_id,
+                "file": batch_file,
+                "batch_index": batch_index,
+                "cluster_index": cluster_index,
+                "gids": batch_gids,
+                "verified_gids": verified_gids,
+                "false_positive_gids": false_positive_gids,
+                "issues": issues,
+            }
+            
+            # 保存到 analysis.jsonl
+            save_analysis_result(sec_dir, analysis_result)
+            
+            try:
+                typer.secho(f"[jarvis-sec] 已将批次 {batch_index} 的分析结果写入 analysis.jsonl（问题: {len(verified_gids)}, 误报: {len(false_positive_gids)}）", fg=typer.colors.GREEN)
+            except Exception:
+                pass
+        except Exception as e:
+            # 报告写入失败不影响主流程
+            try:
+                typer.secho(f"[jarvis-sec] 警告：保存分析结果失败: {e}", fg=typer.colors.YELLOW)
+            except Exception:
+                pass
+    
+    return _append_report
+
+
+def sig_of(c: Dict) -> str:
+    """生成候选问题的签名"""
+    return f"{c.get('language','')}|{c.get('file','')}|{c.get('line','')}|{c.get('pattern','')}"
+
+
+def load_processed_gids_from_issues(sec_dir: Path) -> set:
+    """从 agent_issues.jsonl 中读取已处理的 gid"""
+    processed_gids = set()
+    try:
+        _agent_issues_path = sec_dir / "agent_issues.jsonl"
+        if _agent_issues_path.exists():
+            with _agent_issues_path.open("r", encoding="utf-8", errors="ignore") as f:
+                for line in f:
+                    line = line.strip()
+                    if not line:
+                        continue
+                    try:
+                        issue_obj = json.loads(line)
+                        _gid = int(issue_obj.get("gid", 0))
+                        if _gid >= 1:
+                            processed_gids.add(_gid)
+                    except Exception:
+                        pass
+            if processed_gids:
+                try:
+                    typer.secho(f"[jarvis-sec] 断点恢复：从 agent_issues.jsonl 读取到 {len(processed_gids)} 个已处理的 gid", fg=typer.colors.BLUE)
+                except Exception:
+                    pass
+    except Exception:
+        pass
+    return processed_gids
+
+
+def count_issues_from_file(sec_dir: Path) -> int:
+    """从 analysis.jsonl 读取问题数量"""
+    from jarvis.jarvis_sec.file_manager import get_verified_issue_gids
+    verified_gids = get_verified_issue_gids(sec_dir)
+    return len(verified_gids)
+
+
+def count_issues_from_file_old(sec_dir: Path) -> int:
+    """从 agent_issues.jsonl 中读取当前问题总数（用于状态显示）"""
+    count = 0
+    try:
+        _agent_issues_path = sec_dir / "agent_issues.jsonl"
+        if _agent_issues_path.exists():
+            saved_gids = set()
+            with _agent_issues_path.open("r", encoding="utf-8", errors="ignore") as f:
+                for line in f:
+                    line = line.strip()
+                    if not line:
+                        continue
+                    try:
+                        item = json.loads(line)
+                        gid = item.get("gid", 0)
+                        if gid >= 1 and gid not in saved_gids:
+                            # 只统计验证通过的告警（has_risk: true 且有 verification_notes）
+                            if item.get("has_risk") is True and "verification_notes" in item:
+                                count += 1
+                                saved_gids.add(gid)
+                    except Exception:
+                        pass
+    except Exception:
+        pass
+    return count
+
+
+def load_all_issues_from_file(sec_dir: Path) -> List[Dict]:
+    """从 agent_issues.jsonl 读取所有已保存的告警"""
+    all_issues: List[Dict] = []
+    try:
+        _agent_issues_path = sec_dir / "agent_issues.jsonl"
+        if _agent_issues_path.exists():
+            saved_gids_from_file = set()
+            with _agent_issues_path.open("r", encoding="utf-8", errors="ignore") as f:
+                for line in f:
+                    line = line.strip()
+                    if not line:
+                        continue
+                    try:
+                        item = json.loads(line)
+                        gid = item.get("gid", 0)
+                        if gid >= 1 and gid not in saved_gids_from_file:
+                            # 只保留验证通过的告警（has_risk: true 且有 verification_notes）
+                            if item.get("has_risk") is True and "verification_notes" in item:
+                                all_issues.append(item)
+                                saved_gids_from_file.add(gid)
+                    except Exception:
+                        pass
+            
+            if all_issues:
+                try:
+                    typer.secho(f"[jarvis-sec] 从 agent_issues.jsonl 加载了 {len(all_issues)} 个已保存的告警", fg=typer.colors.BLUE)
+                except Exception:
+                    pass
+        else:
+            try:
+                typer.secho("[jarvis-sec] agent_issues.jsonl 不存在，当前运行未发现任何问题", fg=typer.colors.BLUE)
+            except Exception:
+                pass
+    except Exception as e:
+        # 加载失败不影响主流程
+        try:
+            typer.secho(f"[jarvis-sec] 警告：从 agent_issues.jsonl 加载告警失败: {e}", fg=typer.colors.YELLOW)
+        except Exception:
+            pass
+    return all_issues
+
+
+def load_processed_gids_from_agent_issues(sec_dir: Path) -> set:
+    """从 agent_issues.jsonl 读取已处理的 gid"""
+    processed_gids = set()
+    try:
+        _agent_issues_path = sec_dir / "agent_issues.jsonl"
+        if _agent_issues_path.exists():
+            with _agent_issues_path.open("r", encoding="utf-8", errors="ignore") as f:
+                for line in f:
+                    line = line.strip()
+                    if not line:
+                        continue
+                    try:
+                        issue_obj = json.loads(line)
+                        _gid = int(issue_obj.get("gid", 0))
+                        if _gid >= 1:
+                            processed_gids.add(_gid)
+                    except Exception:
+                        pass
+    except Exception:
+        pass
+    return processed_gids
+