# 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)