152 lines
4.8 KiB
Python
152 lines
4.8 KiB
Python
# Originally from astropy project (http://astropy.org), under BSD
|
|
# 3-clause license.
|
|
|
|
import contextlib
|
|
import socket
|
|
|
|
from matplotlib import cbook
|
|
cbook.warn_deprecated("3.2", name=__name__, obj_type="module",
|
|
alternative="pytest-remotedata")
|
|
|
|
# save original socket method for restoration
|
|
# These are global so that re-calling the turn_off_internet function doesn't
|
|
# overwrite them again
|
|
socket_original = socket.socket
|
|
socket_create_connection = socket.create_connection
|
|
socket_bind = socket.socket.bind
|
|
socket_connect = socket.socket.connect
|
|
|
|
|
|
INTERNET_OFF = False
|
|
|
|
# urllib2 uses a global variable to cache its default "opener" for opening
|
|
# connections for various protocols; we store it off here so we can restore to
|
|
# the default after re-enabling internet use
|
|
_orig_opener = None
|
|
|
|
|
|
# ::1 is apparently another valid name for localhost?
|
|
# it is returned by getaddrinfo when that function is given localhost
|
|
|
|
def check_internet_off(original_function):
|
|
"""
|
|
Wraps ``original_function``, which in most cases is assumed
|
|
to be a `socket.socket` method, to raise an `IOError` for any operations
|
|
on non-local AF_INET sockets.
|
|
"""
|
|
|
|
def new_function(*args, **kwargs):
|
|
if isinstance(args[0], socket.socket):
|
|
if not args[0].family in (socket.AF_INET, socket.AF_INET6):
|
|
# Should be fine in all but some very obscure cases
|
|
# More to the point, we don't want to affect AF_UNIX
|
|
# sockets.
|
|
return original_function(*args, **kwargs)
|
|
host = args[1][0]
|
|
addr_arg = 1
|
|
valid_hosts = ('localhost', '127.0.0.1', '::1')
|
|
else:
|
|
# The only other function this is used to wrap currently is
|
|
# socket.create_connection, which should be passed a 2-tuple, but
|
|
# we'll check just in case
|
|
if not (isinstance(args[0], tuple) and len(args[0]) == 2):
|
|
return original_function(*args, **kwargs)
|
|
|
|
host = args[0][0]
|
|
addr_arg = 0
|
|
valid_hosts = ('localhost', '127.0.0.1')
|
|
|
|
hostname = socket.gethostname()
|
|
fqdn = socket.getfqdn()
|
|
|
|
if host in (hostname, fqdn):
|
|
host = 'localhost'
|
|
new_addr = (host, args[addr_arg][1])
|
|
args = args[:addr_arg] + (new_addr,) + args[addr_arg + 1:]
|
|
|
|
if any(h in host for h in valid_hosts):
|
|
return original_function(*args, **kwargs)
|
|
else:
|
|
raise IOError("An attempt was made to connect to the internet "
|
|
"by a test that was not marked `remote_data`.")
|
|
return new_function
|
|
|
|
|
|
def turn_off_internet(verbose=False):
|
|
"""
|
|
Disable internet access via python by preventing connections from being
|
|
created using the socket module. Presumably this could be worked around by
|
|
using some other means of accessing the internet, but all default python
|
|
modules (urllib, requests, etc.) use socket [citation needed].
|
|
"""
|
|
import urllib.request
|
|
|
|
global INTERNET_OFF
|
|
global _orig_opener
|
|
|
|
if INTERNET_OFF:
|
|
return
|
|
|
|
INTERNET_OFF = True
|
|
|
|
__tracebackhide__ = True
|
|
if verbose:
|
|
print("Internet access disabled")
|
|
|
|
# Update urllib2 to force it not to use any proxies
|
|
# Must use {} here (the default of None will kick off an automatic search
|
|
# for proxies)
|
|
_orig_opener = urllib.request.build_opener()
|
|
no_proxy_handler = urllib.request.ProxyHandler({})
|
|
opener = urllib.request.build_opener(no_proxy_handler)
|
|
urllib.request.install_opener(opener)
|
|
|
|
socket.create_connection = check_internet_off(socket_create_connection)
|
|
socket.socket.bind = check_internet_off(socket_bind)
|
|
socket.socket.connect = check_internet_off(socket_connect)
|
|
|
|
return socket
|
|
|
|
|
|
def turn_on_internet(verbose=False):
|
|
"""
|
|
Restore internet access. Not used, but kept in case it is needed.
|
|
"""
|
|
import urllib.request
|
|
|
|
global INTERNET_OFF
|
|
global _orig_opener
|
|
|
|
if not INTERNET_OFF:
|
|
return
|
|
|
|
INTERNET_OFF = False
|
|
|
|
if verbose:
|
|
print("Internet access enabled")
|
|
|
|
urllib.request.install_opener(_orig_opener)
|
|
|
|
socket.create_connection = socket_create_connection
|
|
socket.socket.bind = socket_bind
|
|
socket.socket.connect = socket_connect
|
|
return socket
|
|
|
|
|
|
@contextlib.contextmanager
|
|
def no_internet(verbose=False):
|
|
"""Context manager to temporarily disable internet access (if not already
|
|
disabled). If it was already disabled before entering the context manager
|
|
(i.e. `turn_off_internet` was called previously) then this is a no-op and
|
|
leaves internet access disabled until a manual call to `turn_on_internet`.
|
|
"""
|
|
|
|
already_disabled = INTERNET_OFF
|
|
|
|
turn_off_internet(verbose=verbose)
|
|
try:
|
|
yield
|
|
finally:
|
|
if not already_disabled:
|
|
turn_on_internet(verbose=verbose)
|