commit

backend/main.py

@@ -62,8 +62,11 @@ class ConnectionManager:
     async def connect(self, session_id: str, websocket: WebSocket):
         await websocket.accept()
         self.active_connections[session_id] = websocket
-        print(f"[WebSocket] 新连接: {session_id}", file=sys.stderr)
+        print(f"[WebSocket] ========== 新连接建立 ==========", file=sys.stderr)
+        print(f"[WebSocket] Session ID: {session_id}", file=sys.stderr)
         print(f"[WebSocket] 当前活跃连接数: {len(self.active_connections)}", file=sys.stderr)
+        print(f"[WebSocket] 连接时间: {__import__('datetime').datetime.now()}", file=sys.stderr)
+        print(f"[WebSocket] ===================================", file=sys.stderr)
             
         # 立即检查缓存消息（不等待）
         if session_id in self.pending_messages:
@@ -91,22 +94,42 @@ class ConnectionManager:
         else:
             print(f"[WebSocket] 没有缓存消息: {session_id}", file=sys.stderr)
     
-    def disconnect(self, session_id: str):
+    def disconnect(self, session_id: str, reason: str = "未知原因"):
+        """断开WebSocket连接并记录原因"""
         if session_id in self.active_connections:
             del self.active_connections[session_id]
-            print(f"[WebSocket] 断开连接: {session_id}", file=sys.stderr)
+            print(f"[WebSocket] ========== 连接断开 ==========", file=sys.stderr)
+            print(f"[WebSocket] Session ID: {session_id}", file=sys.stderr)
+            print(f"[WebSocket] 断开原因: {reason}", file=sys.stderr)
+            print(f"[WebSocket] 剩余活跃连接数: {len(self.active_connections)}", file=sys.stderr)
+            print(f"[WebSocket] ===================================", file=sys.stderr)
         # 清理缓存消息
         if session_id in self.pending_messages:
+            pending_count = len(self.pending_messages[session_id])
+            if pending_count > 0:
+                print(f"[WebSocket] 清理 {pending_count} 条未发送的缓存消息", file=sys.stderr)
             del self.pending_messages[session_id]
     
     async def send_message(self, session_id: str, message: dict):
+        print(f"[WebSocket] ========== 尝试发送消息 ==========", file=sys.stderr)
+        print(f"[WebSocket] Session ID: {session_id}", file=sys.stderr)
+        print(f"[WebSocket] 消息类型: {message.get('type')}", file=sys.stderr)
+        print(f"[WebSocket] 当前活跃连接数: {len(self.active_connections)}", file=sys.stderr)
+        print(f"[WebSocket] 活跃连接session_ids: {list(self.active_connections.keys())}", file=sys.stderr)
+        print(f"[WebSocket] session_id在连接中: {session_id in self.active_connections}", file=sys.stderr)
+        print(f"[WebSocket] ===================================", file=sys.stderr)
+        
         if session_id in self.active_connections:
             try:
                 await self.active_connections[session_id].send_json(message)
                 print(f"[WebSocket] 发送消息到 {session_id}: {message.get('type')}", file=sys.stderr)
             except Exception as e:
-                print(f"[WebSocket] 发送消息失败 {session_id}: {str(e)}", file=sys.stderr)
-                self.disconnect(session_id)
+                print(f"[WebSocket] ========== 发送消息失败 ==========", file=sys.stderr)
+                print(f"[WebSocket] Session ID: {session_id}", file=sys.stderr)
+                print(f"[WebSocket] 失败原因: {str(e)}", file=sys.stderr)
+                print(f"[WebSocket] 消息类型: {message.get('type')}", file=sys.stderr)
+                print(f"[WebSocket] ===================================", file=sys.stderr)
+                self.disconnect(session_id, reason=f"发送消息失败: {str(e)}")
         else:
             # WebSocket还未连接，缓存消息
             print(f"[WebSocket] 连接尚未建立，缓存消息: {session_id}", file=sys.stderr)
@@ -811,6 +834,87 @@ async def save_bind_info(request: dict):
             data=None
         )
 
