diff --git a/kb_workflow/auth/zhihu_login.py b/kb_workflow/auth/zhihu_login.py index 1b9cecf..45d2e52 100644 --- a/kb_workflow/auth/zhihu_login.py +++ b/kb_workflow/auth/zhihu_login.py @@ -34,6 +34,21 @@ class ZhihuLoginResult(BaseModel): 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. @@ -54,6 +69,7 @@ class ZhihuLogin: "canvas[class*='qrcode']", "canvas" ] + VALIDATION_URL = "https://www.zhihu.com/api/v4/me" def __init__( self, @@ -218,3 +234,43 @@ class ZhihuLogin: 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))