@comate/zulu 1.2.1-beta.2 → 1.3.1

package/comate-engine/assets/skills/code-security-comate/scripts/repair_vulnerability.py

@@ -0,0 +1,245 @@
+#!/usr/bin/env python3
+"""
+漏洞修复工具 - 根据扫描结果修复指定文件的漏洞。
+
+用法:
+    python3 repair_vulnerability.py --root-path <项目目录> --username <用户名> --parsed-result <parsed_result.json路径>
+    python3 repair_vulnerability.py --root-path <项目目录> --username <用户名> --vulnerability-info '<漏洞信息JSON>'
+
+支持两种输入方式（二选一）:
+    --parsed-result: 直接传入 parsed_result.json 路径，脚本自动提取普通漏洞并构建修复信息
+    --vulnerability-info: 传入漏洞信息 JSON 字符串
+
+流程:
+    1. 解析漏洞信息，计算待修复文件的哈希
+    2. 调用修复接口
+    3. 如有缺失文件则上传后重试
+    4. 返回包含 diff_content 的修复结果
+"""
+
+import argparse
+import difflib
+import json
+import logging
+import os
+import sys
+import time
+
+import http_client  # noqa: F401 (triggers shared logging config)
+import utils
+
+logger = logging.getLogger("repair")
+
+# 修复轮询最大次数（每次间隔 3 秒，约 10 分钟）
+MAX_REPAIR_POLLS = 200
+REPAIR_POLL_INTERVAL = 3
+MAX_UPLOAD_RETRIES = 10
+
+
+def diff_file_content(file_path, new_content):
+    # type: (str, str) -> str
+    """比较原文件与修复后内容，使用 difflib 生成精确 diff，返回 diff 列表 JSON。
+
+    每个不连续的变更区域（hunk）生成独立的 {from_content, to_content} 对，
+    确保每个 from_content 都是原文件中的连续子串，可以被 str.replace 精确匹配。
+    """
+    with open(file_path, "r", encoding="utf-8") as f:
+        old_content = f.read()
+
+    old_lines = old_content.splitlines(keepends=True)
+    new_lines = new_content.splitlines(keepends=True)
+
+    # 使用 unified_diff 找出所有差异区域，n=0 不提供上下文行
+    diff = list(difflib.unified_diff(old_lines, new_lines, n=0))
+
+    diffs = []  # type: list
+    current_from = []  # type: list
+    current_to = []  # type: list
+
+    def _flush():
+        if current_from or current_to:
+            diffs.append({
+                "from_content": "".join(current_from),
+                "to_content": "".join(current_to),
+            })
+            del current_from[:]
+            del current_to[:]
+
+    for line in diff:
+        if line.startswith("---") or line.startswith("+++"):
+            continue
+        elif line.startswith("@@"):
+            # 新 hunk 开始，先保存上一组
+            _flush()
+        elif line.startswith("-"):
+            current_from.append(line[1:])
+        elif line.startswith("+"):
+            current_to.append(line[1:])
+
+    # 保存最后一组
+    _flush()
+
+    return json.dumps(diffs, ensure_ascii=False, indent=2)
+
+
+def upload_files_for_repair(root_path, missing_files, username, user_id):
+    # type: (str, list, str, str) -> None
+    """上传修复过程中缺失的文件。"""
+    payload = {"files": {}}
+    for name in missing_files:
+        file_path = os.path.join(root_path, name)
+        try:
+            content = utils.read_file_content(file_path)
+            file_hash = utils.calc_sha256(file_path)
+            payload["files"][name] = {"hash": file_hash, "content": content}
+        except Exception as e:
+            print("警告: 读取文件失败 {}: {}".format(file_path, e), file=sys.stderr)
+
+    http_client.put(
+        "{}/api/v1/upload".format(utils.HOST),
+        headers=utils.build_headers(username, user_id),
+        json_body=payload,
+    )
+
+
+def repair_vulnerability(root_path, vulnerability_info, username, user_id):
+    # type: (str, dict, str, str) -> dict
+    """执行漏洞修复，返回包含 diff 的修复结果。"""
+    # 计算待修复文件的哈希
+    for file_info in vulnerability_info.get("files", []):
+        file_path = os.path.join(root_path, file_info["name"])
+        if os.path.isfile(file_path):
+            file_info["hash"] = utils.calc_sha256(file_path)
+
+    repair_type = vulnerability_info.get("type", 2)
+    file_count = len(vulnerability_info.get("files", []))
+    logger.info("repair start: type=%d, files=%d", repair_type, file_count)
+    api_url = "{}/api/v2/repair_file".format(utils.HOST)
+    print("开始修复...", file=sys.stderr)
+
+    poll_count = 0
+    upload_retry_count = 0
+    while poll_count < MAX_REPAIR_POLLS:
+        result = http_client.post(
+            api_url,
+            headers=utils.build_headers(username, user_id),
+            json_body=vulnerability_info,
+        )
+
+        # 如有缺失文件，先上传
+        missing = result.get("data", {}).get("missingFiles", [])
+        if missing:
+            upload_retry_count += 1
+            if upload_retry_count > MAX_UPLOAD_RETRIES:
+                print("错误: 修复上传重试 {} 次后仍有缺失文件".format(MAX_UPLOAD_RETRIES), file=sys.stderr)
+                return {"status": -1, "message": "修复上传重试次数超限"}
+            print("上传缺失文件: {} 个 (重试 {}/{})".format(len(missing), 
+                                                   upload_retry_count, MAX_UPLOAD_RETRIES), file=sys.stderr)
+            upload_files_for_repair(root_path, missing, username, user_id)
+            time.sleep(REPAIR_POLL_INTERVAL)
+            continue
+
+        # status != 0 表示完成
+        if result.get("status") != 0:
+            # 生成 diff
+            files_data = result.get("data", {}).get("files", [])
+            for file_data in files_data:
+                repaired_content = file_data.get("repairedContent", "")
+                if repaired_content:
+                    file_path = os.path.join(root_path, file_data["name"])
+                    if os.path.isfile(file_path):
+                        diff_json = diff_file_content(file_path, repaired_content)
+                        file_data["diff_content"] = diff_json
+                    # 置空防止泄露完整源码
+                    file_data["repairedContent"] = ""
+            return result
+
+        poll_count += 1
+        print("修复中，等待结果... ({}/{})".format(poll_count, MAX_REPAIR_POLLS), file=sys.stderr)
+        time.sleep(REPAIR_POLL_INTERVAL)
+
+    print("错误: 修复超时，已轮询 {} 次仍未完成".format(MAX_REPAIR_POLLS), file=sys.stderr)
+    return {"status": -1, "message": "修复超时"}
+
+
+def build_vulnerability_info(parsed):
+    # type: (dict) -> dict
+    """从 parsed_result.json 构建修复接口所需的 vulnerability-info。"""
+    bundle_hash = parsed.get("bundle_hash", "")
+    common_vuls = parsed.get("common_vuls", [])
+
+    file_map = {}  # type: dict
+    for vul in common_vuls:
+        fname = vul.get("file", "")
+        if not fname:
+            continue
+        if fname not in file_map:
+            file_map[fname] = {"name": fname, "vulList": []}
+        file_map[fname]["vulList"].append({
+            "ruleID": vul.get("ruleID", ""),
+            "line": vul.get("startLine", 0),
+            "hash": vul.get("hash", ""),
+        })
+
+    return {
+        "files": list(file_map.values()),
+        "type": 2,
+        "extra": {
+            "bundleHash": bundle_hash,
+        },
+    }
+
+
+def main():
+    parser = argparse.ArgumentParser(description="代码安全漏洞修复工具")
+    parser.add_argument("--root-path", required=True, help="项目根目录")
+    parser.add_argument("--username", required=True, help="Comate 用户名")
+    parser.add_argument("--vulnerability-info", default=None, help="漏洞信息 JSON 字符串")
+    parser.add_argument("--parsed-result", default=None,
+                        help="解析结果文件路径 (parsed_result.json)，与 --vulnerability-info 二选一")
+    parser.add_argument("--output-dir", default=None, help="结果输出目录，默认为 skill 临时目录")
+    args = parser.parse_args()
+
+    if not args.vulnerability_info and not args.parsed_result:
+        print("错误: 必须提供 --vulnerability-info 或 --parsed-result 之一", file=sys.stderr)
+        sys.exit(1)
+
+    user_id = utils.make_user_id(args.username)
+
+    logger.info("repair_vulnerability start: username=%s, root_path=%s", args.username, args.root_path)
+
+    root_path = os.path.realpath(args.root_path)
+    if not os.path.isdir(root_path):
+        print("错误: 目录不存在 {}".format(root_path), file=sys.stderr)
+        sys.exit(1)
+
+    output_dir = utils.get_output_dir(root_path, args.output_dir)
+
+    if args.parsed_result:
+        try:
+            with open(args.parsed_result, "r", encoding="utf-8") as f:
+                parsed = json.load(f)
+        except Exception as e:
+            print("错误: 读取解析结果失败 {}: {}".format(args.parsed_result, e), file=sys.stderr)
+            sys.exit(1)
+        if not parsed.get("common_vuls"):
+            print("无普通漏洞需要修复", file=sys.stderr)
+            sys.exit(0)
+        vulnerability_info = build_vulnerability_info(parsed)
+    else:
+        try:
+            vulnerability_info = json.loads(args.vulnerability_info)
+        except json.JSONDecodeError as e:
+            print("错误: 漏洞信息 JSON 解析失败: {}".format(e), file=sys.stderr)
+            sys.exit(1)
+
+    result = repair_vulnerability(root_path, vulnerability_info, args.username, user_id)
+
+    output_file = os.path.join(output_dir, "repair_result.json")
+    with open(output_file, "w", encoding="utf-8") as f:
+        json.dump(result, f, ensure_ascii=False, indent=2)
+    print(output_file)
+
+
+if __name__ == "__main__":
+    main()

