From 04563462e46206a49291070d036a0adce5d22556 Mon Sep 17 00:00:00 2001 From: chengchongzhen <15939054361@163.com> Date: Tue, 31 Dec 2024 16:26:09 +0800 Subject: [PATCH] refactor: streamline cookie retrieval and enhance error handling in account management scripts --- cursor_pro_keep_alive.py | 52 +++++++++++++++++++++------------------- get_veri_code.py | 16 ++++--------- 2 files changed, 31 insertions(+), 37 deletions(-) diff --git a/cursor_pro_keep_alive.py b/cursor_pro_keep_alive.py index be037c0..deb3481 100644 --- a/cursor_pro_keep_alive.py +++ b/cursor_pro_keep_alive.py @@ -5,13 +5,10 @@ from license_manager import LicenseManager os.environ["PYTHONVERBOSE"] = "0" os.environ["PYINSTALLER_VERBOSE"] = "0" -from DrissionPage import ChromiumOptions, Chromium -from DrissionPage.common import Keys import re import time import random from cursor_auth_manager import CursorAuthManager -from configparser import ConfigParser import os import sys import logging @@ -141,9 +138,6 @@ def delete_account(browser, tab): handle_turnstile(tab) time.sleep(5) - # tab.get_screenshot('sign-in_success.png') - # print("登录账户截图") - tab.get(settings_url) print("进入设置页面") @@ -188,17 +182,18 @@ def delete_account(browser, tab): def get_cursor_session_token(tab): - """获取cursor session token""" print("开始获取cookie") - cookies = tab.cookies() - cursor_session_token = None - for cookie in cookies: - if cookie["name"] == "WorkosCursorSessionToken": - cursor_session_token = cookie["value"].split("%3A%3A")[1] - break - if not cursor_session_token: + try: + # 获取指定域名的所有cookies + cookies = tab.cookies() + for cookie in cookies: + if cookie.get("name") == "WorkosCursorSessionToken": + return cookie["value"].split("%3A%3A")[1] print("未能获取到CursorSessionToken") - return cursor_session_token + return None + except Exception as e: + print(f"获取cookie失败: {str(e)}") + return None def update_cursor_auth(email=None, access_token=None, refresh_token=None): @@ -261,11 +256,7 @@ def sign_up_account(browser, tab): break if tab.ele("@data-index=0"): code = email_handler.get_verification_code(account) - - if code: - print("获取验证码成功：", code) - else: - print("获取验证码失败，程序退出") + if not code: return False i = 0 @@ -278,7 +269,21 @@ def sign_up_account(browser, tab): print(e) handle_turnstile(tab) - time.sleep(random.uniform(1, 3)) + print("等待进入设置页面") + + # 等待页面加载完成 + max_wait_time = 30 # 最大等待时间（秒） + start_time = time.time() + while time.time() - start_time < max_wait_time: + try: + # 检查页面是否仍在加载 + is_loading = tab.run_js("return document.readyState !== 'complete'") + if not is_loading: + break + time.sleep(1) + except Exception as e: + print(f"检查页面加载状态时出错: {str(e)}") + break print("进入设置页面") tab.get(settings_url) try: @@ -398,12 +403,9 @@ if __name__ == "__main__": tab.get(login_url) - print("开始注册账户") if sign_up_account(browser, tab): token = get_cursor_session_token(tab) - # print(f"CursorSessionToken: {token}") - print("账户注册成功") - if auto_update_cursor_auth: + if token: update_cursor_auth( email=account, access_token=token, refresh_token=token ) diff --git a/get_veri_code.py b/get_veri_code.py index 14833d3..7774f10 100644 --- a/get_veri_code.py +++ b/get_veri_code.py @@ -1,7 +1,6 @@ from DrissionPage.common import Keys import time import re -import logging class EmailVerificationHandler: @@ -10,15 +9,14 @@ class EmailVerificationHandler: self.mail_url = mail_url def get_verification_code(self, email): - """获取邮箱验证码""" username = email.split("@")[0] code = None try: + print("打开邮箱页面") # 打开新标签页访问临时邮箱 tab_mail = self.browser.new_tab(self.mail_url) self.browser.activate_tab(tab_mail) - print("打开邮箱页面") # 输入用户名 self._input_username(tab_mail, username) @@ -33,12 +31,11 @@ class EmailVerificationHandler: tab_mail.close() except Exception as e: - logging.error(f"获取验证码失败: {str(e)}") + print(f"获取验证码失败: {str(e)}") return code def _input_username(self, tab, username): - """输入用户名""" while True: if tab.ele("@id=pre_button"): tab.actions.click("@id=pre_button") @@ -50,17 +47,14 @@ class EmailVerificationHandler: time.sleep(1) def _get_latest_mail_code(self, tab): - """获取最新邮件中的验证码""" code = None while True: new_mail = tab.ele("@class=mail") if new_mail: if new_mail.text: - logging.info(f"最新的邮件：{new_mail.text}") tab.actions.click("@class=mail") break else: - logging.info(str(new_mail)) break time.sleep(1) @@ -71,18 +65,16 @@ class EmailVerificationHandler: ) if verification_code: code = verification_code.group(1) - logging.info(f"验证码：{code}") + print(f"验证码：{code}") else: - logging.warning("未找到验证码") + print("未找到验证码") return code def _cleanup_mail(self, tab): - """清理邮件""" if tab.ele("@id=delete_mail"): tab.actions.click("@id=delete_mail") time.sleep(1) if tab.ele("@id=confirm_mail"): tab.actions.click("@id=confirm_mail") - logging.info("删除邮件")