jarvis-ai-assistant 0.7.16__py3-none-any.whl → 1.0.2__py3-none-any.whl

jarvis/jarvis_sec/utils.py

@@ -1,10 +1,13 @@
 # -*- coding: utf-8 -*-
 """工具函数模块"""
 
-from typing import Dict, List, Optional
-from pathlib import Path
 import json
-import typer
+from pathlib import Path
+from typing import Dict
+from typing import List
+from typing import Optional
+
+from jarvis.jarvis_utils.output import PrettyOutput
 
 from jarvis.jarvis_sec.workflow import direct_scan
 
@@ -16,15 +19,26 @@ def git_restore_if_dirty(repo_root: str) -> int:
     """
     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)
+        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)
+            _sub.run(
+                ["git", "checkout", "--", "."],
+                cwd=str(root),
+                capture_output=True,
+                text=True,
+            )
             return len(lines)
     except Exception:
         pass
@@ -47,17 +61,17 @@ def initialize_analysis_context(
 ) -> 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
 
 
@@ -71,26 +85,33 @@ def load_or_run_heuristic_scan(
 ) -> 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
+    from jarvis.jarvis_sec.file_manager import get_candidates_file
+    from jarvis.jarvis_sec.file_manager import load_candidates
+
     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)
-            })
+            PrettyOutput.auto_print(
+                f"✨ [jarvis-sec] 从 {get_candidates_file(sec_dir)} 恢复启发式扫描",
+                timestamp=True,
+            )
+            _progress_append(
+                {
+                    "event": "pre_scan_resumed",
+                    "path": str(get_candidates_file(sec_dir)),
+                    "issues_found": len(candidates),
+                }
+            )
         except Exception:
             pass
     else:
@@ -98,22 +119,32 @@ def load_or_run_heuristic_scan(
         _heuristic_path = sec_dir / "heuristic_issues.jsonl"
         if _heuristic_path.exists():
             try:
-                typer.secho(f"[jarvis-sec] 从 {_heuristic_path} 恢复启发式扫描（旧格式）", fg=typer.colors.BLUE)
+                PrettyOutput.auto_print(
+                    f"✨ [jarvis-sec] 从 {_heuristic_path} 恢复启发式扫描（旧格式）",
+                    timestamp=True,
+                )
                 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)
-                })
+                _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)
+                PrettyOutput.auto_print(
+                    f"⚠️ [jarvis-sec] 恢复启发式扫描失败，执行完整扫描: {e}",
+                    timestamp=True,
+                )
                 candidates = []  # 重置以便执行完整扫描
-    
+
     if not candidates:
-        _progress_append({"event": "pre_scan_start", "entry_path": entry_path, "languages": langs})
+        _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", [])
@@ -123,36 +154,43 @@ def load_or_run_heuristic_scan(
             current_files=scanned_files,
             total_files=scanned_files,
             issues_found=len(candidates),
-            message=f"启发式扫描完成，发现 {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),
+            }
         )
-        _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)
+            _progress_append(
+                {
+                    "event": "heuristic_report_written",
+                    "path": str(_heuristic_path),
+                    "issues_count": len(candidates),
+                }
+            )
+            PrettyOutput.auto_print(
+                f"✅ [jarvis-sec] 已将 {len(candidates)} 个启发式扫描问题写入 {_heuristic_path}",
+                timestamp=True,
+            )
         except Exception:
             pass
     else:
         # 从断点恢复启发式扫描结果
         status_mgr.update_pre_scan(
             issues_found=len(candidates),
-            message=f"从断点恢复，已发现 {len(candidates)} 个候选问题"
+            message=f"从断点恢复，已发现 {len(candidates)} 个候选问题",
         )
-    
+
     return candidates, summary
 
 
@@ -182,14 +220,17 @@ def compact_candidate(it: Dict) -> Dict:
 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)
-    
+    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
@@ -201,11 +242,15 @@ def prepare_candidates(candidates: List[Dict]) -> List[Dict]:
                     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:
+            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
@@ -215,13 +260,14 @@ def prepare_candidates(candidates: List[Dict]) -> List[Dict]:
                     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)
@@ -230,12 +276,13 @@ def group_candidates_by_file(candidates: List[Dict]) -> Dict[str, List[Dict]]:
 
 def create_report_writer(sec_dir: Path, report_file: Optional[str]):
     """创建报告写入函数"""
-    from jarvis.jarvis_sec.file_manager import save_analysis_result, load_clusters
-    
+    from jarvis.jarvis_sec.file_manager import load_clusters
+    from jarvis.jarvis_sec.file_manager import save_analysis_result
+
     def _append_report(items, source: str, task_id: str, cand: Dict):
         """
         将当前子任务的检测结果追加写入 analysis.jsonl 文件。
