重构登录,减少报错

src/core/login.py

@@ -11,18 +11,6 @@ from src.entity.database import db
 from src.ui import global_signals
 
 
-def handle_request(request, result_container, response_container):
-    if 'loginFlow' in request.url:
-        response = request.response()
-        if response:
-            data = response.json()
-            response_container['data'] = data
-    if 'banner' in request.url:
-        headers = request.headers
-        result_container['headers'] = headers
-        logger.info(f'Successfully got headers: {headers}')
-
-
 def playwright_login(account: Account) -> dict:
     result_container = {'headers': {}}
     response_container = {'data': {}}
@@ -34,6 +22,7 @@ def playwright_login(account: Account) -> dict:
     except Exception as e:
         logger.error(f'Error during login for account {account.username}: {e}', exc_info=True)
         handle_login_failure(account)
+        return {}
 
 
 def perform_login(playwright, account: Account, request_container: dict, response_container: dict):
@@ -41,44 +30,61 @@ def perform_login(playwright, account: Account, request_container: dict, respons
     context = browser.new_context()
     page = context.new_page()
 
-    page.on("request", lambda request: handle_request(request, request_container, response_container))
-
-    page.goto(account.url, timeout=60000)
+    page.goto(account.url, timeout=120000)
 
     # 填充基本的元素
     fill_form_common(page, account)
 
     if account.type == AccountType.jy:
         handle_jy_captcha(page, account)
-        logger.info("验证码处理成功")
-        time.sleep(3)
     else:
-        max_retries = 3
-        for attempt in range(max_retries):
-            if handle_captcha(page):
-                page.locator("div").filter(has_text=re.compile(r"^登录$")).get_by_role("button").click()
-                time.sleep(2)
-                try:
-                    if response_container['data'] and response_container['data'].get('status_code') != 6000:
-                        logger.warning(f"登录发生错误, {response_container['data'].get('message', '未知错误')}")
-                        close_button = page.locator("button.ant-modal-close")
-                        close_button.wait_for()
-                        close_button.click()
-                    else:
-                        break
-                except TimeoutError:
-                    logger.warning(f"等待响应超时，重试 {attempt + 1}/{max_retries}")
-            else:
-                logger.warning(f"验证码输入失败，重试 {attempt + 1}/{max_retries}")
-        else:
-            logger.warning("登录失败，达到最大重试次数")
-    persistence(account, response_container['data'])
-    # 通知app数据更新了
-    global_signals.user_data_updated.emit()
-    page.wait_for_url(f'{account.url}/app/home?showWelcome=false')
+        handle_generic_account(page, response_container)
+
+    # 处理 'banner' 请求并获取响应头
+    handle_banner_request(page, request_container)
+
+    # 仅当 headers 不为空时才保存到数据库中
+    if request_container['headers']:
+        persistence(account, request_container['headers'])
+        global_signals.user_data_updated.emit()
+
     close_resources(page, context, browser)
 
 
+def handle_generic_account(page, response_container):
+    max_retries = 3
+    for attempt in range(max_retries):
+        if handle_captcha(page):
+            with page.expect_request(lambda request: 'loginFlow' in request.url, timeout=30000):
+                page.locator("div").filter(has_text=re.compile(r"^登录$")).get_by_role("button").click()
+
+            # 等待 loginFlow 请求的响应
+            with page.expect_response(lambda response: 'loginFlow' in response.url, timeout=30000) as response_info:
+                response = response_info.value
+                data = response.json()
+                response_container['data'] = data
+                if response_container['data'] and response_container['data'].get('status_code') != 6000:
+                    logger.warning(f"登录发生错误, {response_container['data'].get('message', '未知错误')}")
+                    close_button = page.locator("button.ant-modal-close")
+                    close_button.wait_for()
+                    close_button.click()
+                else:
+                    break
+    else:
+        logger.warning("登录失败，达到最大重试次数")
+
+
+def handle_banner_request(page, request_container):
+    with page.expect_request(lambda request: 'banner' in request.url, timeout=30000) as request_info:
+        banner_request = request_info.value
+
+    with page.expect_response(lambda response: 'banner' in response.url, timeout=30000) as response_info:
+        banner_response = response_info.value
+        headers = banner_request.headers
+        request_container['headers'] = headers
+        logger.info(f'Successfully got headers: {headers}')
+
+
 def fill_form_common(page, account: Account):
     username_input = page.get_by_placeholder('用户名')
     password_input = page.get_by_placeholder('密码')
@@ -167,7 +173,7 @@ def login(account: Account):
         account.headers = headers
         return account
     else:
-        pass
+        return account
 
 
 def persistence(account: Account, headers: dict):