diff --git a/kb_workflow/kb_workflow.py b/kb_workflow/kb_workflow.py index c5fda0f..c89108c 100644 --- a/kb_workflow/kb_workflow.py +++ b/kb_workflow/kb_workflow.py @@ -11,6 +11,9 @@ KB Workflow - Knowledge Base Workflow Orchestrator 使用方式: python kb_workflow.py process "https://www.bilibili.com/video/BVxxx" python kb_workflow.py process "https://..." -o /path/to/output + python kb_workflow.py login # QR code login + python kb_workflow.py login --import-cookie cookies.json + python kb_workflow.py login --status """ import argparse @@ -102,6 +105,68 @@ def process_command( } +def login_command( + import_cookie: Optional[str] = None, + status: bool = False +) -> Dict[str, Any]: + """Execute login command + + Args: + import_cookie: Path to JSON cookie file + status: Check current cookie status + + Returns: + Dictionary with login result + """ + try: + from kb_workflow.auth import ZhihuLogin + from kb_workflow.auth.cookie_manager import CookieManager, parse_json_cookies + from kb_workflow.pipeline.config import get_settings + + settings = get_settings() + manager = CookieManager() + + # Check Cookie status + if status: + current_cookie = settings.zhihu_cookie.get_secret_value() + if not current_cookie: + return {"valid": False, "reason": "No cookie found in ~/.kb_env"} + + async def check_status(): + login = ZhihuLogin(intermediate_path=settings.intermediate_path) + try: + login.browser = login.browser or type('obj', (object,), { + 'browser': None, + 'page': None, + 'start': lambda cookie=None: None, + 'close': lambda: None + })() + login.browser.start(cookie=current_cookie) + return await login.validate_cookie(current_cookie) + finally: + if login.browser and hasattr(login.browser, 'close'): + login.browser.close() + + result = asyncio.run(check_status()) + return result.model_dump() + + # Import Cookie + if import_cookie: + cookie_str = parse_json_cookies(import_cookie) + manager.save_cookie(cookie_str) + return {"success": True, "message": f"Cookie imported from {import_cookie}"} + + # QR code login (synchronous method) + login = ZhihuLogin(intermediate_path=settings.intermediate_path) + result = login.login_by_qrcode() + if result.success and result.cookie_str: + manager.save_cookie(result.cookie_str) + return result.model_dump() + + except Exception as e: + return {"success": False, "error": str(e)} + + def main(): """CLI入口 - 用于命令行调用""" parser = argparse.ArgumentParser( @@ -117,6 +182,15 @@ def main(): # 强制重新执行（忽略已完成状态） python kb_workflow.py process "https://..." --no-skip + + # 知乎登录 (二维码扫码) + python kb_workflow.py login + + # 导入 Cookie + python kb_workflow.py login --import-cookie cookies.json + + # 检查 Cookie 状态 + python kb_workflow.py login --status ''' ) @@ -132,6 +206,13 @@ def main(): process_parser.add_argument('--no-skip', action='store_false', dest='skip_completed', help='Re-run all steps') + # Login subcommand + login_parser = subparsers.add_parser('login', help='Zhihu login authentication') + login_parser.add_argument('--import-cookie', '-i', type=str, + help='Import cookie from JSON file') + login_parser.add_argument('--status', action='store_true', + help='Check current cookie status') + args = parser.parse_args() # Handle process subcommand @@ -142,6 +223,13 @@ def main(): skip_completed=args.skip_completed ) print(json.dumps(result, ensure_ascii=False, indent=2)) + # Handle login subcommand + elif args.command == 'login': + result = login_command( + import_cookie=args.import_cookie, + status=args.status + ) + print(json.dumps(result, ensure_ascii=False, indent=2)) else: parser.print_help() diff --git a/tests/test_auth/test_cli_login.py b/tests/test_auth/test_cli_login.py index 584277a..f6828bf 100644 --- a/tests/test_auth/test_cli_login.py +++ b/tests/test_auth/test_cli_login.py @@ -57,13 +57,14 @@ class TestLoginCommand: self, mock_settings, mock_login_result_success ): """Test 1: login_default_mode - Default starts QR code login.""" - with patch('kb_workflow.kb_workflow.get_settings', return_value=mock_settings), \ - patch('kb_workflow.kb_workflow.CookieManager') as MockManager, \ - patch('kb_workflow.kb_workflow.asyncio.run') as mock_run, \ - patch('kb_workflow.kb_workflow.ZhihuLogin') as MockLogin: + with patch('kb_workflow.pipeline.config.get_settings', return_value=mock_settings), \ + patch('kb_workflow.auth.cookie_manager.CookieManager') as MockManager, \ + patch('kb_workflow.auth.ZhihuLogin') as MockLogin: # Setup mock - mock_run.return_value = mock_login_result_success.model_dump() + mock_login_instance = MagicMock() + mock_login_instance.login_by_qrcode.return_value = mock_login_result_success + MockLogin.return_value = mock_login_instance mock_manager = MagicMock() MockManager.return_value = mock_manager @@ -83,9 +84,9 @@ class TestLoginCommand: {"name": "z_c0", "value": "imported_value", "domain": ".zhihu.com"} ])) - with patch('kb_workflow.kb_workflow.get_settings', return_value=mock_settings), \ - patch('kb_workflow.kb_workflow.CookieManager') as MockManager, \ - patch('kb_workflow.kb_workflow.parse_json_cookies') as mock_parse: + with patch('kb_workflow.pipeline.config.get_settings', return_value=mock_settings), \ + patch('kb_workflow.auth.cookie_manager.CookieManager') as MockManager, \ + patch('kb_workflow.auth.cookie_manager.parse_json_cookies') as mock_parse: mock_parse.return_value = "z_c0=imported_value" mock_manager = MagicMock() @@ -101,7 +102,7 @@ class TestLoginCommand: def test_login_status_no_cookie(self, mock_settings): """Test 3a: login_status - --status with no cookie returns invalid.""" - with patch('kb_workflow.kb_workflow.get_settings', return_value=mock_settings): + with patch('kb_workflow.pipeline.config.get_settings', return_value=mock_settings): # Act result = login_command(status=True) @@ -114,14 +115,14 @@ class TestLoginCommand: self, mock_settings_with_cookie, mock_validation_result_valid ): """Test 3b: login_status - --status with valid cookie returns user info.""" - with patch('kb_workflow.kb_workflow.get_settings', return_value=mock_settings_with_cookie), \ + with patch('kb_workflow.pipeline.config.get_settings', return_value=mock_settings_with_cookie), \ patch('kb_workflow.kb_workflow.asyncio.run') as mock_run, \ - patch('kb_workflow.kb_workflow.ZhihuLogin') as MockLogin: + patch('kb_workflow.auth.ZhihuLogin') as MockLogin: - # Setup mock + # Setup mock - asyncio.run returns the result object directly mock_login_instance = MagicMock() MockLogin.return_value = mock_login_instance - mock_run.return_value = mock_validation_result_valid.model_dump() + mock_run.return_value = mock_validation_result_valid # Return object, not dict # Act result = login_command(status=True) @@ -135,12 +136,14 @@ class TestLoginCommand: self, mock_settings, mock_login_result_success ): """Test 4: login_integration - Successful login saves cookie.""" - with patch('kb_workflow.kb_workflow.get_settings', return_value=mock_settings), \ - patch('kb_workflow.kb_workflow.CookieManager') as MockManager, \ - patch('kb_workflow.kb_workflow.asyncio.run') as mock_run: + with patch('kb_workflow.pipeline.config.get_settings', return_value=mock_settings), \ + patch('kb_workflow.auth.cookie_manager.CookieManager') as MockManager, \ + patch('kb_workflow.auth.ZhihuLogin') as MockLogin: # Setup mock - mock_run.return_value = mock_login_result_success.model_dump() + mock_login_instance = MagicMock() + mock_login_instance.login_by_qrcode.return_value = mock_login_result_success + MockLogin.return_value = mock_login_instance mock_manager = MagicMock() MockManager.return_value = mock_manager @@ -155,7 +158,7 @@ class TestLoginCommand: def test_login_error_handling(self, mock_settings): """Test 5: login handles exceptions gracefully.""" - with patch('kb_workflow.kb_workflow.get_settings', side_effect=Exception("Config error")): + with patch('kb_workflow.pipeline.config.get_settings', side_effect=Exception("Config error")): # Act result = login_command()