+@app.post("/api/xhs/save-login")
+async def save_login(request: dict):
+    """
+    保存验证码登录的信息到Go后端
+    与扫码登录不同，验证码登录返回的是storage_state数据
+    """
+    try:
+        employee_id = request.get('employee_id')
+        storage_state = request.get('storage_state', {})
+        storage_state_path = request.get('storage_state_path', '')
+        user_info = request.get('user_info', {})  # 新增: 获取用户信息
+        
+        if not employee_id:
+            return BaseResponse(
+                code=1,
+                message="employee_id不能为空",
+                data=None
+            )
+        
+        if not storage_state:
+            return BaseResponse(
+                code=1,
+                message="storage_state不能为空",
+                data=None
+            )
+        
+        # 调用Go后端API保存
+        config = get_config()
+        go_backend_url = config.get_str('go_backend.url', 'http://localhost:8080')
+        
+        # 从 storage_state 中提取 cookies
+        cookies_full = storage_state.get('cookies', [])
+        
+        # 构造保存数据
+        save_data = {
+            "employee_id": employee_id,
+            "cookies_full": cookies_full,
+            "storage_state": storage_state,
+            "storage_state_path": storage_state_path,
+            "user_info": user_info  # 新增: 传递用户信息
+        }
+        
+        print(f"[保存验证码登录] employee_id={employee_id}, cookies数量={len(cookies_full)}, 用户={user_info.get('nickname', '未知')}", file=sys.stderr)
+        
+        import aiohttp
+        async with aiohttp.ClientSession() as session:
+            # 获取小程序传来的token
+            auth_header = request.get('Authorization', '')
+            
+            async with session.post(
+                f"{go_backend_url}/api/xhs/save-login",
+                json=save_data,
+                headers={'Authorization': auth_header} if auth_header else {}
+            ) as resp:
+                result = await resp.json()
+                
+                if resp.status == 200 and result.get('code') == 200:
+                    print(f"[保存验证码登录] 保存成功", file=sys.stderr)
+                    return BaseResponse(
+                        code=0,
+                        message="保存成功",
+                        data=result.get('data')
+                    )
+                else:
+                    print(f"[保存验证码登录] 保存失败: {result.get('message')}", file=sys.stderr)
+                    return BaseResponse(
+                        code=1,
+                        message=result.get('message', '保存失败'),
+                        data=None
+                    )
+    
+    except Exception as e:
+        print(f"[保存验证码登录] 异常: {str(e)}", file=sys.stderr)
+        import traceback
+        traceback.print_exc()
+        return BaseResponse(
+            code=1,
+            message=f"保存失败: {str(e)}",
+            data=None
+        )
+
 @app.post("/api/xhs/qrcode/cancel")
 async def cancel_qrcode_login(request: dict):
     """
@@ -1324,6 +1428,7 @@ async def upload_images(files: List[UploadFile] = File(...)):
 async def handle_send_code_ws(session_id: str, phone: str, country_code: str, login_page: str, websocket: WebSocket):
     """
     异步处理WebSocket发送验证码请求
+    返回: (result, service_instance) - result是结果字典，service_instance是XHSLoginService实例
     """
     try:
         print(f"[WebSocket-SendCode] 开始处理: session={session_id}, phone={phone}", file=sys.stderr)
@@ -1343,6 +1448,13 @@ async def handle_send_code_ws(session_id: str, phone: str, country_code: str, lo
             session_id=session_id
         )
         
+        # 将service实例存储到浏览器池，供后续验证码验证使用
+        if session_id in browser_pool.temp_browsers:
+            browser_pool.temp_browsers[session_id]['service'] = request_login_service
+            print(f"[WebSocket-SendCode] 已存储service实例: {session_id}", file=sys.stderr)
+        else:
+            print(f"[WebSocket-SendCode] 警告: session_id {session_id} 不在temp_browsers中", file=sys.stderr)
+        
         # 检查是否需要验证（发送验证码时触发风控）
         if result.get("need_captcha"):
             print(f"[WebSocket-SendCode] 检测到风控，需要扫码", file=sys.stderr)
@@ -1353,7 +1465,8 @@ async def handle_send_code_ws(session_id: str, phone: str, country_code: str, lo
                 "message": result.get("message", "需要扫码验证")
             })
             print(f"[WebSocket-SendCode] 已推送风控信息", file=sys.stderr)
-            return
+            # 返回service实例，供外部启动监听任务
+            return result, request_login_service
         
         if result["success"]:
             print(f"[WebSocket-SendCode] 验证码发送成功", file=sys.stderr)
@@ -1369,6 +1482,9 @@ async def handle_send_code_ws(session_id: str, phone: str, country_code: str, lo
                 "success": False,
                 "message": result.get("error", "发送验证码失败")
             })
+        
+        return result, request_login_service
+        
     except Exception as e:
         print(f"[WebSocket-SendCode] 异常: {str(e)}", file=sys.stderr)
         import traceback
@@ -1381,6 +1497,7 @@ async def handle_send_code_ws(session_id: str, phone: str, country_code: str, lo
             })
         except:
             pass
+        return {"success": False, "error": str(e)}, None
 
 async def handle_verify_code_ws(session_id: str, phone: str, code: str, country_code: str, login_page: str, websocket: WebSocket):
     """
