277 lines
8.7 KiB
Python
277 lines
8.7 KiB
Python
"""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))
|