This commit is contained in:
huangzhenpc
2025-04-02 09:28:25 +08:00
parent 31fe73f998
commit d16f6bdc62
3 changed files with 359 additions and 70 deletions

View File

@@ -31,6 +31,9 @@ class AutoCursorService:
self.config = Config.from_yaml() self.config = Config.from_yaml()
self.db_manager = DatabaseManager(self.config) self.db_manager = DatabaseManager(self.config)
# 添加EmailManager实例用于操作黑名单
self.email_manager = None
# API相关 # API相关
self.api_base_url = "https://cursorapi.nosqli.com/admin/api.AutoCursor" self.api_base_url = "https://cursorapi.nosqli.com/admin/api.AutoCursor"
self.proxy = None self.proxy = None
@@ -94,6 +97,11 @@ class AutoCursorService:
logger.info("初始化自动化服务") logger.info("初始化自动化服务")
await self.db_manager.initialize() await self.db_manager.initialize()
# 初始化EmailManager
from services.email_manager import EmailManager
self.email_manager = EmailManager(self.config, self.db_manager)
await self.email_manager.initialize()
# 检查并设置代理 # 检查并设置代理
if hasattr(self.config, "proxy_config") and self.config.proxy_config: if hasattr(self.config, "proxy_config") and self.config.proxy_config:
if hasattr(self.config.proxy_config, "api_proxy") and self.config.proxy_config.api_proxy: if hasattr(self.config.proxy_config, "api_proxy") and self.config.proxy_config.api_proxy:
@@ -188,6 +196,12 @@ class AutoCursorService:
return 0 return 0
count = 0 count = 0
blacklist_added = 0
# 如果有EmailManager确保黑名单初始化
if self.email_manager and self.email_manager.use_redis:
await self.email_manager._ensure_blacklist_initialized()
for account in accounts: for account in accounts:
try: try:
email = account.get("email", "") email = account.get("email", "")
@@ -199,6 +213,12 @@ class AutoCursorService:
logger.warning(f"账号数据不完整: {account}") logger.warning(f"账号数据不完整: {account}")
continue continue
# 检查邮箱是否已在黑名单中
if self.email_manager and await self.email_manager.is_email_blacklisted(email):
logger.warning(f"跳过黑名单中的邮箱: {email}")
blacklist_added += 1
continue
# 插入数据库 # 插入数据库
insert_query = ''' insert_query = '''
INSERT INTO email_accounts INSERT INTO email_accounts
@@ -221,11 +241,19 @@ class AutoCursorService:
except Exception as e: except Exception as e:
logger.error(f"导入邮箱账号时出错: {e}") logger.error(f"导入邮箱账号时出错: {e}")
if blacklist_added > 0:
logger.warning(f"已跳过 {blacklist_added} 个黑名单中的邮箱")
logger.success(f"成功导入 {count} 个邮箱账号") logger.success(f"成功导入 {count} 个邮箱账号")
return count return count
async def count_pending_accounts(self) -> int: async def count_pending_accounts(self) -> int:
"""统计可用的pending状态账号数量""" """统计可用的pending状态账号数量"""
if self.email_manager:
# 使用EmailManager的方法确保黑名单一致性
return await self.email_manager.count_pending_accounts()
# 兼容旧代码,直接查询
query = """ query = """
SELECT COUNT(*) SELECT COUNT(*)
FROM email_accounts FROM email_accounts

View File

