starmen1265
00245c58cc
- CookieManager saves cookies to ~/.kb_env file - Updates existing ZHIHU_COOKIE line or appends new - Preserves other configuration lines - Sets file permissions to 600 for security - Loads cookies from env file - Clears get_settings cache after saving Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
121 lines
3.4 KiB
Python
121 lines
3.4 KiB
Python
"""Cookie 管理模块
|
|
|
|
处理 Cookie 的导入、转换和持久化。
|
|
"""
|
|
|
|
import json
|
|
from pathlib import Path
|
|
from typing import List, Dict, Any
|
|
|
|
|
|
def parse_json_cookies(json_path: str) -> str:
|
|
"""从 JSON 文件解析 Cookie 为字符串格式
|
|
|
|
支持 Chrome Cookie Editor 导出的 JSON 格式:
|
|
[{"name": "z_c0", "value": "xxx", "domain": ".zhihu.com"}, ...]
|
|
|
|
Args:
|
|
json_path: JSON 文件路径
|
|
|
|
Returns:
|
|
Cookie 字符串，格式: "name1=value1; name2=value2"
|
|
|
|
Raises:
|
|
ValueError: JSON 格式无效或文件不存在
|
|
"""
|
|
path = Path(json_path)
|
|
if not path.exists():
|
|
raise ValueError(f"Cookie file not found: {json_path}")
|
|
|
|
try:
|
|
with open(path, 'r', encoding='utf-8') as f:
|
|
cookies = json.load(f)
|
|
except json.JSONDecodeError as e:
|
|
raise ValueError(f"Invalid JSON format: {e}")
|
|
|
|
if not isinstance(cookies, list):
|
|
raise ValueError("Expected JSON array")
|
|
|
|
# 过滤知乎域名，转换为 name=value 格式
|
|
zhihu_cookies = []
|
|
for c in cookies:
|
|
if not isinstance(c, dict):
|
|
continue
|
|
domain = c.get('domain', '')
|
|
if 'zhihu.com' in domain:
|
|
name = c.get('name', '')
|
|
value = c.get('value', '')
|
|
if name and value:
|
|
zhihu_cookies.append(f"{name}={value}")
|
|
|
|
return '; '.join(zhihu_cookies)
|
|
|
|
|
|
class CookieManager:
|
|
"""Cookie 持久化管理器
|
|
|
|
管理 Cookie 在 ~/.kb_env 文件中的存储和读取。
|
|
"""
|
|
|
|
ENV_FILE_NAME = ".kb_env"
|
|
COOKIE_KEY = "ZHIHU_COOKIE"
|
|
|
|
def __init__(self, env_path: Path = None):
|
|
"""初始化管理器
|
|
|
|
Args:
|
|
env_path: 环境文件路径，默认 ~/.kb_env
|
|
"""
|
|
self.env_path = env_path or (Path.home() / self.ENV_FILE_NAME)
|
|
|
|
def save_cookie(self, cookie_str: str) -> None:
|
|
"""保存 Cookie 到环境文件
|
|
|
|
格式: ZHIHU_COOKIE="name1=value1; name2=value2"
|
|
|
|
Args:
|
|
cookie_str: Cookie 字符串
|
|
"""
|
|
# 读取现有内容
|
|
lines = []
|
|
if self.env_path.exists():
|
|
lines = self.env_path.read_text(encoding='utf-8').splitlines()
|
|
|
|
# 更新或添加 ZHIHU_COOKIE 行
|
|
cookie_line = f'{self.COOKIE_KEY}="{cookie_str}"'
|
|
found = False
|
|
for i, line in enumerate(lines):
|
|
if line.startswith(f"{self.COOKIE_KEY}="):
|
|
lines[i] = cookie_line
|
|
found = True
|
|
break
|
|
if not found:
|
|
lines.append(cookie_line)
|
|
|
|
# 写回文件，设置权限 600
|
|
self.env_path.write_text('\n'.join(lines) + '\n', encoding='utf-8')
|
|
self.env_path.chmod(0o600)
|
|
|
|
# 清除配置缓存
|
|
from kb_workflow.pipeline.config import get_settings
|
|
get_settings.cache_clear()
|
|
|
|
def load_cookie(self) -> str:
|
|
"""从环境文件加载 Cookie
|
|
|
|
Returns:
|
|
Cookie 字符串，不存在时返回空字符串
|
|
"""
|
|
if not self.env_path.exists():
|
|
return ""
|
|
|
|
content = self.env_path.read_text(encoding='utf-8')
|
|
for line in content.splitlines():
|
|
if line.startswith(f"{self.COOKIE_KEY}="):
|
|
# 提取引号内的值
|
|
value = line[len(f"{self.COOKIE_KEY}="):]
|
|
if value.startswith('"') and value.endswith('"'):
|
|
value = value[1:-1]
|
|
return value
|
|
return ""
|