mirror of
https://github.com/SystemRage/py-kms.git
synced 2024-10-18 12:28:16 +02:00
216 lines
10 KiB
Python
216 lines
10 KiB
Python
#!/usr/bin/env python3
|
|
|
|
import os
|
|
import socket
|
|
import selectors
|
|
import ipaddress
|
|
|
|
# https://github.com/python/cpython/blob/master/Lib/socket.py
|
|
def has_dualstack_ipv6():
|
|
""" Return True if the platform supports creating a SOCK_STREAM socket
|
|
which can handle both AF_INET and AF_INET6 (IPv4 / IPv6) connections.
|
|
"""
|
|
if not socket.has_ipv6 or not hasattr(socket._socket, 'IPPROTO_IPV6') or not hasattr(socket._socket, 'IPV6_V6ONLY'):
|
|
return False
|
|
try:
|
|
with socket.socket(socket.AF_INET6, socket.SOCK_STREAM) as sock:
|
|
sock.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, 0)
|
|
return True
|
|
except socket.error:
|
|
return False
|
|
|
|
def create_server_sock(address, *, family = socket.AF_INET, backlog = None, reuse_port = False, dualstack_ipv6 = False):
|
|
""" Convenience function which creates a SOCK_STREAM type socket
|
|
bound to *address* (a 2-tuple (host, port)) and return the socket object.
|
|
Internally it takes care of choosing the right address family (IPv4 or IPv6),depending on
|
|
the host specified in *address* tuple.
|
|
|
|
*family* should be either AF_INET or AF_INET6.
|
|
*backlog* is the queue size passed to socket.listen().
|
|
*reuse_port* dictates whether to use the SO_REUSEPORT socket option.
|
|
*dualstack_ipv6* if True and the platform supports it, it will create an AF_INET6 socket able to accept both IPv4 or IPv6 connections;
|
|
when False it will explicitly disable this option on platforms that enable it by default (e.g. Linux).
|
|
"""
|
|
if reuse_port and not hasattr(socket._socket, "SO_REUSEPORT"):
|
|
raise ValueError("SO_REUSEPORT not supported on this platform")
|
|
|
|
if dualstack_ipv6:
|
|
if not has_dualstack_ipv6():
|
|
raise ValueError("dualstack_ipv6 not supported on this platform")
|
|
if family != socket.AF_INET6:
|
|
raise ValueError("dualstack_ipv6 requires AF_INET6 family")
|
|
|
|
sock = socket.socket(family, socket.SOCK_STREAM)
|
|
try:
|
|
# Note about Windows. We don't set SO_REUSEADDR because:
|
|
# 1) It's unnecessary: bind() will succeed even in case of a
|
|
# previous closed socket on the same address and still in
|
|
# TIME_WAIT state.
|
|
# 2) If set, another socket is free to bind() on the same
|
|
# address, effectively preventing this one from accepting
|
|
# connections. Also, it may set the process in a state where
|
|
# it'll no longer respond to any signals or graceful kills.
|
|
# See: msdn2.microsoft.com/en-us/library/ms740621(VS.85).aspx
|
|
if os.name not in ('nt', 'cygwin') and hasattr(socket._socket, 'SO_REUSEADDR'):
|
|
try:
|
|
sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
|
except socket.error:
|
|
# Fail later on bind(), for platforms which may not
|
|
# support this option.
|
|
pass
|
|
if reuse_port:
|
|
sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)
|
|
if socket.has_ipv6 and family == socket.AF_INET6:
|
|
if dualstack_ipv6:
|
|
sock.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, 0)
|
|
elif hasattr(socket._socket, "IPV6_V6ONLY") and hasattr(socket._socket, "IPPROTO_IPV6"):
|
|
sock.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, 1)
|
|
try:
|
|
sock.bind(address)
|
|
except socket.error as err:
|
|
msg = '%s (while attempting to bind on address %r)' %(err.strerror, address)
|
|
raise socket.error(err.errno, msg) from None
|
|
|
|
if backlog is None:
|
|
sock.listen()
|
|
else:
|
|
sock.listen(backlog)
|
|
return sock
|
|
except socket.error:
|
|
sock.close()
|
|
raise
|
|
|
|
# Giampaolo Rodola' class (license MIT) revisited for py-kms.
|
|
# http://code.activestate.com/recipes/578504-server-supporting-ipv4-and-ipv6/
|
|
class MultipleListener(object):
|
|
""" Listen on multiple addresses specified as a list of
|
|
(`host`, `port`, `backlog`, `reuse_port`) tuples.
|
|
Useful to listen on both IPv4 and IPv6 on those systems where a dual stack
|
|
is not supported natively (Windows and many UNIXes).
|
|
|
|
Calls like settimeout() and setsockopt() will be applied to all sockets.
|
|
Calls like gettimeout() or getsockopt() will refer to the first socket in the list.
|
|
"""
|
|
def __init__(self, addresses = [], want_dual = False):
|
|
self.socks, self.sockmap = [], {}
|
|
completed = False
|
|
self.cant_dual = []
|
|
|
|
try:
|
|
for addr in addresses:
|
|
addr = self.check(addr)
|
|
ip_ver = ipaddress.ip_address(addr[0])
|
|
|
|
if ip_ver.version == 4 and want_dual:
|
|
self.cant_dual.append(addr[0])
|
|
|
|
sock = create_server_sock((addr[0], addr[1]),
|
|
family = (socket.AF_INET if ip_ver.version == 4 else socket.AF_INET6),
|
|
backlog = addr[2],
|
|
reuse_port = addr[3],
|
|
dualstack_ipv6 = (False if ip_ver.version == 4 else want_dual))
|
|
self.socks.append(sock)
|
|
self.sockmap[sock.fileno()] = sock
|
|
|
|
completed = True
|
|
finally:
|
|
if not completed:
|
|
self.close()
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self):
|
|
self.close()
|
|
|
|
def __repr__(self):
|
|
addrs = []
|
|
for sock in self.socks:
|
|
try:
|
|
addrs.append(sock.getsockname())
|
|
except socket.error:
|
|
addrs.append(())
|
|
return "<%s(%r) at %#x>" %(self.__class__.__name__, addrs, id(self))
|
|
|
|
def filenos(self):
|
|
""" Return sockets' file descriptors as a list of integers. """
|
|
return list(self.sockmap.keys())
|
|
|
|
def register(self, pollster):
|
|
for fd in self.filenos():
|
|
pollster.register(fileobj = fd, events = selectors.EVENT_READ)
|
|
|
|
def multicall(self, name, *args, **kwargs):
|
|
for sock in self.socks:
|
|
meth = getattr(sock, name)
|
|
meth(*args, **kwargs)
|
|
|
|
def poll(self):
|
|
""" Return the first readable fd. """
|
|
if hasattr(selectors, 'PollSelector'):
|
|
pollster = selectors.PollSelector
|
|
else:
|
|
pollster = selectors.SelectSelector
|
|
|
|
timeout = self.gettimeout()
|
|
|
|
with pollster() as pollster:
|
|
self.register(pollster)
|
|
fds = pollster.select(timeout)
|
|
|
|
if timeout and fds == []:
|
|
raise socket.timeout('timed out')
|
|
try:
|
|
return fds[0][0].fd
|
|
except IndexError:
|
|
# non-blocking socket
|
|
pass
|
|
|
|
def accept(self):
|
|
""" Accept a connection from the first socket which is ready to do so. """
|
|
fd = self.poll()
|
|
sock = (self.sockmap[fd] if fd else self.socks[0])
|
|
return sock.accept()
|
|
|
|
def getsockname(self):
|
|
""" Return first registered socket's own address. """
|
|
return self.socks[0].getsockname()
|
|
|
|
def getsockopt(self, level, optname, buflen = 0):
|
|
""" Return first registered socket's options. """
|
|
return self.socks[0].getsockopt(level, optname, buflen)
|
|
|
|
def gettimeout(self):
|
|
""" Return first registered socket's timeout. """
|
|
return self.socks[0].gettimeout()
|
|
|
|
def settimeout(self, timeout):
|
|
""" Set timeout for all registered sockets. """
|
|
self.multicall('settimeout', timeout)
|
|
|
|
def setblocking(self, flag):
|
|
""" Set non-blocking mode for all registered sockets. """
|
|
self.multicall('setblocking', flag)
|
|
|
|
def setsockopt(self, level, optname, value):
|
|
""" Set option for all registered sockets. """
|
|
self.multicall('setsockopt', level, optname, value)
|
|
|
|
def shutdown(self, how):
|
|
""" Shut down all registered sockets. """
|
|
self.multicall('shutdown', how)
|
|
|
|
def close(self):
|
|
""" Close all registered sockets. """
|
|
self.multicall('close')
|
|
self.socks, self.sockmap = [], {}
|
|
|
|
def check(self, address):
|
|
if len(address) == 1:
|
|
raise socket.error("missing `host` or `port` parameter.")
|
|
if len(address) == 2:
|
|
address += (None, True,)
|
|
elif len(address) == 3:
|
|
address += (True,)
|
|
return address
|