记一个简单测试端口的python脚本

脚本的简单功能：

• 支持IPv4和IPv6
• 支持TCP和UDP端口

import socket, ipaddress
from enum import Enum


class Mode(Enum):
    TCP = 0
    UDP = 1

    def __repr__(self):
        return self.name

    __str__ = __repr__


def is_port_open(ip: str, port: int, mode: Mode = Mode.TCP, timeout: int = 5):
    sock = None
    ip = ip.strip()
    try:
        version = ipaddress.ip_address(ip).version
        if version == 4:
            af = socket.AF_INET
        elif version == 6:
            af = socket.AF_INET6
        else:
            print("IP地址不合法")
            exit(0)
        if mode == Mode.TCP:
            # 创建一个 TCP socket
            skt = socket.SOCK_STREAM
        elif mode == Mode.UDP:
            # 创建一个 UDP socket
            skt = socket.SOCK_DGRAM
        else:
            print("模式错误,端口无法测试")
            exit(1)
        sock = socket.socket(af, skt)
        # 设置一个超时值，以防止程序无限期地等待连接
        sock.settimeout(timeout)
        if mode == Mode.UDP:
            sock.sendto(b"Hello, World!", (ip, port))
            sock.recvfrom(1024)
        else:
            # 尝试连接到指定的主机和端口
            sock.connect((ip, port))
        # 如果连接成功，则端口是开放的
        return True
    except socket.error as e:
        print("socket连接失败:", e)
        # 如果连接失败，则端口是关闭的
        return False
    except ValueError as e:
        print(e)
        exit(1)
    finally:
        if sock:
            sock.close()


host = "10.0.0.10"
p = 22
mode = Mode.TCP
if is_port_open(host, p, mode):
    print(f"{host}的{mode}端口{p}是open")
else:
    print(f"{host}的{mode}端口{p}是closed")