refactor: streamline cookie retrieval and enhance error handling in account management scripts

This commit is contained in:
chengchongzhen
2024-12-31 16:26:09 +08:00
parent 50e69cd51c
commit 04563462e4
2 changed files with 31 additions and 37 deletions

View File

@@ -5,13 +5,10 @@ from license_manager import LicenseManager
os.environ["PYTHONVERBOSE"] = "0" os.environ["PYTHONVERBOSE"] = "0"
os.environ["PYINSTALLER_VERBOSE"] = "0" os.environ["PYINSTALLER_VERBOSE"] = "0"
from DrissionPage import ChromiumOptions, Chromium
from DrissionPage.common import Keys
import re import re
import time import time
import random import random
from cursor_auth_manager import CursorAuthManager from cursor_auth_manager import CursorAuthManager
from configparser import ConfigParser
import os import os
import sys import sys
import logging import logging
@@ -141,9 +138,6 @@ def delete_account(browser, tab):
handle_turnstile(tab) handle_turnstile(tab)
time.sleep(5) time.sleep(5)
# tab.get_screenshot('sign-in_success.png')
# print("登录账户截图")
tab.get(settings_url) tab.get(settings_url)
print("进入设置页面") print("进入设置页面")
@@ -188,17 +182,18 @@ def delete_account(browser, tab):
def get_cursor_session_token(tab): def get_cursor_session_token(tab):
"""获取cursor session token"""
print("开始获取cookie") print("开始获取cookie")
cookies = tab.cookies() try:
cursor_session_token = None # 获取指定域名的所有cookies
for cookie in cookies: cookies = tab.cookies()
if cookie["name"] == "WorkosCursorSessionToken": for cookie in cookies:
cursor_session_token = cookie["value"].split("%3A%3A")[1] if cookie.get("name") == "WorkosCursorSessionToken":
break return cookie["value"].split("%3A%3A")[1]
if not cursor_session_token:
print("未能获取到CursorSessionToken") 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): def update_cursor_auth(email=None, access_token=None, refresh_token=None):
@@ -261,11 +256,7 @@ def sign_up_account(browser, tab):
break break
if tab.ele("@data-index=0"): if tab.ele("@data-index=0"):
code = email_handler.get_verification_code(account) code = email_handler.get_verification_code(account)
if not code:
if code:
print("获取验证码成功:", code)
else:
print("获取验证码失败,程序退出")
return False return False
i = 0 i = 0
@@ -278,7 +269,21 @@ def sign_up_account(browser, tab):
print(e) print(e)
handle_turnstile(tab) 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("进入设置页面") print("进入设置页面")
tab.get(settings_url) tab.get(settings_url)
try: try:
@@ -398,12 +403,9 @@ if __name__ == "__main__":
tab.get(login_url) tab.get(login_url)
print("开始注册账户")
if sign_up_account(browser, tab): if sign_up_account(browser, tab):
token = get_cursor_session_token(tab) token = get_cursor_session_token(tab)
# print(f"CursorSessionToken: {token}") if token:
print("账户注册成功")
if auto_update_cursor_auth:
update_cursor_auth( update_cursor_auth(
email=account, access_token=token, refresh_token=token email=account, access_token=token, refresh_token=token
) )

View File

@@ -1,7 +1,6 @@
from DrissionPage.common import Keys from DrissionPage.common import Keys
import time import time
import re import re
import logging
class EmailVerificationHandler: class EmailVerificationHandler:
@@ -10,15 +9,14 @@ class EmailVerificationHandler:
self.mail_url = mail_url self.mail_url = mail_url
def get_verification_code(self, email): def get_verification_code(self, email):
"""获取邮箱验证码"""
username = email.split("@")[0] username = email.split("@")[0]
code = None code = None
try: try:
print("打开邮箱页面")
# 打开新标签页访问临时邮箱 # 打开新标签页访问临时邮箱
tab_mail = self.browser.new_tab(self.mail_url) tab_mail = self.browser.new_tab(self.mail_url)
self.browser.activate_tab(tab_mail) self.browser.activate_tab(tab_mail)
print("打开邮箱页面")
# 输入用户名 # 输入用户名
self._input_username(tab_mail, username) self._input_username(tab_mail, username)
@@ -33,12 +31,11 @@ class EmailVerificationHandler:
tab_mail.close() tab_mail.close()
except Exception as e: except Exception as e:
logging.error(f"获取验证码失败: {str(e)}") print(f"获取验证码失败: {str(e)}")
return code return code
def _input_username(self, tab, username): def _input_username(self, tab, username):
"""输入用户名"""
while True: while True:
if tab.ele("@id=pre_button"): if tab.ele("@id=pre_button"):
tab.actions.click("@id=pre_button") tab.actions.click("@id=pre_button")
@@ -50,17 +47,14 @@ class EmailVerificationHandler:
time.sleep(1) time.sleep(1)
def _get_latest_mail_code(self, tab): def _get_latest_mail_code(self, tab):
"""获取最新邮件中的验证码"""
code = None code = None
while True: while True:
new_mail = tab.ele("@class=mail") new_mail = tab.ele("@class=mail")
if new_mail: if new_mail:
if new_mail.text: if new_mail.text:
logging.info(f"最新的邮件:{new_mail.text}")
tab.actions.click("@class=mail") tab.actions.click("@class=mail")
break break
else: else:
logging.info(str(new_mail))
break break
time.sleep(1) time.sleep(1)
@@ -71,18 +65,16 @@ class EmailVerificationHandler:
) )
if verification_code: if verification_code:
code = verification_code.group(1) code = verification_code.group(1)
logging.info(f"验证码:{code}") print(f"验证码:{code}")
else: else:
logging.warning("未找到验证码") print("未找到验证码")
return code return code
def _cleanup_mail(self, tab): def _cleanup_mail(self, tab):
"""清理邮件"""
if tab.ele("@id=delete_mail"): if tab.ele("@id=delete_mail"):
tab.actions.click("@id=delete_mail") tab.actions.click("@id=delete_mail")
time.sleep(1) time.sleep(1)
if tab.ele("@id=confirm_mail"): if tab.ele("@id=confirm_mail"):
tab.actions.click("@id=confirm_mail") tab.actions.click("@id=confirm_mail")
logging.info("删除邮件")