Python 封装 socket 为 [TCP/UDP/MULTICAST] 服务端

在新线程中创建 TCP/UDP/MULTICAST 协议的服务端套接字，接收客户端的连接请求或数据，并调用 on_recv 回调函数处理数据。

#!/usr/bin/env python
# -*- coding: utf-8 -*-
import socket
import threading


class ServerSocket:
    def __init__(self, *, protocol: str, port: int, group: str = '', on_recv: 'function'):
        """服务端套接字

        在新线程中创建 TCP/UDP/MULTICAST 协议的服务端套接字，接收客户端的
        连接请求或数据，并调用 on_recv 回调函数处理数据。

        TCP 断开连接的情况：
        
        - TCP 正常断开
            + 客户端主动断开连接
            + 通信期间正常交换数据 (若服务端返回了响应, 则客户端应该接收响应)
        - TCP 连接已重置
            + 客户端主动断开连接
            + 服务端返回了响应，但客户端未接收
        - TCP 连接已终止
            + 未通信完毕就已经断开了连接

        Args:
            protocol (str): 协议
            port (int): 端口号
            group (str, optional): 组播地址. Defaults to ''.
            on_recv (function, optional): 接收到数据时的回调函数, 参数为 (data: bytes, client_name: str). Defaults to None.

        Raises:
            ValueError: 无效的端口号, 应为 [1-65535]
            ValueError: 无效的协议类型, 应为 [TCP, UDP, MULTICAST]
        """
        if port < 1 or port > 65535:
            raise ValueError(f'ServerSocket 无效的端口号 "{port}"')
        if protocol not in ['TCP', 'UDP', 'MULTICAST']:
            raise ValueError(f'ServerSocket 无效的协议类型 "{protocol}"')
        if protocol == 'MULTICAST' and not group:
            raise ValueError(f'ServerSocket 组播协议必须指定组播地址')
        if protocol != 'MULTICAST' and group:
            raise ValueError(f'ServerSocket 协议类型 "{protocol}" 请勿设置 group 参数')


        self.protocol = protocol
        self.port = port
        self.group = group
        self.on_recv = on_recv
        self.sock: socket.socket | None = None
        self.tcp_sub_socks: list[socket.socket] = []
        self.thread: threading.Thread | None = None
        self.__active = False

    def __str__(self) -> str:
        if self.protocol == 'MULTICAST':
            return f'ServerSocket({self.protocol}, {self.group}:{self.port})'
        return f'ServerSocket({self.protocol}, {self.port})'

    def __create_socket(self) -> None:
        match self.protocol:
            case 'TCP':
                self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                self.sock.bind(('0.0.0.0', self.port))
                self.sock.listen(10)
            case 'UDP':
                self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
                self.sock.bind(('0.0.0.0', self.port))
            case 'MULTICAST':
                self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
                self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
                self.sock.bind(('0.0.0.0', self.port))
                self.sock.setsockopt(
                    socket.IPPROTO_IP,
                    socket.IP_ADD_MEMBERSHIP,
                    socket.inet_aton(self.group) + socket.INADDR_ANY.to_bytes(4, byteorder='big'))

    def __send_back(self, client_addr: tuple[str, int], client_sock: socket.socket | None = None) -> 'function':
        def send_back(data: bytes):
            if self.protocol == 'TCP':
                return client_sock.sendto(data, client_addr)
            return self.sock.sendto(data, client_addr)

        return send_back

    def __tcp_sub_thread(self, client_sock: socket.socket, client_addr: tuple[str, int]) -> None:
        while self.is_active():
            try:
                if not (data := client_sock.recv(1024)):
                    print(f'{self} TCP 子线程 {client_addr} 正常断开')
                    break

                self.on_recv and self.on_recv(
                    data=data,
                    client_addr=client_addr,
                    send_back=self.__send_back(client_addr, client_sock),
                )
            except ConnectionResetError:
                print(f'{self} TCP 子线程 {client_addr} 连接已重置')
                break
            except ConnectionAbortedError:
                print(f'{self} TCP 子线程 {client_addr} 连接已终止')
                break
            except Exception as e:
                if self.is_active():
                    print(f'{self} TCP 子线程 {client_addr} 异常: \n{e}')
                break
        if self.is_active(): # 断开或异常
            self.tcp_sub_socks.remove(client_sock)
            client_sock.close()

    def __main_thread(self) -> None:
        self.__active = True

        while self.is_active():
            try:
                if self.protocol == 'TCP':
                        client_sock, client_addr = self.sock.accept()
                        self.tcp_sub_socks.append(client_sock)
                        threading.Thread(target=self.__tcp_sub_thread, args=(client_sock, client_addr), daemon=True).start()
                else:
                    data, client_addr = self.sock.recvfrom(1024)

                    self.on_recv and self.on_recv(
                        data=data,
                        client_addr=client_addr,
                        send_back=self.__send_back(client_addr),
                    )
            except Exception as e:
                if self.is_active():
                    print(f'{self} 主线程异常 : \n{e}')
                    break

    def start(self) -> bool:
        """启动服务端

        将在新线程中运行，直到调用 close() 关闭，TCP 协议下会创建子线程处理 TCP 连接

        Returns:
            bool: 是否启动成功
        """
        try:
            self.__create_socket()
        except Exception as e:
            print(f'{self} 创建失败: \n{e}')
            return False

        self.thread = threading.Thread(target=self.__main_thread, daemon=True)
        self.thread.start()
        return True

    def close(self) -> bool:
        """关闭服务端

        Returns:
            bool: 是否关闭成功
        """
        try:
            self.__active = False
            if self.protocol == 'TCP':
                for client_sock in self.tcp_sub_socks:
                    client_sock.shutdown(socket.SHUT_RDWR)
                    client_sock.close()
                self.tcp_sub_socks.clear()
            else:
                self.sock.shutdown(socket.SHUT_RDWR)
            self.sock.close()
        except Exception as e:
            print(f'{self} 关闭失败: \n{e}')
            return False
        return True

    def is_active(self) -> bool:
        """返回服务端是否处于活动状态

        Returns:
            bool: 是否处于活动状态
        """
        return self.__active


quit = False
def on_recv(data, client_addr, send_back):
    global quit

    print(f'收到数据: {data} 来自 {client_addr}')
    send_back(b'ok')
    if data == b'q':
        quit = True
        # server.close()



if __name__ == '__main__':
    from time import sleep
    # server = ServerSocket(protocol='TCP', port=60000, on_recv=on_recv)
    # server = ServerSocket(protocol='UDP', port=60000, on_recv=on_recv)
    server = ServerSocket(protocol='MULTICAST', group='224.1.1.1', port=65000, on_recv=on_recv)
    server.start()

    while server.is_active():
        print('等待数据...')
        if quit:
            print('退出')
            server.close()
        sleep(1)