security_utils.py 17 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589
  1. #!/usr/bin/env python3
  2. """
  3. 安全工具函数库
  4. 用于webnovel-writer系统的通用安全函数
  5. 创建时间: 2026-01-02
  6. 创建原因: 安全审计发现路径遍历和命令注入漏洞
  7. 修复方案: 集中管理所有安全相关的输入清理函数
  8. """
  9. import json
  10. import os
  11. import re
  12. import sys
  13. import tempfile
  14. from pathlib import Path
  15. from runtime_compat import enable_windows_utf8_stdio
  16. from typing import Any, Dict, Optional, Union
  17. # 尝试导入 filelock(可选依赖)
  18. try:
  19. from filelock import FileLock
  20. HAS_FILELOCK = True
  21. except ImportError:
  22. HAS_FILELOCK = False
  23. def sanitize_filename(name: str, max_length: int = 100) -> str:
  24. """
  25. 清理文件名,防止路径遍历攻击 (CWE-22)
  26. 安全关键函数 - 修复extract_entities.py路径遍历漏洞
  27. Args:
  28. name: 原始文件名(可能包含路径遍历字符)
  29. max_length: 文件名最大长度(默认100字符)
  30. Returns:
  31. 安全的文件名(仅包含基本文件名,移除所有路径信息)
  32. 示例:
  33. >>> sanitize_filename("../../../etc/passwd")
  34. 'passwd'
  35. >>> sanitize_filename("C:\\Windows\\System32")
  36. 'System32'
  37. >>> sanitize_filename("正常角色名")
  38. '正常角色名'
  39. 安全验证:
  40. - ✅ 防止目录遍历(../、..\\)
  41. - ✅ 防止绝对路径(/、C:\\)
  42. - ✅ 移除特殊字符
  43. - ✅ 长度限制
  44. """
  45. # Step 1: 仅保留基础文件名(移除所有路径)
  46. safe_name = os.path.basename(name)
  47. # Step 2: 移除路径分隔符(双重保险)
  48. safe_name = safe_name.replace('/', '_').replace('\\', '_')
  49. # Step 3: 只保留安全字符
  50. # 允许:中文(\u4e00-\u9fff)、字母(a-zA-Z)、数字(0-9)、下划线(_)、连字符(-)
  51. safe_name = re.sub(r'[^\w\u4e00-\u9fff-]', '_', safe_name)
  52. # Step 4: 移除连续的下划线(美化)
  53. safe_name = re.sub(r'_+', '_', safe_name)
  54. # Step 5: 长度限制
  55. if len(safe_name) > max_length:
  56. safe_name = safe_name[:max_length]
  57. # Step 6: 移除首尾下划线
  58. safe_name = safe_name.strip('_')
  59. # Step 7: 确保非空(防御性编程)
  60. if not safe_name:
  61. safe_name = "unnamed_entity"
  62. return safe_name
  63. def sanitize_commit_message(message: str, max_length: int = 200) -> str:
  64. """
  65. 清理Git提交消息,防止命令注入 (CWE-77)
  66. 安全关键函数 - 修复backup_manager.py命令注入漏洞
  67. Args:
  68. message: 原始提交消息(可能包含Git标志)
  69. max_length: 消息最大长度(默认200字符)
  70. Returns:
  71. 安全的提交消息(移除Git特殊标志和危险字符)
  72. 示例:
  73. >>> sanitize_commit_message("Test\\n--author='Attacker'")
  74. 'Test author Attacker'
  75. >>> sanitize_commit_message("--amend Chapter 1")
  76. 'amend Chapter 1'
  77. 安全验证:
  78. - ✅ 防止多行注入(换行符)
  79. - ✅ 防止Git标志注入(--xxx)
  80. - ✅ 防止参数分隔符混淆(引号)
  81. - ✅ 防止单字母标志(-x)
  82. """
  83. # Step 1: 移除换行符(防止多行参数注入)
  84. safe_msg = message.replace('\n', ' ').replace('\r', ' ')
  85. # Step 2: 移除Git特殊标志(--开头的参数)
  86. safe_msg = re.sub(r'--[\w-]+', '', safe_msg)
  87. # Step 3: 移除引号(防止参数分隔符混淆)
  88. safe_msg = safe_msg.replace("'", "").replace('"', '')
  89. # Step 4: 移除前导的-(防止单字母标志如-m)
  90. safe_msg = safe_msg.lstrip('-')
  91. # Step 5: 移除连续空格(美化)
  92. safe_msg = re.sub(r'\s+', ' ', safe_msg)
  93. # Step 6: 长度限制
  94. if len(safe_msg) > max_length:
  95. safe_msg = safe_msg[:max_length]
  96. # Step 7: 移除首尾空格
  97. safe_msg = safe_msg.strip()
  98. # Step 8: 确保非空
  99. if not safe_msg:
  100. safe_msg = "Untitled commit"
  101. return safe_msg
  102. def create_secure_directory(path: str, mode: int = 0o700) -> Path:
  103. """
  104. 创建安全目录(仅所有者可访问)
  105. 安全关键函数 - 修复文件权限配置缺失漏洞
  106. Args:
  107. path: 目录路径
  108. mode: 权限模式(默认0o700,仅所有者可读写执行)
  109. Returns:
  110. Path对象
  111. 示例:
  112. >>> create_secure_directory('.webnovel')
  113. PosixPath('.webnovel') # drwx------ (700)
  114. 安全验证:
  115. - ✅ 仅所有者可访问(0o700)
  116. - ✅ 防止同组用户读取
  117. - ✅ 跨平台兼容(Windows/Linux/macOS)
  118. """
  119. path_obj = Path(path)
  120. # Windows 上传入 mode 会触发不可预期的 ACL 行为(实测会导致目录创建后立刻无法访问)。
  121. # 因此在 Windows 下不传 mode,保持默认继承权限;在类 Unix 系统才使用 mode。
  122. if os.name == 'nt':
  123. os.makedirs(path, exist_ok=True)
  124. else:
  125. os.makedirs(path, mode=mode, exist_ok=True)
  126. # 双重保险:显式设置权限(某些系统可能忽略makedirs的mode参数)
  127. if os.name != 'nt': # Unix系统(Linux/macOS)
  128. os.chmod(path, mode)
  129. return path_obj
  130. def create_secure_file(file_path: str, content: str, mode: int = 0o600) -> None:
  131. """
  132. 创建安全文件(仅所有者可读写)
  133. Args:
  134. file_path: 文件路径
  135. content: 文件内容
  136. mode: 权限模式(默认0o600,仅所有者可读写)
  137. 安全验证:
  138. - ✅ 仅所有者可读写(0o600)
  139. - ✅ 防止其他用户访问
  140. """
  141. # 创建文件
  142. with open(file_path, 'w', encoding='utf-8') as f:
  143. f.write(content)
  144. # 设置权限(仅Unix系统)
  145. if os.name != 'nt':
  146. os.chmod(file_path, mode)
  147. def validate_integer_input(value: str, field_name: str) -> int:
  148. """
  149. 验证并转换整数输入(严格模式)
  150. 安全关键函数 - 修复update_state.py弱验证漏洞
  151. Args:
  152. value: 输入值(字符串)
  153. field_name: 字段名称(用于错误消息)
  154. Returns:
  155. 转换后的整数
  156. Raises:
  157. ValueError: 输入不是有效整数
  158. 示例:
  159. >>> validate_integer_input("123", "chapter_num")
  160. 123
  161. >>> validate_integer_input("abc", "level")
  162. ValueError: ❌ 错误:level 必须是整数,收到: abc
  163. """
  164. try:
  165. return int(value)
  166. except ValueError:
  167. print(f"❌ 错误:{field_name} 必须是整数,收到: {value}", file=sys.stderr)
  168. raise ValueError(f"Invalid integer input for {field_name}: {value}")
  169. # ============================================================================
  170. # Git 环境检测(优雅降级支持)
  171. # ============================================================================
  172. # 缓存 Git 可用性检测结果
  173. _git_available: Optional[bool] = None
  174. def is_git_available() -> bool:
  175. """
  176. 检测 Git 是否可用
  177. Returns:
  178. bool: Git 是否可用
  179. 说明:
  180. - 检测结果会被缓存,避免重复检测
  181. - 用于支持在无 Git 环境下优雅降级
  182. """
  183. global _git_available
  184. if _git_available is not None:
  185. return _git_available
  186. import subprocess
  187. try:
  188. result = subprocess.run(
  189. ["git", "--version"],
  190. capture_output=True,
  191. text=True,
  192. timeout=5
  193. )
  194. _git_available = result.returncode == 0
  195. except (FileNotFoundError, subprocess.TimeoutExpired, OSError):
  196. _git_available = False
  197. return _git_available
  198. def is_git_repo(path: Union[str, Path]) -> bool:
  199. """
  200. 检测指定目录是否是 Git 仓库
  201. Args:
  202. path: 目录路径
  203. Returns:
  204. bool: 是否是 Git 仓库
  205. """
  206. if not is_git_available():
  207. return False
  208. path = Path(path)
  209. git_dir = path / ".git"
  210. return git_dir.exists() and git_dir.is_dir()
  211. def git_graceful_operation(
  212. args: list,
  213. cwd: Union[str, Path],
  214. *,
  215. fallback_msg: str = "Git 不可用,跳过版本控制操作"
  216. ) -> tuple:
  217. """
  218. 优雅执行 Git 操作(Git 不可用时静默降级)
  219. Args:
  220. args: Git 命令参数(不含 'git')
  221. cwd: 工作目录
  222. fallback_msg: 降级时的提示消息
  223. Returns:
  224. (success: bool, output: str, was_skipped: bool)
  225. - success: 操作是否成功
  226. - output: 输出内容
  227. - was_skipped: 是否因 Git 不可用而跳过
  228. 示例:
  229. >>> success, output, skipped = git_graceful_operation(
  230. ... ["add", "."], cwd="/path/to/project"
  231. ... )
  232. >>> if skipped:
  233. ... print("Git not available, using fallback")
  234. """
  235. if not is_git_available():
  236. print(f"⚠️ {fallback_msg}", file=sys.stderr)
  237. return False, "", True
  238. import subprocess
  239. try:
  240. result = subprocess.run(
  241. ["git"] + args,
  242. cwd=cwd,
  243. capture_output=True,
  244. text=True,
  245. encoding='utf-8',
  246. timeout=60
  247. )
  248. return result.returncode == 0, result.stdout, False
  249. except subprocess.TimeoutExpired:
  250. print(f"⚠️ Git 操作超时: git {' '.join(args)}", file=sys.stderr)
  251. return False, "", False
  252. except OSError as e:
  253. print(f"⚠️ Git 操作失败: {e}", file=sys.stderr)
  254. return False, "", False
  255. # ============================================================================
  256. # 原子化文件写入(防止并发冲突和数据损坏)
  257. # ============================================================================
  258. class AtomicWriteError(Exception):
  259. """原子写入失败异常"""
  260. pass
  261. def atomic_write_json(
  262. file_path: Union[str, Path],
  263. data: Dict[str, Any],
  264. *,
  265. use_lock: bool = True,
  266. backup: bool = True,
  267. indent: int = 2
  268. ) -> None:
  269. """
  270. 原子化写入 JSON 文件,防止并发冲突和数据损坏 (CWE-362, CWE-367)
  271. 安全关键函数 - 修复 state.json 并发写入风险
  272. 实现策略:
  273. 1. 写入临时文件(同目录,确保同文件系统)
  274. 2. 可选:使用 filelock 获取排他锁
  275. 3. 可选:备份原文件
  276. 4. 原子重命名(os.replace 在 POSIX 上是原子的)
  277. Args:
  278. file_path: 目标文件路径
  279. data: 要写入的字典数据
  280. use_lock: 是否使用文件锁(需要 filelock 库)
  281. backup: 是否在写入前备份原文件
  282. indent: JSON 缩进(默认 2)
  283. Raises:
  284. AtomicWriteError: 写入失败时抛出
  285. 示例:
  286. >>> atomic_write_json('.webnovel/state.json', {'progress': {'chapter': 10}})
  287. 安全验证:
  288. - ✅ 防止写入中断导致的数据损坏(先写临时文件)
  289. - ✅ 防止并发写入冲突(filelock)
  290. - ✅ 支持回滚(备份机制)
  291. - ✅ 跨平台兼容
  292. """
  293. file_path = Path(file_path)
  294. parent_dir = file_path.parent
  295. parent_dir.mkdir(parents=True, exist_ok=True)
  296. # 准备 JSON 内容
  297. try:
  298. json_content = json.dumps(data, ensure_ascii=False, indent=indent)
  299. except (TypeError, ValueError) as e:
  300. raise AtomicWriteError(f"JSON 序列化失败: {e}")
  301. # 锁文件路径
  302. lock_path = file_path.with_suffix(file_path.suffix + '.lock')
  303. backup_path = file_path.with_suffix(file_path.suffix + '.bak')
  304. # 创建临时文件(同目录确保同文件系统,os.replace 才能原子操作)
  305. fd, temp_path = tempfile.mkstemp(
  306. suffix='.tmp',
  307. prefix=file_path.stem + '_',
  308. dir=parent_dir
  309. )
  310. try:
  311. # Step 1: 写入临时文件
  312. with os.fdopen(fd, 'w', encoding='utf-8') as f:
  313. f.write(json_content)
  314. f.flush()
  315. os.fsync(f.fileno()) # 确保写入磁盘
  316. # Step 2: 获取锁(如果可用且启用)
  317. lock = None
  318. if use_lock and HAS_FILELOCK:
  319. lock = FileLock(str(lock_path), timeout=10)
  320. lock.acquire()
  321. try:
  322. # Step 3: 备份原文件(如果存在且启用备份)
  323. if backup and file_path.exists():
  324. try:
  325. import shutil
  326. shutil.copy2(file_path, backup_path)
  327. except OSError:
  328. pass # 备份失败不阻止写入
  329. # Step 4: 原子重命名
  330. os.replace(temp_path, file_path)
  331. temp_path = None # 标记已成功,不需要清理
  332. finally:
  333. if lock is not None:
  334. lock.release()
  335. except Exception as e:
  336. raise AtomicWriteError(f"原子写入失败: {e}")
  337. finally:
  338. # 清理:删除临时文件(如果仍存在说明写入失败)
  339. if temp_path is not None:
  340. try:
  341. os.unlink(temp_path)
  342. except OSError:
  343. pass
  344. def read_json_safe(
  345. file_path: Union[str, Path],
  346. default: Optional[Dict[str, Any]] = None
  347. ) -> Dict[str, Any]:
  348. """
  349. 安全读取 JSON 文件(带默认值和错误处理)
  350. Args:
  351. file_path: 文件路径
  352. default: 文件不存在或解析失败时的默认值
  353. Returns:
  354. 解析后的字典,或默认值
  355. 示例:
  356. >>> state = read_json_safe('.webnovel/state.json', {})
  357. """
  358. file_path = Path(file_path)
  359. if default is None:
  360. default = {}
  361. if not file_path.exists():
  362. return default
  363. try:
  364. with open(file_path, 'r', encoding='utf-8') as f:
  365. return json.load(f)
  366. except (json.JSONDecodeError, OSError) as e:
  367. print(f"⚠️ 读取 JSON 失败 ({file_path}): {e}", file=sys.stderr)
  368. return default
  369. def restore_from_backup(file_path: Union[str, Path]) -> bool:
  370. """
  371. 从备份恢复文件
  372. Args:
  373. file_path: 原文件路径
  374. Returns:
  375. 是否成功恢复
  376. 示例:
  377. >>> restore_from_backup('.webnovel/state.json')
  378. True
  379. """
  380. file_path = Path(file_path)
  381. backup_path = file_path.with_suffix(file_path.suffix + '.bak')
  382. if not backup_path.exists():
  383. print(f"⚠️ 备份文件不存在: {backup_path}", file=sys.stderr)
  384. return False
  385. try:
  386. import shutil
  387. shutil.copy2(backup_path, file_path)
  388. print(f"✅ 已从备份恢复: {file_path}")
  389. return True
  390. except OSError as e:
  391. print(f"❌ 恢复失败: {e}", file=sys.stderr)
  392. return False
  393. # ============================================================================
  394. # 单元测试(内置自检)
  395. # ============================================================================
  396. def _run_self_tests():
  397. """运行内置安全测试"""
  398. print("🔍 运行安全工具函数自检...")
  399. # Test 1: sanitize_filename
  400. assert sanitize_filename("../../../etc/passwd") == "passwd", "路径遍历测试失败"
  401. assert sanitize_filename("C:\\Windows\\System32") == "System32", "Windows路径测试失败"
  402. assert sanitize_filename("正常角色名") == "正常角色名", "中文测试失败"
  403. assert sanitize_filename("/tmp/../../../../../etc/hosts") == "hosts", "复杂路径遍历测试失败"
  404. assert sanitize_filename("test///file...name") == "file_name", "特殊字符测试失败" # . 会被替换
  405. print(" ✅ sanitize_filename: 所有测试通过")
  406. # Test 2: sanitize_commit_message
  407. result = sanitize_commit_message("Test\n--author='Attacker'")
  408. assert "\n" not in result, "换行符未移除"
  409. assert "--author" not in result, "Git标志未移除"
  410. assert "Attacker" in result, "内容被错误移除"
  411. assert sanitize_commit_message("--amend Chapter 1") == "Chapter 1", "Git标志测试失败" # --amend被完全移除
  412. assert "'" not in sanitize_commit_message("Test'message"), "引号测试失败"
  413. assert sanitize_commit_message("-m Test") == "m Test", "单字母标志测试失败" # -m被移除后是"m Test"
  414. print(" ✅ sanitize_commit_message: 所有测试通过")
  415. # Test 3: validate_integer_input
  416. assert validate_integer_input("123", "test") == 123, "整数验证测试失败"
  417. try:
  418. validate_integer_input("abc", "test")
  419. assert False, "应该抛出ValueError"
  420. except ValueError:
  421. pass
  422. print(" ✅ validate_integer_input: 所有测试通过")
  423. # Test 4: atomic_write_json
  424. import tempfile as tf
  425. test_dir = Path(tf.mkdtemp())
  426. test_file = test_dir / "test_state.json"
  427. # 写入测试
  428. test_data = {"chapter": 10, "中文键": "中文值"}
  429. atomic_write_json(test_file, test_data, use_lock=False, backup=False)
  430. assert test_file.exists(), "原子写入未创建文件"
  431. # 读取验证
  432. with open(test_file, 'r', encoding='utf-8') as f:
  433. loaded = json.load(f)
  434. assert loaded == test_data, "原子写入数据不匹配"
  435. # 备份测试
  436. atomic_write_json(test_file, {"updated": True}, use_lock=False, backup=True)
  437. backup_file = test_file.with_suffix('.json.bak')
  438. assert backup_file.exists(), "备份未创建"
  439. # 恢复测试
  440. restore_from_backup(test_file)
  441. with open(test_file, 'r', encoding='utf-8') as f:
  442. restored = json.load(f)
  443. assert restored == test_data, "恢复数据不匹配"
  444. # 清理
  445. import shutil
  446. shutil.rmtree(test_dir)
  447. print(" ✅ atomic_write_json: 所有测试通过")
  448. if HAS_FILELOCK:
  449. print(" ℹ️ filelock 可用,已启用文件锁支持")
  450. else:
  451. print(" ⚠️ filelock 未安装,文件锁功能不可用")
  452. print("\n✅ 所有安全工具函数测试通过!")
  453. if __name__ == "__main__":
  454. # Windows UTF-8 编码修复(必须在打印前执行)
  455. if sys.platform == "win32":
  456. enable_windows_utf8_stdio()
  457. # 运行自检测试
  458. _run_self_tests()