ssss

auto_cursor_service.py

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

main.py

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

services/email_manager.py

@@ -6,6 +6,7 @@ from typing import Dict, List, Optional
 
 import aiohttp
 from loguru import logger
+import aiomysql
 
 from core.config import Config
 from core.database import DatabaseManager
@@ -34,106 +35,360 @@ class EmailManager:
             "Verify your email address",
             "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]:
         """批量获取未使用的邮箱账号"""
         logger.info(f"尝试获取 {num} 个未使用的邮箱账号")
         
-        # 使用简单条件获取未使用的账号
+        # 如果使用Redis，确保黑名单已初始化
+        if self.use_redis:
+            await self._ensure_blacklist_initialized()
+        
+        # 1. 先从数据库中获取候选账号
         select_query = """
             SELECT id, email, password, client_id, refresh_token 
             FROM email_accounts 
             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
         """
-        accounts = await self.db.fetch_all(select_query, (num,))
+        # 多获取一些候选账号，防止有些被排除
+        candidate_accounts = await self.db.fetch_all(select_query, (num * 2,))
         
-        if not accounts:
-            logger.debug("没有找到符合条件的账号")
+        if not candidate_accounts:
+            logger.debug("没有找到符合条件的候选账号")
             return []
             
-        # 2. 提取账号ID列表
-        account_ids = [account['id'] for account in accounts]
+        logger.debug(f"找到 {len(candidate_accounts)} 个候选账号")
         
-        # 3. 更新这些账号的状态，设置注册开始时间戳，避免其他进程获取
-        if account_ids:
-            placeholders = ', '.join(['%s' for _ in account_ids])
-            update_query = f"""
+        # 2. 筛选并锁定账号
+        result_accounts = []
+        for account in candidate_accounts:
+            # 检查邮箱是否在黑名单中
+            if await self.is_email_blacklisted(account['email']):
+                logger.debug(f"邮箱 {account['email']} 在黑名单中，跳过")
+                continue
+                
+            # 尝试锁定账号
+            if await self.lock_account(account['id']):
+                # 添加到结果列表
+                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 = 1, updated_at = CURRENT_TIMESTAMP
-                WHERE id IN ({placeholders})
+                SET in_use = 0 
+                WHERE in_use = 1 
+                AND updated_at < DATE_SUB(NOW(), INTERVAL 30 MINUTE)
             """
-            await self.db.execute(update_query, tuple(account_ids))
+            affected = await self.db.execute(cleanup_query)
+            if affected > 0:
+                logger.info(f"已清理 {affected} 个长时间锁定的账号")
                 
-            # 记录被锁定的账号信息
-            logger.info(f"已锁定 {len(account_ids)} 个账号用于注册: {account_ids}")
-        
-        # 4. 返回账号数据
-        logger.debug(f"实际获取到 {len(accounts)} 个账号")
-        return [
-            EmailAccount(
-                id=row['id'],
-                email=row['email'],
-                password=row['password'],
-                client_id=row['client_id'],
-                refresh_token=row['refresh_token'],
-                in_use=True
-            )
-            for row in accounts
-        ]
+                # 如果使用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):
         """更新账号状态"""
-        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))
+        try:
+            # 获取账号邮箱信息（用于黑名单）
+            account_query = "SELECT email, status as current_status FROM email_accounts WHERE id = %s"
+            account_info = await self.db.fetch_one(account_query, (account_id,))
+            
+            if not account_info:
+                logger.error(f"账号 {account_id} 不存在")
+                return
+                
+            # 检查状态变更是否合理
+            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):
-        """更新账号信息"""
-        query = '''
-            UPDATE email_accounts 
-            SET 
-                cursor_password = %s, 
-                cursor_cookie = %s, 
-                cursor_token = %s,
-                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))
+        """更新账号信息（注册成功）"""
+        try:
+            # 获取账号邮箱信息（用于黑名单）
+            account_query = "SELECT email FROM email_accounts WHERE id = %s"
+            account_info = await self.db.fetch_one(account_query, (account_id,))
+            
+            # 更新数据库
+            query = '''
+                UPDATE email_accounts 
+                SET 
+                    cursor_password = %s, 
+                    cursor_cookie = %s, 
+                    cursor_token = %s,
+                    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):
         """释放账号"""
-        query = '''
-            UPDATE email_accounts 
-            SET in_use = 0, updated_at = CURRENT_TIMESTAMP
-            WHERE id = %s
-        '''
-        await self.db.execute(query, (account_id,))
+        try:
+            await self.unlock_account(account_id)
+            logger.info(f"账号 {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:
         """统计可用的pending状态账号数量"""
+        if self.use_redis:
+            await self._ensure_blacklist_initialized()
+            
         query = """
             SELECT COUNT(*) 
             FROM email_accounts 
             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)
         if result:
             return result.get("COUNT(*)", 0)