@@ -35,6 +35,8 @@ class CursorRegister:
async def initialize(self): async def initialize(self):
"""初始化数据库""" """初始化数据库"""
await self.db_manager.initialize() await self.db_manager.initialize()
# 确保EmailManager完成初始化
await self.email_manager.initialize()
async def cleanup(self): async def cleanup(self):
"""清理资源""" """清理资源"""
@@ -160,6 +162,10 @@ async def main():
register.logger.info(f"当前总进度: 已注册 {total_registered} 个账号") register.logger.info(f"当前总进度: 已注册 {total_registered} 个账号")
# 批次结束后等待3秒确保所有数据库更新都已完成
register.logger.info("本批次注册完成,等待数据库状态完全更新...")
await asyncio.sleep(3)
# 如果本批次注册失败率过高,暂停一段时间 # 如果本批次注册失败率过高,暂停一段时间
if successful < batch_size * 0.5 and successful > 0: # 成功率低于50%但不为零 if successful < batch_size * 0.5 and successful > 0: # 成功率低于50%但不为零
register.logger.warning("本批次成功率过低暂停60秒后继续") register.logger.warning("本批次成功率过低暂停60秒后继续")

View File

@@ -6,6 +6,7 @@ from typing import Dict, List, Optional
import aiohttp import aiohttp
from loguru import logger from loguru import logger
import aiomysql
from core.config import Config from core.config import Config
from core.database import DatabaseManager from core.database import DatabaseManager
@@ -34,106 +35,360 @@ class EmailManager:
"Verify your email address", "Verify your email address",
"Complete code challenge", "Complete code challenge",
] ]
# Redis相关配置
self.use_redis = False
if hasattr(self.db, 'redis') and self.db.redis:
self.use_redis = True
logger.info("Redis可用将使用Redis进行邮箱状态管理")
else:
logger.warning("Redis不可用将使用MySQL进行邮箱状态管理")
# Redis键前缀
self.redis_prefix = "emailmanager:"
# 锁和黑名单的过期时间(秒)
self.lock_timeout = 600 # 10分钟
self.blacklist_timeout = 86400 * 30 # 30天
# 黑名单初始化标记
self.blacklist_initialized = False
async def initialize(self):
"""初始化EmailManager确保在使用前完成必要的设置"""
if self.use_redis:
await self._ensure_blacklist_initialized()
self.blacklist_initialized = True
async def _ensure_blacklist_initialized(self):
"""确保黑名单已经初始化"""
if not self.use_redis:
return False
blacklist_key = f"{self.redis_prefix}blacklist:initialized"
initialized = await self.db.redis.exists(blacklist_key)
if not initialized:
logger.info("初始化Redis邮箱黑名单...")
# 查询所有已成功或不可用的邮箱
query = """
SELECT DISTINCT email
FROM email_accounts
WHERE status = 'success' OR status = 'unavailable'
"""
results = await self.db.fetch_all(query)
if results:
# 批量添加到黑名单
blacklist_key = f"{self.redis_prefix}blacklist:emails"
emails = [row['email'] for row in results]
if emails:
pipeline = self.db.redis.pipeline()
for email in emails:
pipeline.sadd(blacklist_key, email)
pipeline.expire(blacklist_key, self.blacklist_timeout)
await pipeline.execute()
logger.info(f"已将 {len(emails)} 个邮箱添加到黑名单")
# 标记为已初始化
await self.db.redis.setex(f"{self.redis_prefix}blacklist:initialized", self.blacklist_timeout, "1")
return True
return True
async def is_email_blacklisted(self, email: str) -> bool:
"""检查邮箱是否在黑名单中"""
if not self.use_redis:
# 回退到数据库查询
query = """
SELECT 1 FROM email_accounts
WHERE email = %s AND (status = 'success' OR status = 'unavailable')
LIMIT 1
"""
result = await self.db.fetch_one(query, (email,))
return result is not None
# 使用Redis SET存储黑名单
blacklist_key = f"{self.redis_prefix}blacklist:emails"
return await self.db.redis.sismember(blacklist_key, email)
async def add_email_to_blacklist(self, email: str):
"""将邮箱添加到黑名单"""
if not self.use_redis:
return False
blacklist_key = f"{self.redis_prefix}blacklist:emails"
await self.db.redis.sadd(blacklist_key, email)
await self.db.redis.expire(blacklist_key, self.blacklist_timeout)
logger.debug(f"邮箱 {email} 已添加到黑名单")
return True
async def lock_account(self, account_id: int) -> bool:
"""锁定账号"""
if not self.use_redis:
# 如果不使用Redis通过数据库更新来锁定
try:
query = """
UPDATE email_accounts
SET in_use = 1, updated_at = CURRENT_TIMESTAMP
WHERE id = %s AND in_use = 0
"""
affected = await self.db.execute(query, (account_id,))
return affected > 0
except Exception as e:
logger.error(f"通过数据库锁定账号 {account_id} 失败: {e}")
return False
# 使用Redis实现分布式锁
lock_key = f"{self.redis_prefix}lock:account:{account_id}"
locked = await self.db.redis.setnx(lock_key, "1")
if locked:
# 锁定成功,设置过期时间
await self.db.redis.expire(lock_key, self.lock_timeout)
# 同时更新数据库状态
try:
update_query = """
UPDATE email_accounts
SET in_use = 1, updated_at = CURRENT_TIMESTAMP
WHERE id = %s
"""
await self.db.execute(update_query, (account_id,))
except Exception as e:
# 如果数据库更新失败释放Redis锁
logger.error(f"锁定账号 {account_id} 后更新数据库失败: {e}")
await self.db.redis.delete(lock_key)
return False
logger.debug(f"账号 {account_id} 已锁定")
return locked
async def unlock_account(self, account_id: int) -> bool:
"""解锁账号"""
# 无论是否使用Redis都更新数据库
try:
update_query = """
UPDATE email_accounts
SET in_use = 0, updated_at = CURRENT_TIMESTAMP
WHERE id = %s
"""
await self.db.execute(update_query, (account_id,))
except Exception as e:
logger.error(f"解锁账号 {account_id} 更新数据库失败: {e}")
if not self.use_redis:
return True
# 删除Redis锁
lock_key = f"{self.redis_prefix}lock:account:{account_id}"
deleted = await self.db.redis.delete(lock_key)
logger.debug(f"账号 {account_id} 锁已释放")
return deleted > 0
async def batch_get_accounts(self, num: int) -> List[EmailAccount]: async def batch_get_accounts(self, num: int) -> List[EmailAccount]:
"""批量获取未使用的邮箱账号""" """批量获取未使用的邮箱账号"""
logger.info(f"尝试获取 {num} 个未使用的邮箱账号") logger.info(f"尝试获取 {num} 个未使用的邮箱账号")
# 使用简单条件获取未使用的账号 # 如果使用Redis确保黑名单已初始化
if self.use_redis:
await self._ensure_blacklist_initialized()
# 1. 先从数据库中获取候选账号
select_query = """ select_query = """
SELECT id, email, password, client_id, refresh_token SELECT id, email, password, client_id, refresh_token
FROM email_accounts FROM email_accounts
WHERE in_use = 0 AND sold = 0 AND status = 'pending' WHERE in_use = 0 AND sold = 0 AND status = 'pending'
AND email NOT IN (
SELECT email FROM email_accounts
WHERE status = 'success' OR status = 'unavailable'
)
LIMIT %s LIMIT %s
""" """
accounts = await self.db.fetch_all(select_query, (num,)) # 多获取一些候选账号,防止有些被排除
candidate_accounts = await self.db.fetch_all(select_query, (num * 2,))
if not accounts: if not candidate_accounts:
logger.debug("没有找到符合条件的账号") logger.debug("没有找到符合条件的候选账号")
return [] return []
# 2. 提取账号ID列表
account_ids = [account['id'] for account in accounts]
# 3. 更新这些账号的状态,设置注册开始时间戳,避免其他进程获取
if account_ids:
placeholders = ', '.join(['%s' for _ in account_ids])
update_query = f"""
UPDATE email_accounts
SET in_use = 1, updated_at = CURRENT_TIMESTAMP
WHERE id IN ({placeholders})
"""
await self.db.execute(update_query, tuple(account_ids))
# 记录被锁定的账号信息 logger.debug(f"找到 {len(candidate_accounts)} 个候选账号")
logger.info(f"已锁定 {len(account_ids)} 个账号用于注册: {account_ids}")
# 4. 返回账号数据 # 2. 筛选并锁定账号
logger.debug(f"实际获取到 {len(accounts)} 个账号") result_accounts = []
return [ for account in candidate_accounts:
EmailAccount( # 检查邮箱是否在黑名单中
id=row['id'], if await self.is_email_blacklisted(account['email']):
email=row['email'], logger.debug(f"邮箱 {account['email']} 在黑名单中,跳过")
password=row['password'], continue
client_id=row['client_id'],
refresh_token=row['refresh_token'], # 尝试锁定账号
in_use=True if await self.lock_account(account['id']):
) # 添加到结果列表
for row in accounts result_accounts.append(EmailAccount(
] id=account['id'],
email=account['email'],
password=account['password'],
client_id=account['client_id'],
refresh_token=account['refresh_token'],
in_use=True
))
# 如果已经获取足够的账号,退出循环
if len(result_accounts) >= num:
break
else:
logger.debug(f"账号 {account['id']} 锁定失败,可能被其他进程使用")
logger.info(f"实际获取到 {len(result_accounts)} 个可用账号")
# 如果账号数量不足,尝试清理长时间锁定但未更新的账号
if len(result_accounts) < num and len(result_accounts) < len(candidate_accounts):
logger.warning("可用账号不足,尝试清理长时间锁定的账号")
await self._cleanup_stuck_accounts()
return result_accounts
async def _cleanup_stuck_accounts(self):
"""清理长时间锁定但未更新的账号"""
try:
# 清理超过30分钟未更新且仍标记为in_use=1的账号
cleanup_query = """
UPDATE email_accounts
SET in_use = 0
WHERE in_use = 1
AND updated_at < DATE_SUB(NOW(), INTERVAL 30 MINUTE)
"""
affected = await self.db.execute(cleanup_query)
if affected > 0:
logger.info(f"已清理 {affected} 个长时间锁定的账号")
# 如果使用Redis同时清理对应的锁
if self.use_redis:
# 清理可能存在的Redis锁但这需要知道具体的account_id
# 这里简化处理,依赖锁的自动过期机制
pass
except Exception as e:
logger.error(f"清理账号时出错: {e}")
async def update_account_status(self, account_id: int, status: str): async def update_account_status(self, account_id: int, status: str):
"""更新账号状态""" """更新账号状态"""
query = ''' try:
UPDATE email_accounts # 获取账号邮箱信息(用于黑名单)
SET account_query = "SELECT email, status as current_status FROM email_accounts WHERE id = %s"
status = %s, account_info = await self.db.fetch_one(account_query, (account_id,))
in_use = 0,
updated_at = CURRENT_TIMESTAMP if not account_info:
WHERE id = %s logger.error(f"账号 {account_id} 不存在")
''' return
await self.db.execute(query, (status, account_id))
# 检查状态变更是否合理
current_status = account_info.get('current_status')
if current_status == 'success' and status != 'success':
logger.warning(f"警告: 尝试将成功账号 {account_id} 状态改为 {status},这可能是不正确的操作")
# 如果已经是success状态不允许降级为其他状态
# return
# 更新数据库状态
query = '''
UPDATE email_accounts
SET
status = %s,
in_use = 0,
updated_at = CURRENT_TIMESTAMP
WHERE id = %s
'''
await self.db.execute(query, (status, account_id))
logger.info(f"账号 {account_id} 状态已更新为 {status}")
# 如果是success或unavailable状态添加到黑名单
if account_info and (status == 'success' or status == 'unavailable'):
email = account_info['email']
logger.debug(f"将邮箱 {email} 添加到黑名单 (状态: {status})")
await self.add_email_to_blacklist(email)
# 解锁账号
await self.unlock_account(account_id)
# 清除数据库缓存
if self.db.redis:
await self.db.clear_cache("db:*email_accounts*")
except Exception as e:
logger.error(f"更新账号 {account_id} 状态为 {status} 时出错: {e}")
# 确保无论如何都解锁账号
try:
await self.unlock_account(account_id)
except Exception as unlock_error:
logger.error(f"尝试解锁账号 {account_id} 时出错: {unlock_error}")
raise
async def update_account(self, account_id: int, cursor_password: str, cursor_cookie: str, cursor_token: str): async def update_account(self, account_id: int, cursor_password: str, cursor_cookie: str, cursor_token: str):
"""更新账号信息""" """更新账号信息(注册成功)"""
query = ''' try:
UPDATE email_accounts # 获取账号邮箱信息(用于黑名单)
SET account_query = "SELECT email FROM email_accounts WHERE id = %s"
cursor_password = %s, account_info = await self.db.fetch_one(account_query, (account_id,))
cursor_cookie = %s,
cursor_token = %s, # 更新数据库
in_use = 0, query = '''
sold = 1, UPDATE email_accounts
status = 'success', SET
updated_at = CURRENT_TIMESTAMP cursor_password = %s,
WHERE id = %s cursor_cookie = %s,
''' cursor_token = %s,
await self.db.execute(query, (cursor_password, cursor_cookie, cursor_token, account_id)) in_use = 0,
sold = 1,
status = 'success',
updated_at = CURRENT_TIMESTAMP
WHERE id = %s
'''
await self.db.execute(query, (cursor_password, cursor_cookie, cursor_token, account_id))
logger.info(f"账号 {account_id} 更新为注册成功")
# 添加到黑名单
if account_info:
email = account_info['email']
logger.debug(f"将邮箱 {email} 添加到黑名单 (注册成功)")
await self.add_email_to_blacklist(email)
# 解锁账号
await self.unlock_account(account_id)
# 清除数据库缓存
if self.db.redis:
await self.db.clear_cache("db:*email_accounts*")
except Exception as e:
logger.error(f"更新账号 {account_id} 信息时出错: {e}")
# 确保无论如何都解锁账号
try:
await self.unlock_account(account_id)
except:
pass
raise
async def release_account(self, account_id: int): async def release_account(self, account_id: int):
"""释放账号""" """释放账号"""
query = ''' try:
UPDATE email_accounts await self.unlock_account(account_id)
SET in_use = 0, updated_at = CURRENT_TIMESTAMP logger.info(f"账号 {account_id} 已释放")
WHERE id = %s
''' # 清除数据库缓存
await self.db.execute(query, (account_id,)) if self.db.redis:
await self.db.clear_cache("db:*email_accounts*")
except Exception as e:
logger.error(f"释放账号 {account_id} 时出错: {e}")
raise
async def count_pending_accounts(self) -> int: async def count_pending_accounts(self) -> int:
"""统计可用的pending状态账号数量""" """统计可用的pending状态账号数量"""
if self.use_redis:
await self._ensure_blacklist_initialized()
query = """ query = """
SELECT COUNT(*) SELECT COUNT(*)
FROM email_accounts FROM email_accounts
WHERE status = 'pending' AND in_use = 0 AND sold = 0 WHERE status = 'pending' AND in_use = 0 AND sold = 0
AND email NOT IN (
SELECT email FROM email_accounts
WHERE status = 'success' OR status = 'unavailable'
)
""" """
# 注:这里不使用黑名单过滤,因为数据量可能很大,
# 但实际获取账号时会应用黑名单过滤
result = await self.db.fetch_one(query) result = await self.db.fetch_one(query)
if result: if result:
return result.get("COUNT(*)", 0) return result.get("COUNT(*)", 0)