-        
+
         参数:
         - items: 验证通过的问题列表（has_risk: true）
         - source: 来源（"analysis_only" 或 "verified"）
@@ -244,26 +291,31 @@ def create_report_writer(sec_dir: Path, report_file: Optional[str]):
         """
         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 = (
+                    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)
+                    PrettyOutput.auto_print(
+                        f"✅ [jarvis-sec] 已将 {len(items)} 个问题写入 {path}（旧格式）",
+                        timestamp=True,
+                    )
                 except Exception:
                     pass
                 return
-            
+
             # 从批次中提取 file 和 gids
             batch_file = candidates[0].get("file") if candidates else ""
             batch_gids = []
@@ -274,25 +326,25 @@ def create_report_writer(sec_dir: Path, report_file: Optional[str]):
                         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:
@@ -301,18 +353,18 @@ def create_report_writer(sec_dir: Path, report_file: Optional[str]):
                     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:
@@ -326,17 +378,21 @@ def create_report_writer(sec_dir: Path, report_file: Optional[str]):
                             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:
+                    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,
@@ -348,27 +404,30 @@ def create_report_writer(sec_dir: Path, report_file: Optional[str]):
                 "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)
+                PrettyOutput.auto_print(
+                    f"✅ [jarvis-sec] 已将批次 {batch_index} 的分析结果写入 analysis.jsonl（问题: {len(verified_gids)}, 误报: {len(false_positive_gids)}）",
+                    timestamp=True,
+                )
             except Exception:
                 pass
         except Exception as e:
             # 报告写入失败不影响主流程
             try:
-                typer.secho(f"[jarvis-sec] 警告：保存分析结果失败: {e}", fg=typer.colors.YELLOW)
+                PrettyOutput.auto_print(f"⚠️ [jarvis-sec] 警告：保存分析结果失败: {e}")
             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','')}"
+    return f"{c.get('language', '')}|{c.get('file', '')}|{c.get('line', '')}|{c.get('pattern', '')}"
 
 
 def load_processed_gids_from_issues(sec_dir: Path) -> set:
@@ -391,7 +450,10 @@ def load_processed_gids_from_issues(sec_dir: Path) -> set:
                         pass
             if processed_gids:
                 try:
-                    typer.secho(f"[jarvis-sec] 断点恢复：从 agent_issues.jsonl 读取到 {len(processed_gids)} 个已处理的 gid", fg=typer.colors.BLUE)
+                    PrettyOutput.auto_print(
+                        f"✨ [jarvis-sec] 断点恢复：从 agent_issues.jsonl 读取到 {len(processed_gids)} 个已处理的 gid",
+                        timestamp=True,
+                    )
                 except Exception:
                     pass
     except Exception:
@@ -402,6 +464,7 @@ def load_processed_gids_from_issues(sec_dir: Path) -> set:
 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)
 
@@ -423,7 +486,10 @@ def count_issues_from_file_old(sec_dir: Path) -> int:
                         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:
+                            if (
+                                item.get("has_risk") is True
+                                and "verification_notes" in item
+                            ):
                                 count += 1
                                 saved_gids.add(gid)
                     except Exception:
@@ -450,26 +516,38 @@ def load_all_issues_from_file(sec_dir: Path) -> List[Dict]:
                         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:
+                            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)
+                    PrettyOutput.auto_print(
+                        f"✨ [jarvis-sec] 从 agent_issues.jsonl 加载了 {len(all_issues)} 个已保存的告警",
+                        timestamp=True,
+                    )
                 except Exception:
                     pass
         else:
             try:
-                typer.secho("[jarvis-sec] agent_issues.jsonl 不存在，当前运行未发现任何问题", fg=typer.colors.BLUE)
+                PrettyOutput.auto_print(
+                    "✨ [jarvis-sec] agent_issues.jsonl 不存在，当前运行未发现任何问题",
+                    timestamp=True,
+                )
             except Exception:
                 pass
     except Exception as e:
         # 加载失败不影响主流程
         try:
-            typer.secho(f"[jarvis-sec] 警告：从 agent_issues.jsonl 加载告警失败: {e}", fg=typer.colors.YELLOW)
+            PrettyOutput.auto_print(
+                f"⚠️ [jarvis-sec] 警告：从 agent_issues.jsonl 加载告警失败: {e}",
+                timestamp=True,
+            )
         except Exception:
             pass
     return all_issues
@@ -496,4 +574,3 @@ def load_processed_gids_from_agent_issues(sec_dir: Path) -> set:
     except Exception:
         pass
     return processed_gids
-