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