kb-skills/kb_workflow/auth/zhihu_login.py

"""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 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"
    ]

    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()