mirror of
https://github.com/fromchat-messenger/web.git
synced 2026-09-22 19:15:08 +03:00
325 lines
11 KiB
Python
325 lines
11 KiB
Python
from datetime import datetime
|
|
import logging
|
|
from fastapi import APIRouter, Depends, HTTPException, WebSocket, WebSocketDisconnect
|
|
from fastapi.security import HTTPAuthorizationCredentials
|
|
from sqlalchemy.orm import Session
|
|
from dependencies import get_current_user, get_db
|
|
from constants import OWNER_USERNAME
|
|
from models import Message, SendMessageRequest, EditMessageRequest, ReplyMessageRequest, User, DMEnvelope
|
|
|
|
router = APIRouter()
|
|
logger = logging.getLogger("uvicorn.error")
|
|
|
|
def convert_message(msg: Message) -> dict:
|
|
return {
|
|
"id": msg.id,
|
|
"content": msg.content,
|
|
"timestamp": msg.timestamp.isoformat(),
|
|
"is_read": msg.is_read,
|
|
"is_edited": msg.is_edited,
|
|
"username": msg.author.username,
|
|
"profile_picture": msg.author.profile_picture,
|
|
"reply_to": convert_message(msg.reply_to) if msg.reply_to else None
|
|
}
|
|
|
|
|
|
@router.post("/send_message")
|
|
async def send_message(
|
|
request: SendMessageRequest,
|
|
current_user: User = Depends(get_current_user),
|
|
db: Session = Depends(get_db)
|
|
):
|
|
if not request.content.strip():
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail="No content provided"
|
|
)
|
|
|
|
if len(request.content.strip()) > 4096:
|
|
raise HTTPException(
|
|
status_code=400,
|
|
detail="Message too long"
|
|
)
|
|
|
|
new_message = Message(
|
|
content=request.content.strip(),
|
|
user_id=current_user.id,
|
|
timestamp=datetime.now()
|
|
)
|
|
|
|
db.add(new_message)
|
|
db.commit()
|
|
db.refresh(new_message)
|
|
|
|
return {"status": "success", "message": convert_message(new_message)}
|
|
|
|
|
|
@router.get("/get_messages")
|
|
async def get_messages(db: Session = Depends(get_db)):
|
|
messages = db.query(Message).order_by(Message.timestamp.asc()).all()
|
|
|
|
messages_data = []
|
|
for msg in messages:
|
|
messages_data.append(convert_message(msg))
|
|
|
|
return {
|
|
"status": "success",
|
|
"messages": messages_data
|
|
}
|
|
|
|
|
|
@router.post("/dm/send")
|
|
async def dm_send(payload: dict, current_user: User = Depends(get_current_user), db: Session = Depends(get_db)):
|
|
required = ["recipientId", "iv", "ciphertext", "salt", "iv2", "wrappedMk"]
|
|
for key in required:
|
|
if key not in payload:
|
|
raise HTTPException(status_code=400, detail=f"Missing {key}")
|
|
env = DMEnvelope(
|
|
sender_id=current_user.id,
|
|
recipient_id=int(payload["recipientId"]),
|
|
iv_b64=payload["iv"],
|
|
ciphertext_b64=payload["ciphertext"],
|
|
salt_b64=payload["salt"],
|
|
iv2_b64=payload["iv2"],
|
|
wrapped_mk_b64=payload["wrappedMk"],
|
|
)
|
|
db.add(env)
|
|
db.commit()
|
|
db.refresh(env)
|
|
return {"status": "ok", "id": env.id}
|
|
|
|
|
|
@router.get("/dm/fetch")
|
|
async def dm_fetch(since: int | None = None, current_user: User = Depends(get_current_user), db: Session = Depends(get_db)):
|
|
q = db.query(DMEnvelope).filter(DMEnvelope.recipient_id == current_user.id)
|
|
if since:
|
|
q = q.filter(DMEnvelope.id > since)
|
|
envs = q.order_by(DMEnvelope.id.asc()).all()
|
|
return {
|
|
"status": "ok",
|
|
"messages": [
|
|
{
|
|
"id": e.id,
|
|
"senderId": e.sender_id,
|
|
"recipientId": e.recipient_id,
|
|
"iv": e.iv_b64,
|
|
"ciphertext": e.ciphertext_b64,
|
|
"salt": e.salt_b64,
|
|
"iv2": e.iv2_b64,
|
|
"wrappedMk": e.wrapped_mk_b64,
|
|
"timestamp": e.timestamp.isoformat(),
|
|
}
|
|
for e in envs
|
|
]
|
|
}
|
|
|
|
|
|
@router.put("/edit_message/{message_id}")
|
|
async def edit_message(
|
|
message_id: int,
|
|
request: EditMessageRequest,
|
|
current_user: User = Depends(get_current_user),
|
|
db: Session = Depends(get_db)
|
|
):
|
|
message = db.query(Message).filter(Message.id == message_id).first()
|
|
|
|
if not message:
|
|
raise HTTPException(status_code=404, detail="Message not found")
|
|
|
|
if message.user_id != current_user.id:
|
|
raise HTTPException(status_code=403, detail="You can only edit your own messages")
|
|
|
|
if not request.content.strip():
|
|
raise HTTPException(status_code=400, detail="Message content cannot be empty")
|
|
|
|
message.content = request.content.strip()
|
|
message.is_edited = True
|
|
|
|
db.commit()
|
|
db.refresh(message)
|
|
|
|
return {"status": "success", "message": convert_message(message)}
|
|
|
|
|
|
@router.delete("/delete_message/{message_id}")
|
|
async def delete_message(
|
|
message_id: int,
|
|
current_user: User = Depends(get_current_user),
|
|
db: Session = Depends(get_db)
|
|
):
|
|
message = db.query(Message).filter(Message.id == message_id).first()
|
|
|
|
if not message:
|
|
raise HTTPException(status_code=404, detail="Message not found")
|
|
|
|
# Allow owner to delete any message
|
|
if current_user.username != OWNER_USERNAME and message.user_id != current_user.id:
|
|
raise HTTPException(status_code=403, detail="You can only delete your own messages")
|
|
|
|
db.delete(message)
|
|
db.commit()
|
|
|
|
return {"status": "success", "message_id": message_id}
|
|
|
|
|
|
@router.post("/reply_message")
|
|
async def reply_message(
|
|
request: ReplyMessageRequest,
|
|
current_user: User = Depends(get_current_user),
|
|
db: Session = Depends(get_db)
|
|
):
|
|
# Check if the message being replied to exists
|
|
original_message = db.query(Message).filter(Message.id == request.reply_to_id).first()
|
|
if not original_message:
|
|
raise HTTPException(status_code=404, detail="Original message not found")
|
|
|
|
if not request.content.strip():
|
|
raise HTTPException(status_code=400, detail="No content provided")
|
|
|
|
new_message = Message(
|
|
content=request.content.strip(),
|
|
user_id=current_user.id,
|
|
timestamp=datetime.now(),
|
|
reply_to_id=request.reply_to_id
|
|
)
|
|
|
|
db.add(new_message)
|
|
db.commit()
|
|
db.refresh(new_message)
|
|
|
|
return {"status": "success", "message": convert_message(new_message)}
|
|
|
|
|
|
class MessaggingSocketManager:
|
|
def __init__(self) -> None:
|
|
self.connections: list[WebSocket] = []
|
|
|
|
async def send_error(self, websocket: WebSocket, type: str, e: HTTPException):
|
|
await websocket.send_json({"type": type, "error": {"code": e.status_code, "detail": e.detail}})
|
|
|
|
async def handle_connection(self, websocket: WebSocket, db: Session):
|
|
while True:
|
|
data = await websocket.receive_json()
|
|
type = data["type"]
|
|
|
|
def get_current_user_inner() -> User | None:
|
|
if data["credentials"]:
|
|
return get_current_user(
|
|
HTTPAuthorizationCredentials(
|
|
scheme=data["credentials"]["scheme"],
|
|
credentials=data["credentials"]["credentials"]
|
|
),
|
|
db
|
|
)
|
|
else:
|
|
return None
|
|
|
|
if type == "ping":
|
|
await websocket.send_json({"type": "ping", "data": {"status": "success"}})
|
|
elif type == "getMessages":
|
|
try:
|
|
current_user = get_current_user_inner()
|
|
if not current_user:
|
|
raise HTTPException(401)
|
|
|
|
await websocket.send_json({"type": type, "data": await get_messages(current_user, db)})
|
|
except HTTPException as e:
|
|
await self.send_error(websocket, type, e)
|
|
elif type == "sendMessage":
|
|
try:
|
|
current_user = get_current_user_inner()
|
|
if not current_user:
|
|
raise HTTPException(401)
|
|
|
|
request: SendMessageRequest = SendMessageRequest.model_validate(data["data"])
|
|
|
|
response = await send_message(request, current_user, db)
|
|
await self.broadcast({
|
|
"type": "newMessage",
|
|
"data": response["message"]
|
|
})
|
|
|
|
await websocket.send_json({"type": type, "data": response})
|
|
except HTTPException as e:
|
|
await self.send_error(websocket, type, e)
|
|
elif type == "editMessage":
|
|
try:
|
|
current_user = get_current_user_inner()
|
|
if not current_user:
|
|
raise HTTPException(401)
|
|
|
|
message_id = data["data"]["message_id"]
|
|
request: EditMessageRequest = EditMessageRequest.model_validate(data["data"])
|
|
|
|
response = await edit_message(message_id, request, current_user, db)
|
|
await self.broadcast({
|
|
"type": "messageEdited",
|
|
"data": response["message"]
|
|
})
|
|
|
|
await websocket.send_json({"type": type, "data": response})
|
|
except HTTPException as e:
|
|
await self.send_error(websocket, type, e)
|
|
elif type == "deleteMessage":
|
|
try:
|
|
current_user = get_current_user_inner()
|
|
if not current_user:
|
|
raise HTTPException(401)
|
|
|
|
message_id = data["data"]["message_id"]
|
|
response = await delete_message(message_id, current_user, db)
|
|
await self.broadcast({
|
|
"type": "messageDeleted",
|
|
"data": {"message_id": message_id}
|
|
})
|
|
|
|
await websocket.send_json({"type": type, "data": response})
|
|
except HTTPException as e:
|
|
await self.send_error(websocket, type, e)
|
|
elif type == "replyMessage":
|
|
try:
|
|
current_user = get_current_user_inner()
|
|
if not current_user:
|
|
raise HTTPException(401)
|
|
|
|
request: ReplyMessageRequest = ReplyMessageRequest.model_validate(data["data"])
|
|
response = await reply_message(request, current_user, db)
|
|
await self.broadcast({
|
|
"type": "newMessage",
|
|
"data": response["message"]
|
|
})
|
|
|
|
await websocket.send_json({"type": type, "data": response})
|
|
except HTTPException as e:
|
|
await self.send_error(websocket, type, e)
|
|
else:
|
|
await websocket.send_json({"type": type, "error": {"code": 400, "detail": "Invalid type"}})
|
|
|
|
async def disconnect(self, websocket: WebSocket, code: int = 1000, message: str | None = None):
|
|
try:
|
|
await websocket.close(code=code, reason=message)
|
|
finally:
|
|
self.connections.remove(websocket)
|
|
|
|
async def connect(self, websocket: WebSocket, db: Session):
|
|
await websocket.accept()
|
|
self.connections.append(websocket)
|
|
try:
|
|
await self.handle_connection(websocket, db)
|
|
except WebSocketDisconnect as e:
|
|
logger.info(f"WebSocket disconnected with code {e.code}: {e.reason}")
|
|
finally:
|
|
self.connections.remove(websocket)
|
|
|
|
async def broadcast(self, message: dict):
|
|
for websocket in self.connections:
|
|
await websocket.send_json(message)
|
|
|
|
messagingManager = MessaggingSocketManager()
|
|
|
|
@router.websocket("/chat/ws")
|
|
async def chat_websocket(
|
|
websocket: WebSocket,
|
|
db: Session = Depends(get_db)
|
|
):
|
|
await messagingManager.connect(websocket, db) |