即时通讯工具（基础演示框架）

服务端

# coding=utf-8

import asyncio

import bcrypt

import jwt

import json

import websockets

from sqlalchemy.ext.declarative import declarative_base

from sqlalchemy import create_engine, Column, Integer, String, DateTime

from sqlalchemy.orm import sessionmaker, relationship

 

Base = declarative_base()

 

# 用户模型

class User(Base):

    __tablename__ = 'users'

    id = Column(Integer, primary_key=True)

    username = Column(String, unique=True)

    password_hash = Column(String)

 

# 消息模型

class Message(Base):

    __tablename__ = 'messages'

    id = Column(Integer, primary_key=True)

    sender_id = Column(Integer)

    receiver_id = Column(Integer)

    content = Column(String)

    timestamp = Column(DateTime)

 

# 数据库配置

DATABASE_URL = 'sqlite:///chat.db'

engine = create_engine(DATABASE_URL)

Base.metadata.create_all(engine)

Session = sessionmaker(bind=engine)

session = Session()

 

# WebSocket通信处理

connected_clients = {}

 

async def handle_client(websocket, path):

    global connected_clients

    

    # 从WebSocket获取认证信息并验证

    try:

        auth_data = await websocket.recv()

        auth_info = json.loads(auth_data)

        username = auth_info['username']

        password = auth_info['password']

        

        # 进行数据库查询和用户认证

        user = session.query(User).filter_by(username=username).first()

        if user and bcrypt.checkpw(password.encode('utf-8'), user.password_hash.encode('utf-8')):

            connected_clients[username] = websocket

            print(f"{username} connected")

            # 处理客户端消息

            async for message_data in websocket:

                message_info = json.loads(message_data)

                await forward_message(message_info)

        else:

            raise Exception("Authentication failed")

    except websockets.ConnectionClosed:

        print(f"{username} connection closed")

    finally:

        if username in connected_clients:

            del connected_clients[username]

 

# 转发消息到接收方

async def forward_message(message_info):

    receiver_ws = connected_clients.get(message_info['receiver_username'])

    if receiver_ws:

        await receiver_ws.send(message_info['content'])

 

# 启动WebSocket服务器

async def main():

    ssl_context = None # 需要创建SSL上下文来支持加密通信

    async with websockets.serve(handle_client, "localhost", 6789, ssl=ssl_context):

        await asyncio.Future() # 运行直到Interrupt

 

# 错误处理和日志记录省略

 

# 运行服务器

if __name__ == '__main__':

    asyncio.run(main())

---

客户端

import asyncio

import websockets

import json

 

async def chat_client():

    uri = "ws://localhost:6789"

    

    username = input("Enter your username: ")

    password = input("Enter your password: ")

 

    async with websockets.connect(uri) as websocket:

        # 发送认证信息

        auth_info = json.dumps({'username': username, 'password': password})

        await websocket.send(auth_info)

        

        # 处理接收和发送消息

        send_task = asyncio.ensure_future(send_message(websocket, username))

        receive_task = asyncio.ensure_future(receive_message(websocket))

        done, pending = await asyncio.wait(

            [send_task, receive_task],

            return_when=asyncio.FIRST_COMPLETED,

        )

        for task in pending:

            task.cancel()

 

async def send_message(websocket, username):

    try:

        while True:

            message = input("Enter your message (or 'quit' to exit): ")

            if message.lower() == 'quit':

                break

 

            receiver_username = input("Enter the receiver's username: ")

            # 发送消息

            message_info = json.dumps({

                'receiver_username': receiver_username,

                'content': message

            })

            await websocket.send(message_info)

    except websockets.ConnectionClosed:

        pass

 

async def receive_message(websocket):

    try:

        while True:

            message_data = await websocket.recv()

            print(f"\nReceived message: {message_data}")

    except websockets.ConnectionClosed:

        pass

 

if __name__ == "__main__":

    asyncio.get_event_loop().run_until_complete(chat_client())