import json from channels.generic.websocket import AsyncWebsocketConsumer from channels.db import database_sync_to_async from django.contrib.auth import get_user_model from django.utils import timezone User = get_user_model() class ChatConsumer(AsyncWebsocketConsumer): async def connect(self): self.user = self.scope.get('user') if not self.user or self.user.is_anonymous: token = self.scope.get('query_string', b'').decode() if 'token=' in token: token = token.split('token=')[1].split('&')[0] self.user = await self.get_user_from_token(token) if not self.user or self.user.is_anonymous: await self.close() return self.user_group = f'user_{self.user.id}' await self.channel_layer.group_add(self.user_group, self.channel_name) await self.accept() await self.broadcast_presence(True) async def disconnect(self, close_code): if hasattr(self, 'user_group'): await self.broadcast_presence(False) await self.channel_layer.group_discard(self.user_group, self.channel_name) async def receive(self, text_data): try: data = json.loads(text_data) except json.JSONDecodeError: return msg_type = data.get('type', '') if msg_type == 'chat_message': await self.handle_chat_message(data) elif msg_type == 'typing': await self.handle_typing(data) elif msg_type == 'read_receipt': await self.handle_read_receipt(data) async def handle_chat_message(self, data): conversation_id = data.get('conversation_id') content = data.get('content', '').strip() msg_type = data.get('msg_type', 'text') file_url = data.get('file_url', '') reply_to_id = data.get('reply_to') if not content or not conversation_id: return # 检查是否被对方拉黑 is_blocked = await self.check_if_blocked(conversation_id) if is_blocked: # 发送错误提示给发送者 await self.send(text_data=json.dumps({ 'type': 'error', 'code': 'BLOCKED_BY_USER', 'message': '你已被对方拉黑，无法发送消息', })) return message = await self.save_message(conversation_id, content, msg_type, file_url, reply_to_id) if not message: return participants = await self.get_participants(conversation_id) message_data = await self.serialize_message(message) for participant_id in participants: await self.channel_layer.group_send( f'user_{participant_id}', { 'type': 'chat.message', 'conversation_id': conversation_id, 'message': message_data, } ) async def handle_typing(self, data): conversation_id = data.get('conversation_id') if not conversation_id: return participants = await self.get_participants(conversation_id) for participant_id in participants: if participant_id != self.user.id: await self.channel_layer.group_send( f'user_{participant_id}', { 'type': 'typing.indicator', 'conversation_id': conversation_id, 'user_id': self.user.id, } ) async def handle_read_receipt(self, data): conversation_id = data.get('conversation_id') if not conversation_id: return await self.update_last_read(conversation_id) participants = await self.get_participants(conversation_id) for participant_id in participants: if participant_id != self.user.id: await self.channel_layer.group_send( f'user_{participant_id}', { 'type': 'read.receipt', 'conversation_id': conversation_id, 'user_id': self.user.id, } ) async def broadcast_presence(self, online): friend_ids = await self.get_friend_ids() for friend_id in friend_ids: await self.channel_layer.group_send( f'user_{friend_id}', { 'type': 'presence.update', 'user_id': self.user.id, 'online': online, } ) async def chat_message(self, event): await self.send(text_data=json.dumps({ 'type': 'chat_message', 'conversation_id': event['conversation_id'], 'message': event['message'], })) async def typing_indicator(self, event): await self.send(text_data=json.dumps({ 'type': 'typing', 'conversation_id': event['conversation_id'], 'user_id': event['user_id'], })) async def read_receipt(self, event): await self.send(text_data=json.dumps({ 'type': 'read_receipt', 'conversation_id': event['conversation_id'], 'user_id': event['user_id'], })) async def presence_update(self, event): await self.send(text_data=json.dumps({ 'type': 'presence', 'user_id': event['user_id'], 'online': event['online'], })) @database_sync_to_async def get_user_from_token(self, token): from rest_framework_simplejwt.tokens import AccessToken try: access_token = AccessToken(token) user_id = access_token['user_id'] return User.objects.get(id=user_id) except Exception: return None @database_sync_to_async def save_message(self, conversation_id, content, msg_type, file_url, reply_to_id): from .models import Conversation, ConversationParticipant, ChatMessage try: conversation = Conversation.objects.get(pk=conversation_id) if not ConversationParticipant.objects.filter(conversation=conversation, user=self.user).exists(): return None reply_to = None if reply_to_id: try: reply_to = ChatMessage.objects.get(pk=reply_to_id, conversation=conversation) except ChatMessage.DoesNotExist: pass return ChatMessage.objects.create( conversation=conversation, sender=self.user, content=content, msg_type=msg_type, file_url=file_url, reply_to=reply_to, ) except Exception: return None @database_sync_to_async def get_participants(self, conversation_id): from .models import ConversationParticipant return list(ConversationParticipant.objects.filter( conversation_id=conversation_id ).values_list('user_id', flat=True)) @database_sync_to_async def serialize_message(self, message): from .serializers import ChatMessageSerializer return ChatMessageSerializer(message).data @database_sync_to_async def get_friend_ids(self): from .models import Friendship from django.db.models import Q friendships = Friendship.objects.filter( Q(user1=self.user) | Q(user2=self.user) ) return [f.user2_id if f.user1 == self.user else f.user1_id for f in friendships] @database_sync_to_async def update_last_read(self, conversation_id): from .models import ConversationParticipant ConversationParticipant.objects.filter( conversation_id=conversation_id, user=self.user ).update(last_read_at=timezone.now()) @database_sync_to_async def check_if_blocked(self, conversation_id): """检查当前用户是否被会话中的其他用户拉黑""" from .models import ConversationParticipant from user.models import Blacklist participants = ConversationParticipant.objects.filter( conversation_id=conversation_id ).exclude(user=self.user).values_list('user_id', flat=True) return Blacklist.objects.filter( user_id__in=participants, blocked_user=self.user ).exists()