from fastapi import HTTPException, status, APIRouter, WebSocket, WebSocketDisconnect, Query, Depends from app.core.security import test_token from typing import Dict from datetime import datetime, timezone import json from sqlalchemy.orm import Session from app.db import models from firebase_admin import messaging, credentials, exceptions import firebase_admin from app.core.config import config cred = credentials.Certificate(config.FIREBASE_CREDENTIALS_PATH) firebase_admin.initialize_app(cred) # бд def get_db(): db = models.SessionLocal() try: yield db finally: db.close() wsRouter = APIRouter( prefix='/ws' ) @wsRouter.websocket("") async def websocket_endpoint(websocket: WebSocket, token: str = Query(None), db: Session = Depends(get_db)): if token is None: await websocket.close(code=status.WS_1008_POLICY_VIOLATION) return try: user_id = int(await test_token(token=token)) except HTTPException: await websocket.close(code=status.WS_1008_POLICY_VIOLATION) return print("ПОДКЛЮЧЕНИЕ") await manager.connect(user_id, websocket) print("ПОДКЛЮЧЕНО") db.query(models.User).filter(models.User.id == user_id).update({"last_online": datetime.now(timezone.utc)}, synchronize_session="fetch") db.commit() await manager.broadcast({ "type": "user_online", "user_id": user_id, }) try: while True: print("ОЖИДАНИЕ СООБЩЕНИЙ") data = await websocket.receive_text() message_data = json.loads(data) print(f"DEBUG: Получены данные: {message_data}") db.query(models.User).filter(models.User.id == user_id).update({"last_online": datetime.now(timezone.utc)}, synchronize_session="fetch") db.commit() if message_data.get("type") == "private_message": user = db.query(models.User).filter( models.User.id == user_id).first() receiver_id = message_data.get("receiver_id") temp_id = message_data.get("temp_id") content = message_data.get("content") content50 = message_data.get("content50") if receiver_id is None or content is None: await websocket.send_json({ "type": "error", "detail": "receiver_id/content required", }) continue try: receiver_id = int(receiver_id) except (TypeError, ValueError): await websocket.send_json({ "type": "error", "detail": "receiver_id must be int", }) continue new_msg = models.Message( sender_id=user_id, receiver_id=receiver_id, content=content, reply_to_id=message_data.get("reply_to_id"), reply_to_text=message_data.get("reply_to_text") ) db.add(new_msg) db.commit() db.refresh(new_msg) # ACK отправителю: сервер принял и сохранил сообщение (нужно для статусов клиента). await manager.send_personal_message({ "type": "message_sent", "temp_id": temp_id, "server_id": new_msg.id, "timestamp": (new_msg.timestamp or datetime.now()).isoformat(), }, str(user_id)) # Если получатель оффлайн — отправим пуш (если есть токен и ключи). if user.public_key: receiver = db.query(models.User).filter( models.User.id == receiver_id).first() if receiver.fcm_token: send_fcm_notification( receiver.fcm_token, user_id, user.first_name, user.public_key, content50 if content50 else content, datetime.now(), ) # Формируем пакет для получателя outgoing_message = { "id": new_msg.id, "type": "private_message", "sender_id": user_id, "receiver_id": receiver_id, "content": message_data.get("content"), "timestamp": (new_msg.timestamp or datetime.now()).isoformat(), "reply_to_id": new_msg.reply_to_id, "reply_to_text": new_msg.reply_to_text, } # Пересылаем получателю, если он в сети sent_to_receiver = await manager.send_personal_message(outgoing_message, str(receiver_id)) # Если сообщение реально ушло по сокету получателю — отмечаем delivered_at. if sent_to_receiver: try: delivered_at = datetime.now() new_msg.delivered_at = delivered_at db.add(new_msg) db.commit() await manager.send_personal_message({ "type": "message_delivered", "message_id": new_msg.id, "timestamp": delivered_at.isoformat(), }, str(user_id)) except Exception: db.rollback() elif message_data.get("type") == "edit_message": message_id = message_data.get("message_id") content = message_data.get("content") if message_id is None or content is None: await websocket.send_json({ "type": "error", "detail": "message_id/content required", }) continue try: message_id = int(message_id) except (TypeError, ValueError): await websocket.send_json({ "type": "error", "detail": "message_id must be int", }) continue msg = db.query(models.Message).filter( models.Message.id == message_id).first() if msg is None or msg.sender_id != user_id: continue try: msg.content = content msg.edited_at = datetime.now() db.add(msg) db.commit() except Exception: db.rollback() continue event = { "type": "message_edited", "message_id": msg.id, "content": msg.content, "edited_at": msg.edited_at.isoformat() if msg.edited_at else None, } await manager.send_personal_message(event, str(msg.receiver_id)) await manager.send_personal_message(event, str(msg.sender_id)) elif message_data.get("type") == "delete_message": message_id = message_data.get("message_id") if message_id is None: await websocket.send_json({ "type": "error", "detail": "message_id required", }) continue try: message_id = int(message_id) except (TypeError, ValueError): await websocket.send_json({ "type": "error", "detail": "message_id must be int", }) continue msg = db.query(models.Message).filter( models.Message.id == message_id).first() if msg is None or msg.sender_id != user_id: continue receiver_id = msg.receiver_id try: db.delete(msg) db.commit() except Exception: db.rollback() continue event = { "type": "message_deleted", "message_id": message_id, } await manager.send_personal_message(event, str(receiver_id)) await manager.send_personal_message(event, str(user_id)) elif message_data.get("type") == "read_receipt": message_id = message_data.get("message_id") try: message_id = int(message_id) except (TypeError, ValueError): continue msg = db.query(models.Message).filter( models.Message.id == message_id).first() if msg is None: continue # Безопасность: read_receipt может отправлять только получатель. if int(msg.receiver_id) != int(user_id): continue # Сохраняем read_at в БД try: read_at = datetime.now() msg.read_at = read_at db.add(msg) db.commit() except Exception: db.rollback() sender_id = int(msg.sender_id) await manager.send_personal_message({ "type": "message_read", "message_id": message_id, "timestamp": read_at.isoformat() if 'read_at' in locals() else datetime.now().isoformat(), }, str(sender_id)) elif message_data.get("type") == "typing": receiver_id = message_data.get("receiver_id") if receiver_id is None: continue try: receiver_id = int(receiver_id) except (TypeError, ValueError): continue await manager.send_personal_message({ "type": "typing", "sender_id": user_id, }, str(receiver_id)) elif message_data.get("type") == "stop_typing": receiver_id = message_data.get("receiver_id") if receiver_id is None: continue try: receiver_id = int(receiver_id) except (TypeError, ValueError): continue await manager.send_personal_message({ "type": "stop_typing", "sender_id": user_id, }, str(receiver_id)) except WebSocketDisconnect: pass finally: manager.disconnect(user_id) db.query(models.User).filter(models.User.id == user_id).update( {"last_online": datetime.now(timezone.utc)}, synchronize_session="fetch") db.commit() print("ОТКЛЮЧЕНИЕ") await manager.broadcast({ "type": "user_offline", "user_id": user_id, }) def send_fcm_notification(token, user_id, username, public_key, encrypted_text, timestamp): print( f"DEBUG: Отправляем FCM уведомление пользователю {user_id} с токеном {token}") message = messaging.Message( data={ "type": "enc_message", "sender_id": str(user_id), "username": username, "public_key": public_key, "content": encrypted_text, # Зашифрованный текст "timestamp": timestamp.isoformat(), }, android=messaging.AndroidConfig( priority='high', ), token=token, ) try: response = messaging.send(message) print('Successfully sent message:', response) except Exception as e: print('Unexpected error sending push:', e) class ConnectionManager: def __init__(self): # Храним активные соединения: {user_id: websocket} self.active_connections: Dict[str, WebSocket] = {} async def connect(self, user_id: int, websocket: WebSocket): await websocket.accept() self.active_connections[str(user_id)] = websocket def disconnect(self, user_id: int): key = str(user_id) if key in self.active_connections: del self.active_connections[key] async def send_personal_message(self, message: dict, user_id: str) -> bool: if str(user_id) in self.active_connections: try: await self.active_connections[str(user_id)].send_json(message) print('Sent to socket') return True except Exception as e: print(f'Failed to send to socket: {e}') return False else: print('User not active') return False async def broadcast(self, message: dict): # Рассылка вообще всем (например, системное уведомление) for connection in self.active_connections.values(): await connection.send_json(message) manager = ConnectionManager()