1#!/usr/bin/env python
2# This software released into the public domain. Anyone is free to copy,
3# modify, publish, use, compile, sell, or distribute this software,
4# either in source code form or as a compiled binary, for any purpose,
5# commercial or non-commercial, and by any means.
6
7import socket
8import ctypes
9import os
10
11
12class sockaddr(ctypes.Structure):
13    _fields_ = [("sa_family", ctypes.c_short),
14                ("__pad1", ctypes.c_ushort),
15                ("ipv4_addr", ctypes.c_byte * 4),
16                ("ipv6_addr", ctypes.c_byte * 16),
17                ("__pad2", ctypes.c_ulong)]
18
19if hasattr(ctypes, 'windll'):
20    WSAStringToAddressA = ctypes.windll.ws2_32.WSAStringToAddressA
21    WSAAddressToStringA = ctypes.windll.ws2_32.WSAAddressToStringA
22else:
23    def not_windows():
24        raise SystemError(
25            "Invalid platform. ctypes.windll must be available."
26        )
27    WSAStringToAddressA = not_windows
28    WSAAddressToStringA = not_windows
29
30
31def inet_pton(address_family, ip_string):
32    addr = sockaddr()
33    addr.sa_family = address_family
34    addr_size = ctypes.c_int(ctypes.sizeof(addr))
35
36    if WSAStringToAddressA(
37            ip_string,
38            address_family,
39            None,
40            ctypes.byref(addr),
41            ctypes.byref(addr_size)
42    ) != 0:
43        raise socket.error(ctypes.FormatError())
44
45    if address_family == socket.AF_INET:
46        return ctypes.string_at(addr.ipv4_addr, 4)
47    if address_family == socket.AF_INET6:
48        return ctypes.string_at(addr.ipv6_addr, 16)
49
50    raise socket.error('unknown address family')
51
52
53def inet_ntop(address_family, packed_ip):
54    addr = sockaddr()
55    addr.sa_family = address_family
56    addr_size = ctypes.c_int(ctypes.sizeof(addr))
57    ip_string = ctypes.create_string_buffer(128)
58    ip_string_size = ctypes.c_int(ctypes.sizeof(ip_string))
59
60    if address_family == socket.AF_INET:
61        if len(packed_ip) != ctypes.sizeof(addr.ipv4_addr):
62            raise socket.error('packed IP wrong length for inet_ntoa')
63        ctypes.memmove(addr.ipv4_addr, packed_ip, 4)
64    elif address_family == socket.AF_INET6:
65        if len(packed_ip) != ctypes.sizeof(addr.ipv6_addr):
66            raise socket.error('packed IP wrong length for inet_ntoa')
67        ctypes.memmove(addr.ipv6_addr, packed_ip, 16)
68    else:
69        raise socket.error('unknown address family')
70
71    if WSAAddressToStringA(
72            ctypes.byref(addr),
73            addr_size,
74            None,
75            ip_string,
76            ctypes.byref(ip_string_size)
77    ) != 0:
78        raise socket.error(ctypes.FormatError())
79
80    return ip_string[:ip_string_size.value - 1]
81
82# Adding our two functions to the socket library
83if os.name == 'nt':
84    socket.inet_pton = inet_pton
85    socket.inet_ntop = inet_ntop
86