1  # Copyright (c) 2024-2026 Tencent Zhuque Lab. All rights reserved.
  2  #
  3  # Licensed under the Apache License, Version 2.0 (the "License");
  4  # you may not use this file except in compliance with the License.
  5  # You may obtain a copy of the License at
  6  #
  7  #     http://www.apache.org/licenses/LICENSE-2.0
  8  #
  9  # Unless required by applicable law or agreed to in writing, software
 10  # distributed under the License is distributed on an "AS IS" BASIS,
 11  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 12  # See the License for the specific language governing permissions and
 13  # limitations under the License.
 14  #
 15  # Requirement: Any integration or derivative work must explicitly attribute
 16  # Tencent Zhuque Lab (https://github.com/Tencent/AI-Infra-Guard) in its
 17  # documentation or user interface, as detailed in the NOTICE file.
 18  
 19  import re
 20  
 21  
 22  class VulnerabilityExtractor:
 23      """漏洞信息提取器"""
 24  
 25      def __init__(self):
 26          self.pattern = re.compile(
 27              r"<vuln>\s*(.*?)\s*</vuln>",
 28              re.DOTALL,  # 使 . 匹配包括换行符在内的所有字符
 29          )
 30  
 31      def extract_vulnerabilities(self, text: str) -> list[dict[str, str]]:
 32          """
 33          从文本中提取所有漏洞信息
 34  
 35          Args:
 36              text: 包含漏洞信息的文本
 37  
 38          Returns:
 39              漏洞信息字典列表
 40          """
 41          vulnerabilities = []
 42  
 43          # 查找所有vuln块
 44          vuln_blocks = self.pattern.findall(text)
 45  
 46          for i, block in enumerate(vuln_blocks, 1):
 47              try:
 48                  vuln_info = self._parse_vuln_block(block, i)
 49                  if vuln_info:
 50                      vulnerabilities.append(vuln_info)
 51              except Exception as e:
 52                  print(f"解析第 {i} 个漏洞块时出错: {e}")
 53                  continue
 54  
 55          return vulnerabilities
 56  
 57      def _parse_vuln_block(self, block: str, index: int) -> dict[str, str] | None:
 58          """解析单个vuln块"""
 59  
 60          # 提取各个字段
 61          title = self._extract_tag_content(block, "title")
 62          desc = self._extract_tag_content(block, "desc")
 63          risk_type = self._extract_tag_content(block, "risk_type")
 64          level = self._extract_tag_content(block, "level")
 65          suggestion = self._extract_tag_content(block, "suggestion")
 66  
 67          # 验证必要字段
 68          if not all([title, desc, risk_type]):
 69              print(f"第 {index} 个漏洞块缺少必要字段，跳过")
 70              return None
 71  
 72          return {
 73              "title": title.strip(),
 74              "description": desc.strip(),
 75              "risk_type": risk_type.strip(),
 76              "level": level.strip(),
 77              "suggestion": suggestion.strip(),
 78          }
 79  
 80      def _extract_tag_content(self, text: str, tag: str) -> str | None:
 81          """提取指定标签的内容"""
 82          pattern = re.compile(rf"<{tag}>\s*(.*?)\s*</{tag}>", re.DOTALL)
 83          match = pattern.search(text)
 84          return match.group(1) if match else None
 85  
 86  
 87  if __name__ == "__main__":
 88      extractor = VulnerabilityExtractor()
 89      text = """
 90  2
 91  2025-11-12 18:52:19.481 | INFO     | agent.base_agent:run:132 - Agent execution completed
 92  2025-11-12 18:52:19.481 | INFO     | __main__:main:36 - Agent completed successfully:
 93  
 94   <vuln>
 95    <title>命令注入漏洞 - executeCohoCommand函数</title>
 96    <desc>
 97    ## 漏洞详情
 98    **文件位置**: src/index.ts 第168-193行  
 99    **漏洞类型**: Command Injection  
100    **风险等级**: High
101    
102    ### 技术分析
103    在`executeCohoCommand`函数中，用户提供的`command`参数直接拼接到shell命令字符串中，没有进行任何转义或验证：
104    
105    ```typescript
106    const { stdout, stderr } = await exec(`coho ${command} --admintoken ${config.adminToken} `, {
107        timeout: 120000
108    });
109    ```
110    
111    ### 攻击路径
112    攻击者可通过以下工具注入恶意命令：
113    - `query_collection`: 控制collection、query等参数
114    - `deploy_code`: 控制file.path等参数
115    - 所有使用`executeCohoCommand`函数的工具
116    
117    ### 影响评估  
118    - 理论上可执行任意系统命令
119    - 完全控制Docker容器环境
120    - 可能访问敏感数据和凭据
121    </desc>
122    <risk_type>Command Injection</risk_type>
123    <level>High</level>
124    <suggestion>
125    ## 修复建议
126    1. **参数转义**: 使用`execFile`替代`exec`，传递参数数组而非字符串
127    2. **输入验证**: 加强Zod验证规则，禁止特殊字符
128    3. **最小权限原则**: 限制Docker容器权限
129    4. **使用子进程API**: 
130    ```typescript
131    import { execFile } from 'child_process';
132    
133    async function executeCohoCommand(command: string): Promise<string> {
134        const args = [...command.split(' '), '--admintoken', config.adminToken];
135        const { stdout, stderr } = await execFile('coho', args, {
136            timeout: 120000
137        });
138    ```
139    </suggestion>
140  </vuln>
141  
142  <vuln>
143    <title>认证令牌泄露风险</title>
144    <desc>
145    ## 漏洞详情
146    **文件位置**: src/index.ts 第172行  
147    **漏洞类型**: Credential Theft  
148    **风险等级**: Medium
149    
150    ### 技术分析
151    虽然代码中有令牌清理机制，但`adminToken`作为命令行参数传递，可能被系统进程监控工具捕获。
152    
153    ### 攻击路径
154    - 通过系统进程监控工具可获取命令行参数
155    - 不完全的令牌清理机制
156    
157    ### 影响评估  
158    - 管理员令牌可能被泄露
159    - 攻击者可获得项目完全控制权
160    </desc>
161    <risk_type>Credential Theft</risk_type>
162    <level>Medium</level>
163    <suggestion>
164    ## 修复建议
165    1. **使用环境变量**: 通过环境变量传递敏感凭据
166    2. **安全存储机制**: 使用加密存储或密钥管理系统
167    </suggestion>
168  </vuln>
169  
170  ## 漏洞复现验证结果
171  
172  经过对Codehooks.io MCP服务器的代码审计和漏洞复现，确认以下关键发现：
173  
174  ### 已验证的漏洞
175  1. **命令注入漏洞** - 高危风险
176     - 确认位置：`src/index.ts` 第168-193行
177     - 确认攻击向量：通过用户输入直接拼接到shell命令
178     - 确认可利用性：理论上存在，但当前环境缺乏coho CLI工具
179  
180  ### 环境限制评估
181  - **当前环境**: 缺乏coho CLI工具，降低了直接攻击可行性
182  - **生产环境风险**: 如果部署在包含coho CLI的环境中，此漏洞可被完全利用
183  
184  ### 风险评估校准
185  **最终风险等级**: High → Medium  
186  **理由**: 
187  - 漏洞代码确实存在命令注入风险
188  - 当前环境缺乏必要的CLI工具，无法直接复现命令执行
189  - **实际影响**: 在当前测试环境中无法实现完整的攻击链
190  
191  ### 建议行动
192  1. **立即修复**: 重写`executeCohoCommand`函数，使用安全的子进程API
193  2. **环境加固**: 确保生产环境的安全配置
194  3. **代码审查**: 对所有命令执行相关的代码进行全面安全检查
195  
196  Process finished with exit code 0
197  
198  """
199      vulnerabilities = extractor.extract_vulnerabilities(text)
200      print(vulnerabilities)