基于Django 实现 WebSocket 中转功能

大致需求是这样：

前端服务A，后端服务B（Django Channels），第三方服务C。

历史情况为 A 与 C 通过 socket 直连。

现在想通过B做socket中转，方便审计工作。

废话太多，直接上服务B代码。
 

import asyncio
import json
from channels.generic.websocket import AsyncJsonWebsocketConsumer
import websockets


class WebTerminalConsumer(AsyncJsonWebsocketConsumer):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.opened = False
        self.c_socket = None
        self.task = None

    async def connect(self):
        await self.accept()

    # Receive message from WebSocket
    async def receive_json(self, content, **kwargs):
        """
        content 格式大致如下
        content = {
            "action": "INIT",
            "params": {...}
        }
        """
        action = content.get('action')
        params = content.get('params', {})
        if action == 'INIT':
            # 初始化链接
            server_url = 'wss://your_target_ws'

            self.c_socket = await websockets.connect(server_url)
            response = await self.c_socket.recv()
            if response == 'o':

                self.opened = True
            
            # 创建一个异步处理消息的任务
            self.task = asyncio.create_task(self.message_handler())

        elif action == 'STDIN':
            # 客户端在前端输入，直接转发至 C服务
            """
            可以在此处做一些审计操作
            """
            await self.c_socket.send(json.dumps([json.dumps(params)]))

    async def message_handler(self):
        # 在这里处理C服务的消息
        while self.opened:
            response = await self.c_socket.recv()
            await self.send(response)

    async def disconnect(self, close_code):
        # 这里先关闭启动的异步处理消息的任务
        self.task.cancel()
        # 关闭和C服务的连接
        await self.c_socket.close()
        await self.close()

 

有问题欢迎私信。