@@ -1404,7 +1521,7 @@ async def handle_verify_code_ws(session_id: str, phone: str, code: str, country_
         request_login_service = browser_data['service']
         
         # 调用登录服务验证登录
-        result = await request_login_service.login_with_code(
+        result = await request_login_service.login(
             phone=phone,
             code=code,
             country_code=country_code,
@@ -1483,6 +1600,43 @@ async def websocket_login(websocket: WebSocket, session_id: str):
     """
     await ws_manager.connect(session_id, websocket)
     
+    # 启动Redis订阅任务
+    import asyncio
+    import redis.asyncio as aioredis
+    import json
+    from config import get_config
+    
+    config = get_config()
+    redis_host = config.get_str('redis.host', 'localhost')
+    redis_port = config.get_int('redis.port', 6379)
+    redis_password = config.get_str('redis.password', '')
+    
+    # 创建Redis订阅客户端
+    redis_url = f"redis://:{redis_password}@{redis_host}:{redis_port}" if redis_password else f"redis://{redis_host}:{redis_port}"
+    redis_client = await aioredis.from_url(redis_url, decode_responses=True)
+    pubsub = redis_client.pubsub()
+    channel = f"ws_message:{session_id}"
+    await pubsub.subscribe(channel)
+    print(f"[WebSocket] 已订阅Redis频道: {channel}", file=sys.stderr)
+    
+    # 启动后台任务监听Redis消息
+    async def redis_subscriber():
+        try:
+            async for message in pubsub.listen():
+                if message['type'] == 'message':
+                    try:
+                        data = json.loads(message['data'])
+                        print(f"[WebSocket] 从Redis收到消息: {data}", file=sys.stderr)
+                        await websocket.send_json(data)
+                        print(f"[WebSocket] 已转发消息到前端: {session_id}", file=sys.stderr)
+                    except Exception as e:
+                        print(f"[WebSocket] 处理Redis消息失败: {str(e)}", file=sys.stderr)
+        except Exception as e:
+            print(f"[WebSocket] Redis订阅异常: {str(e)}", file=sys.stderr)
+    
+    # 在后台启动Redis监听
+    redis_task = asyncio.create_task(redis_subscriber())
+    
     try:
         # 保持连接，等待消息或断开
         while True:
@@ -1496,7 +1650,6 @@ async def websocket_login(websocket: WebSocket, session_id: str):
             else:
                 # 尝试解析JSON消息
                 try:
-                    import json
                     msg = json.loads(data)
                     msg_type = msg.get('type', 'unknown')
                     print(f"[WebSocket] 解析消息类型: {msg_type}", file=sys.stderr)
@@ -1519,8 +1672,15 @@ async def websocket_login(websocket: WebSocket, session_id: str):
                         login_page = msg.get('login_page', 'creator')
                         print(f"[WebSocket] 收到发送验证码请求: phone={phone}", file=sys.stderr)
                         
-                        # 启动异步任务处理发送验证码
-                        asyncio.create_task(handle_send_code_ws(session_id, phone, country_code, login_page, websocket))
+                        # 直接处理发送验证码（不使用create_task）
+                        result, service_instance = await handle_send_code_ws(session_id, phone, country_code, login_page, websocket)
+                        
+                        # 如果需要扫码，在当前协程中启动监听
+                        if result.get("need_captcha") and service_instance:
+                            print(f"[WebSocket] 在主协程中启动扫码监听: {session_id}", file=sys.stderr)
+                            # 使用create_task在后台监听，但不阻塞当前消息循环
+                            asyncio.create_task(service_instance._monitor_qrcode_scan(session_id))
+                            print(f"[WebSocket] 已启动扫码监听任务", file=sys.stderr)
                     
                     # 处理验证码验证消息
                     elif msg_type == 'verify_code':
@@ -1536,12 +1696,38 @@ async def websocket_login(websocket: WebSocket, session_id: str):
                 except json.JSONDecodeError:
                     print(f"[WebSocket] 无法解析为JSON: {data}", file=sys.stderr)
                 
-    except WebSocketDisconnect:
-        ws_manager.disconnect(session_id)
-        print(f"[WebSocket] 客户端断开: {session_id}", file=sys.stderr)
+    except WebSocketDisconnect as e:
+        reason = f"客户端主动断开连接 (code: {e.code if hasattr(e, 'code') else 'unknown'})"
+        ws_manager.disconnect(session_id, reason=reason)
     except Exception as e:
-        ws_manager.disconnect(session_id)
-        print(f"[WebSocket] 连接异常 {session_id}: {str(e)}", file=sys.stderr)
+        reason = f"连接异常: {type(e).__name__} - {str(e)}"
+        ws_manager.disconnect(session_id, reason=reason)
+    finally:
+        # 清理Redis订阅
+        try:
+            redis_task.cancel()
+            await pubsub.unsubscribe(channel)
+            await pubsub.close()
+            await redis_client.close()
+            print(f"[WebSocket] 已取消Redis订阅: {channel}", file=sys.stderr)
+        except:
+            pass
+        
+        # 释放浏览器实例
+        try:
+            # 检查是否有临时浏览器需要释放
+            if session_id in browser_pool.temp_browsers:
+                print(f"[WebSocket] 检测到未释放的临时浏览器，开始清理: {session_id}", file=sys.stderr)
+                await browser_pool.release_temp_browser(session_id)
+                print(f"[WebSocket] 已释放临时浏览器: {session_id}", file=sys.stderr)
+            
+            # 检查是否有扫码页面需要释放
+            if session_id in browser_pool.qrcode_pages:
+                print(f"[WebSocket] 检测到未释放的扫码页面，开始清理: {session_id}", file=sys.stderr)
+                await browser_pool.release_qrcode_page(session_id)
+                print(f"[WebSocket] 已释放扫码页面: {session_id}", file=sys.stderr)
+        except Exception as e:
+            print(f"[WebSocket] 释放浏览器异常: {str(e)}", file=sys.stderr)
 
 if __name__ == "__main__":
     import uvicorn