#!/usr/bin/env python
# This software released into the public domain. Anyone is free to copy,
# modify, publish, use, compile, sell, or distribute this software,
# either in source code form or as a compiled binary, for any purpose,
# commercial or non-commercial, and by any means.

import socket
import ctypes
import os

# 定义一个结构体，用于存储socket地址信息
class sockaddr(ctypes.Structure):
    _fields_ = [("sa_family", ctypes.c_short),   # 地址族（例如AF_INET或AF_INET6）
                ("__pad1", ctypes.c_ushort),     # 填充字段
                ("ipv4_addr", ctypes.c_byte * 4), # IPv4地址（4个字节）
                ("ipv6_addr", ctypes.c_byte * 16),# IPv6地址（16个字节）
                ("__pad2", ctypes.c_ulong)]     # 填充字段

# 根据操作系统的不同，导入不同的库
if hasattr(ctypes, 'windll'):
    WSAStringToAddressA = ctypes.windll.ws2_32.WSAStringToAddressA
    WSAAddressToStringA = ctypes.windll.ws2_32.WSAAddressToStringA
else:
    def not_windows():
        raise SystemError(
            "Invalid platform. ctypes.windll must be available."
        )
    WSAStringToAddressA = not_windows
    WSAAddressToStringA = not_windows

# inet_pton函数：将IP字符串转换为二进制格式
def inet_pton(address_family, ip_string):
    addr = sockaddr()  # 创建sockaddr实例
    addr.sa_family = address_family  # 设置地址族
    addr_size = ctypes.c_int(ctypes.sizeof(addr))  # 获取地址结构体大小

    # 使用WSAStringToAddressA函数将IP字符串转换为地址结构体
    if WSAStringToAddressA(
            ip_string,
            address_family,
            None,
            ctypes.byref(addr),
            ctypes.byref(addr_size)
    ) != 0:
        raise socket.error(ctypes.FormatError())

    # 根据地址族返回对应的二进制IP地址
    if address_family == socket.AF_INET:
        return ctypes.string_at(addr.ipv4_addr, 4)
    if address_family == socket.AF_INET6:
        return ctypes.string_at(addr.ipv6_addr, 16)

    raise socket.error('unknown address family')

# inet_ntop函数：将二进制格式的IP地址转换为字符串
def inet_ntop(address_family, packed_ip):
    addr = sockaddr()  # 创建sockaddr实例
    addr.sa_family = address_family  # 设置地址族
    addr_size = ctypes.c_int(ctypes.sizeof(addr))  # 获取地址结构体大小
    ip_string = ctypes.create_string_buffer(128)  # 创建字符串缓冲区
    ip_string_size = ctypes.c_int(ctypes.sizeof(ip_string))  # 获取字符串缓冲区大小

    # 根据地址族将二进制IP地址复制到地址结构体中
    if address_family == socket.AF_INET:
        if len(packed_ip) != ctypes.sizeof(addr.ipv4_addr):
            raise socket.error('packed IP wrong length for inet_ntoa')
        ctypes.memmove(addr.ipv4_addr, packed_ip, 4)
    elif address_family == socket.AF_INET6:
        if len(packed_ip) != ctypes.sizeof(addr.ipv6_addr):
            raise socket.error('packed IP wrong length for inet_ntoa')
        ctypes.memmove(addr.ipv6_addr, packed_ip, 16)
    else:
        raise socket.error('unknown address family')

    # 使用WSAAddressToStringA函数将地址结构体转换为IP字符串
    if WSAAddressToStringA(
            ctypes.byref(addr),
            addr_size,
            None,
            ip_string,
            ctypes.byref(ip_string_size)
    ) != 0:
        raise socket.error(ctypes.FormatError())

    return ip_string[:ip_string_size.value - 1]

# 如果当前操作系统是Windows，将自定义的inet_pton和inet_ntop函数添加到socket库中
if os.name == 'nt':
    socket.inet_pton = inet_pton
    socket.inet_ntop = inet_ntop