Tornado实现本地代理

• http

• websocket

import json
import uuid

import tornado.httpserver
import tornado.ioloop
import tornado.web
import tornado.websocket


class MainHandler(tornado.web.RequestHandler):
    def get(self):
        # 获取 HTTP GET 请求的参数
        params = {
            "username": self.get_argument("username", default=""),
            "account": self.get_argument("account", default=""),
        }

        # 获取 WebSocketHandler 实例列表
        ws_handler_map = self.application.settings["ws_handler_map"]

        # 发送消息到所有 WebSocket 客户端
        ws_handler = ws_handler_map.get("websocket_id")
        if ws_handler:
            ws_handler.send_message(json.dumps(params))

        # 写入 HTTP 响应
        self.write(json.dumps(params))


class WebSocketHandler(tornado.websocket.WebSocketHandler):
    def initialize(self):
        # 初始化一个全局的 WebSocketHandler 实例列表
        if "ws_handlers" not in self.application.settings:
            self.application.settings["ws_handlers"] = []
        self.application.settings["ws_handlers"].append(self)

    def check_origin(self, origin):
        return True

    def open(self):
        # 关闭旧的连接
        ws_handler_map = self.application.settings.setdefault("ws_handler_map", {})
        if ws_handler_map.get("websocket_id"):
            ws_handler_map["websocket_id"].close()

        self.client_id = str(uuid.uuid4())
        ws_handler_map["websocket_id"] = self

        print(f"WebSocket Opened for client_id: {self.client_id}")

    def on_message(self, message):
        # 在收到消息时，向客户端发送响应
        try:
            data = tornado.escape.json_decode(message)
            user_id = data["user_id"]
            content = data["content"]
            response = f"Hello, {user_id}! You sent: {content}"
            self.write_message(response)
        except json.JSONDecodeError:
            self.write_message("Error: Invalid JSON format")

    def send_message(self, message):
        self.write_message(message)

    def on_close(self):
        ws_handler_map = self.application.settings.get("ws_handler_map", {})
        ws_handler_map.pop("websocket_id", None)
        print(f"WebSocket Closed for client_id: {self.client_id}")


def make_app():
    return tornado.web.Application(
        [
            (r'/', MainHandler),
            (r'/websocket', WebSocketHandler),
        ],
        debug=True,  # 可以在调试时启用
    )


if __name__ == "__main__":
    try:
        app = make_app()
        server = tornado.httpserver.HTTPServer(app)
        server.listen(9000)

        # HTTPS server can be configured with SSL options if needed
        ws_server = tornado.httpserver.HTTPServer(app)
        ws_server.add_sockets(tornado.netutil.bind_sockets(9001))

        tornado.ioloop.IOLoop.current().start()
    except KeyboardInterrupt:
        # 捕获 Ctrl+C 中断信号，执行清理逻辑
        tornado.ioloop.IOLoop.current().stop()
    except Exception as e:
        print(f"Unexpected error: {e}")