diff --git a/kb_workflow/auth/cookie_manager.py b/kb_workflow/auth/cookie_manager.py index 32765cf..b9927bf 100644 --- a/kb_workflow/auth/cookie_manager.py +++ b/kb_workflow/auth/cookie_manager.py @@ -49,3 +49,72 @@ def parse_json_cookies(json_path: str) -> str: 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 ""