自动同步: 2026-03-18 02:10

skills/kb_workflow/auth/__init__.py

@@ -0,0 +1,8 @@
+"""Auth module for KB Workflow
+
+Provides authentication and cookie management functionality.
+"""
+
+from .zhihu_login import ZhihuLogin, ZhihuLoginResult
+
+__all__ = ["ZhihuLogin", "ZhihuLoginResult"]

skills/kb_workflow/auth/cookie_manager.py

@@ -0,0 +1,120 @@
+"""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 ""

skills/kb_workflow/auth/zhihu_login.py

@@ -0,0 +1,276 @@
+"""Zhihu Login via QR Code
+
+Implements QR code login flow for Zhihu:
+1. Navigate to login page
+2. Capture QR code from canvas element
+3. Poll for login cookie (z_c0)
+4. Extract cookie string on success
+
+Reference: MediaCrawler implementation pattern
+"""
+
+from datetime import datetime
+from pathlib import Path
+from typing import Optional, List
+
+from pydantic import BaseModel, Field
+from tenacity import Retrying, stop_after_delay, wait_fixed, RetryError
+
+from kb_workflow.metadata.zhihu_extractor import PlaywrightBrowser
+
+
+class ZhihuLoginResult(BaseModel):
+    """Result of QR code login attempt.
+
+    Attributes:
+        success: Whether login was successful
+        cookie_str: Cookie string if login succeeded, empty otherwise
+        qrcode_path: Path to saved QR code image
+        error_message: Error message if login failed
+    """
+    success: bool = Field(..., description="Whether login was successful")
+    cookie_str: str = Field(default="", description="Cookie string")
+    qrcode_path: Optional[Path] = Field(None, description="Path to saved QR code image")
+    error_message: Optional[str] = Field(None, description="Error message if failed")
+
+
+class CookieValidationResult(BaseModel):
+    """Cookie validation result.
+
+    Attributes:
+        valid: Whether the cookie is valid
+        username: User's display name if valid
+        user_id: User's ID if valid
+        reason: Reason for invalid cookie
+    """
+    valid: bool = Field(..., description="Whether cookie is valid")
+    username: Optional[str] = Field(None, description="User display name")
+    user_id: Optional[str] = Field(None, description="User ID")
+    reason: Optional[str] = Field(None, description="Reason if invalid")
+
+
+class ZhihuLogin:
+    """Zhihu QR Code Login Handler.
+
+    Handles the QR code login flow for Zhihu:
+    - Navigate to login page
+    - Capture QR code screenshot
+    - Poll for login cookie
+    - Extract cookie string
+
+    Attributes:
+        LOGIN_URL: Zhihu login page URL
+        CANVAS_SELECTORS: List of selectors to find QR code canvas
+    """
+
+    LOGIN_URL = "https://www.zhihu.com/signin"
+    CANVAS_SELECTORS = [
+        "canvas.SignIn-qrcode",
+        "canvas[class*='qrcode']",
+        "canvas"
+    ]
+    VALIDATION_URL = "https://www.zhihu.com/api/v4/me"
+
+    def __init__(
+        self,
+        intermediate_path: str,
+        headless: bool = True,
+        timeout: int = 60
+    ):
+        """Initialize ZhihuLogin.
+
+        Args:
+            intermediate_path: Path to intermediate directory for QR code storage
+            headless: Whether to run browser in headless mode
+            timeout: Login timeout in seconds
+        """
+        self.intermediate_path = Path(intermediate_path)
+        self.headless = headless
+        self.timeout = timeout
+        self.browser: Optional[PlaywrightBrowser] = None
+
+    def _ensure_intermediate_dir(self) -> None:
+        """Ensure intermediate directory exists."""
+        self.intermediate_path.mkdir(parents=True, exist_ok=True)
+
+    def _find_qrcode_canvas(self):
+        """Find QR code canvas element.
+
+        Tries multiple selectors to find the QR code canvas.
+
+        Returns:
+            Canvas element handle or None
+        """
+        for selector in self.CANVAS_SELECTORS:
+            try:
+                canvas = self.browser.wait_for_selector(selector, timeout=5)
+                if canvas:
+                    return canvas
+            except Exception:
+                continue
+        return None
+
+    def _save_qrcode_image(self, screenshot: bytes) -> Path:
+        """Save QR code screenshot to intermediate directory.
+
+        Args:
+            screenshot: PNG bytes from canvas screenshot
+
+        Returns:
+            Path to saved QR code image
+        """
+        self._ensure_intermediate_dir()
+
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        filename = f"zhihu_qrcode_{timestamp}.png"
+        filepath = self.intermediate_path / filename
+
+        filepath.write_bytes(screenshot)
+        return filepath
+
+    def check_login_cookie(self) -> bool:
+        """Check if login cookie (z_c0) exists.
+
+        Returns:
+            True if z_c0 cookie exists, False otherwise
+        """
+        cookies = self.browser.context.cookies()
+        for cookie in cookies:
+            if cookie.get("name") == "z_c0":
+                return True
+        return False
+
+    def get_cookies_str(self) -> str:
+        """Get cookie string from browser context.
+
+        Returns:
+            Cookie string in format "name1=value1; name2=value2"
+        """
+        cookies = self.browser.context.cookies()
+        cookie_pairs = [
+            f"{cookie['name']}={cookie['value']}"
+            for cookie in cookies
+        ]
+        return "; ".join(cookie_pairs)
+
+    def _wait_for_login(self) -> None:
+        """Wait for user to scan QR code and login.
+
+        Uses tenacity Retrying to poll for login cookie with configurable timeout.
+
+        Raises:
+            RetryError: If login timeout exceeded
+        """
+        for attempt in Retrying(
+            stop=stop_after_delay(self.timeout),
+            wait=wait_fixed(1),
+            reraise=True
+        ):
+            with attempt:
+                if not self.check_login_cookie():
+                    raise Exception("Not logged in yet")
+
+    def login_by_qrcode(self) -> ZhihuLoginResult:
+        """Execute QR code login flow.
+
+        Main entry point for QR code login:
+        1. Start browser
+        2. Navigate to login page
+        3. Find and capture QR code
+        4. Wait for user to scan
+        5. Extract cookies on success
+
+        Returns:
+            ZhihuLoginResult with success status and cookies
+        """
+        qrcode_path = None
+
+        try:
+            # Start browser
+            self.browser = PlaywrightBrowser(headless=self.headless, timeout=30)
+            self.browser.start()
+
+            # Navigate to login page
+            self.browser.goto(self.LOGIN_URL, wait_until="networkidle")
+
+            # Find QR code canvas
+            canvas = self._find_qrcode_canvas()
+            if not canvas:
+                return ZhihuLoginResult(
+                    success=False,
+                    error_message="Could not find QR code canvas on login page"
+                )
+
+            # Capture screenshot
+            screenshot = canvas.screenshot()
+            qrcode_path = self._save_qrcode_image(screenshot)
+
+            # Wait for login
+            try:
+                self._wait_for_login()
+            except RetryError:
+                return ZhihuLoginResult(
+                    success=False,
+                    qrcode_path=qrcode_path,
+                    error_message=f"Login timeout after {self.timeout} seconds"
+                )
+
+            # Extract cookies
+            cookie_str = self.get_cookies_str()
+
+            return ZhihuLoginResult(
+                success=True,
+                cookie_str=cookie_str,
+                qrcode_path=qrcode_path
+            )
+
+        except Exception as e:
+            return ZhihuLoginResult(
+                success=False,
+                qrcode_path=qrcode_path,
+                error_message=str(e)
+            )
+
+        finally:
+            if self.browser:
+                self.browser.close()
+
+    async def validate_cookie(self, cookie_str: str = None) -> CookieValidationResult:
+        """Validate cookie validity.
+
+        Checks cookie by requesting Zhihu user info API.
+
+        Args:
+            cookie_str: Cookie string to validate. If None, uses already loaded cookie.
+
+        Returns:
+            CookieValidationResult with validation status and user info
+        """
+        # If cookie_str provided, need to restart browser with new cookie
+        if cookie_str:
+            if self.browser.browser:
+                self.browser.close()
+            self.browser.start(cookie=cookie_str)
+        elif not self.browser.page:
+            return CookieValidationResult(valid=False, reason="No cookie provided")
+
+        try:
+            # Use Playwright's request API
+            response = await self.browser.page.request.get(self.VALIDATION_URL)
+
+            if response.status == 200:
+                user_info = await response.json()
+                return CookieValidationResult(
+                    valid=True,
+                    username=user_info.get("name", ""),
+                    user_id=str(user_info.get("id", ""))
+                )
+            elif response.status == 401:
+                return CookieValidationResult(valid=False, reason="Cookie expired")
+            else:
+                return CookieValidationResult(
+                    valid=False,
+                    reason=f"Validation failed with status {response.status}"
+                )
+        except Exception as e:
+            return CookieValidationResult(valid=False, reason=str(e))