145 lines
5.1 KiB
Python
145 lines
5.1 KiB
Python
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
|
||
import json
|
||
from sqlalchemy.orm import Session
|
||
from app.db import models
|
||
from firebase_admin import messaging, credentials, exceptions
|
||
import firebase_admin
|
||
|
||
cred = credentials.Certificate("chepuhagram-6ca5d-firebase-adminsdk-fbsvc-cf8a5ad2f3.json")
|
||
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 = 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("ПОДКЛЮЧЕНО")
|
||
try:
|
||
while True:
|
||
print("ОЖИДАНИЕ СООБЩЕНИЙ")
|
||
data = await websocket.receive_text()
|
||
message_data = json.loads(data)
|
||
print(f"DEBUG: Получены данные: {message_data}")
|
||
|
||
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")
|
||
content = message_data.get("content")
|
||
content50 = message_data.get("content50")
|
||
new_msg = models.Message(
|
||
sender_id=user_id,
|
||
receiver_id=receiver_id,
|
||
content=content
|
||
)
|
||
db.add(new_msg)
|
||
db.commit()
|
||
db.refresh(new_msg)
|
||
|
||
if receiver_id not in manager.active_connections and 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,
|
||
"reciever_id": receiver_id,
|
||
"content": message_data.get("content"),
|
||
"timestamp": datetime.now().isoformat()
|
||
}
|
||
|
||
# Пересылаем получателю, если он в сети
|
||
await manager.send_personal_message(outgoing_message, str(receiver_id))
|
||
except WebSocketDisconnect:
|
||
pass
|
||
finally:
|
||
manager.disconnect(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: str, websocket: WebSocket):
|
||
await websocket.accept()
|
||
self.active_connections[user_id] = websocket
|
||
|
||
def disconnect(self, user_id: str):
|
||
if user_id in self.active_connections:
|
||
del self.active_connections[user_id]
|
||
|
||
async def send_personal_message(self, message: dict, user_id: str):
|
||
if str(user_id) in self.active_connections:
|
||
await self.active_connections[str(user_id)].send_json(message)
|
||
print('Sent to socket')
|
||
else:
|
||
print('User not active')
|
||
|
||
async def broadcast(self, message: dict):
|
||
# Рассылка вообще всем (например, системное уведомление)
|
||
for connection in self.active_connections.values():
|
||
await connection.send_json(message)
|
||
|
||
|
||
manager = ConnectionManager()
|