package/comate-engine/assets/skills/code-security-comate/scripts/report_chat.py

@@ -0,0 +1,145 @@
+#!/usr/bin/env python3
+"""
+对话结果上报工具 - 扫描成功后向服务端上报本次对话信息。
+
+用法:
+    python3 report_chat.py --username <用户名> --chat-id <对话ID> 
+    --scan-result <扫描结果JSON文件> [--git-url <仓库URL>] [--git-branch <分支>] 
+    [--ide <IDE名称>] [--query <用户输入>]
+
+接口: POST /api/v1/chats
+"""
+
+import argparse
+import json
+import logging
+import os
+import sys
+
+import http_client  # noqa: F401 (triggers shared logging config)
+import utils
+
+logger = logging.getLogger("report")
+
+
+def build_vuls_from_scan_result(scan_result, root_path=""):
+    # type: (dict, str) -> list
+    """从扫描结果中提取漏洞信息，构造 vuls 数组。"""
+    vuls = []
+    data = scan_result.get("data", {})
+    # 兼容两种结构: data.sarif.runs 或 data.runs
+    runs = data.get("sarif", {}).get("runs", []) or data.get("runs", [])
+    if not runs:
+        return vuls
+
+    # 按文件聚合漏洞
+    file_vuls = {}  # type: dict
+    for run in runs:
+        # 兼容 "results"(标准SARIF) 和 "result"(服务端实际返回)
+        results = run.get("results", []) or run.get("result", [])
+        for result in results:
+            locations = result.get("locations", [])
+            vul_hash = result.get("properties", {}).get("hash", "")
+            if not vul_hash:
+                continue
+            for loc in locations:
+                file_path = loc.get("physicalLocation", {}).get("artifactLocation", {}).get("uri", "")
+                if not file_path:
+                    continue
+                if file_path not in file_vuls:
+                    file_vuls[file_path] = {
+                        "fileName": file_path,
+                        "fileHash": "",
+                        "vuls": [],
+                    }
+                file_vuls[file_path]["vuls"].append({"vulHash": vul_hash})
+
+    # 填充 fileHash
+    bundle_files = scan_result.get("data", {}).get("bundleFiles", {})
+    for file_path, file_info in file_vuls.items():
+        fh = bundle_files.get(file_path, "")
+        if not fh and root_path:
+            fh = utils.calc_sha256_safe(os.path.join(root_path, file_path))
+        file_info["fileHash"] = fh
+
+    vuls = list(file_vuls.values())
+    return vuls
+
+
+def report_chat(chat_id, scan_result, git_url, git_branch, ide, query, root_path="", username="", user_id=""):
+    # type: (str, dict, str, str, str, str, str, str, str) -> dict
+    """上报对话结果到服务端。"""
+    # git_url/git_branch 为空时，自动从 root_path 获取
+    if root_path and (not git_url or not git_branch):
+        git_info = utils.get_git_info(root_path)
+        if not git_url:
+            git_url = git_info["gitURL"]
+        if not git_branch:
+            git_branch = git_info["gitBranch"]
+
+    vuls = build_vuls_from_scan_result(scan_result, root_path)
+
+    body = {
+        "type": 1,
+        "gitInfo": {
+            "url": git_url,
+            "branch": git_branch,
+        },
+        "ide": ide,
+        "query": query,
+        "status": 0,
+        "errMessage": "",
+        "vuls": vuls,
+    }
+
+    headers = utils.build_headers(username, user_id, chat_id)
+    logger.info("report_chat: chat_id=%s, vuls_count=%d, git_url=%s, git_branch=%s",
+                chat_id, len(vuls), git_url, git_branch)
+    return http_client.post("{}/api/v1/chats".format(utils.HOST), headers=headers, json_body=body)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="对话结果上报工具")
+    parser.add_argument("--username", required=True, help="Comate 用户名")
+    parser.add_argument("--chat-id", required=True, help="对话唯一标识 (COMATE_SESSION_ID)")
+    parser.add_argument("--scan-result", required=True, help="扫描结果 JSON 文件路径")
+    parser.add_argument("--root-path", default="", help="项目根目录，用于计算文件哈希")
+    parser.add_argument("--git-url", default="", help="Git 仓库 URL")
+    parser.add_argument("--git-branch", default="", help="Git 分支")
+    parser.add_argument("--ide", default="", help="IDE 名称")
+    parser.add_argument("--query", default="", help="用户输入的查询文本")
+    args = parser.parse_args()
+
+    user_id = utils.make_user_id(args.username)
+
+    logger.info("report_chat start: chat_id=%s, username=%s, scan_result=%s",
+                args.chat_id, args.username, args.scan_result)
+
+    # 读取扫描结果
+    scan_result = {}
+    try:
+        with open(args.scan_result, "r", encoding="utf-8") as f:
+            scan_result = json.load(f)
+    except Exception as e:
+        print("警告: 读取扫描结果失败 {}: {}".format(args.scan_result, e), file=sys.stderr)
+
+    try:
+        result = report_chat(
+            chat_id=args.chat_id,
+            scan_result=scan_result,
+            git_url=args.git_url,
+            git_branch=args.git_branch,
+            ide=args.ide,
+            query=args.query,
+            root_path=args.root_path,
+            username=args.username,
+            user_id=user_id,
+        )
+        print(json.dumps(result, ensure_ascii=False))
+    except Exception as e:
+        print("上报失败: {}".format(e), file=sys.stderr)
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()