disable_internet.py 4.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151
  1. # Originally from astropy project (http://astropy.org), under BSD
  2. # 3-clause license.
  3. import contextlib
  4. import socket
  5. from matplotlib import cbook
  6. cbook.warn_deprecated("3.2", name=__name__, obj_type="module",
  7. alternative="pytest-remotedata")
  8. # save original socket method for restoration
  9. # These are global so that re-calling the turn_off_internet function doesn't
  10. # overwrite them again
  11. socket_original = socket.socket
  12. socket_create_connection = socket.create_connection
  13. socket_bind = socket.socket.bind
  14. socket_connect = socket.socket.connect
  15. INTERNET_OFF = False
  16. # urllib2 uses a global variable to cache its default "opener" for opening
  17. # connections for various protocols; we store it off here so we can restore to
  18. # the default after re-enabling internet use
  19. _orig_opener = None
  20. # ::1 is apparently another valid name for localhost?
  21. # it is returned by getaddrinfo when that function is given localhost
  22. def check_internet_off(original_function):
  23. """
  24. Wraps ``original_function``, which in most cases is assumed
  25. to be a `socket.socket` method, to raise an `IOError` for any operations
  26. on non-local AF_INET sockets.
  27. """
  28. def new_function(*args, **kwargs):
  29. if isinstance(args[0], socket.socket):
  30. if not args[0].family in (socket.AF_INET, socket.AF_INET6):
  31. # Should be fine in all but some very obscure cases
  32. # More to the point, we don't want to affect AF_UNIX
  33. # sockets.
  34. return original_function(*args, **kwargs)
  35. host = args[1][0]
  36. addr_arg = 1
  37. valid_hosts = ('localhost', '127.0.0.1', '::1')
  38. else:
  39. # The only other function this is used to wrap currently is
  40. # socket.create_connection, which should be passed a 2-tuple, but
  41. # we'll check just in case
  42. if not (isinstance(args[0], tuple) and len(args[0]) == 2):
  43. return original_function(*args, **kwargs)
  44. host = args[0][0]
  45. addr_arg = 0
  46. valid_hosts = ('localhost', '127.0.0.1')
  47. hostname = socket.gethostname()
  48. fqdn = socket.getfqdn()
  49. if host in (hostname, fqdn):
  50. host = 'localhost'
  51. new_addr = (host, args[addr_arg][1])
  52. args = args[:addr_arg] + (new_addr,) + args[addr_arg + 1:]
  53. if any(h in host for h in valid_hosts):
  54. return original_function(*args, **kwargs)
  55. else:
  56. raise IOError("An attempt was made to connect to the internet "
  57. "by a test that was not marked `remote_data`.")
  58. return new_function
  59. def turn_off_internet(verbose=False):
  60. """
  61. Disable internet access via python by preventing connections from being
  62. created using the socket module. Presumably this could be worked around by
  63. using some other means of accessing the internet, but all default python
  64. modules (urllib, requests, etc.) use socket [citation needed].
  65. """
  66. import urllib.request
  67. global INTERNET_OFF
  68. global _orig_opener
  69. if INTERNET_OFF:
  70. return
  71. INTERNET_OFF = True
  72. __tracebackhide__ = True
  73. if verbose:
  74. print("Internet access disabled")
  75. # Update urllib2 to force it not to use any proxies
  76. # Must use {} here (the default of None will kick off an automatic search
  77. # for proxies)
  78. _orig_opener = urllib.request.build_opener()
  79. no_proxy_handler = urllib.request.ProxyHandler({})
  80. opener = urllib.request.build_opener(no_proxy_handler)
  81. urllib.request.install_opener(opener)
  82. socket.create_connection = check_internet_off(socket_create_connection)
  83. socket.socket.bind = check_internet_off(socket_bind)
  84. socket.socket.connect = check_internet_off(socket_connect)
  85. return socket
  86. def turn_on_internet(verbose=False):
  87. """
  88. Restore internet access. Not used, but kept in case it is needed.
  89. """
  90. import urllib.request
  91. global INTERNET_OFF
  92. global _orig_opener
  93. if not INTERNET_OFF:
  94. return
  95. INTERNET_OFF = False
  96. if verbose:
  97. print("Internet access enabled")
  98. urllib.request.install_opener(_orig_opener)
  99. socket.create_connection = socket_create_connection
  100. socket.socket.bind = socket_bind
  101. socket.socket.connect = socket_connect
  102. return socket
  103. @contextlib.contextmanager
  104. def no_internet(verbose=False):
  105. """Context manager to temporarily disable internet access (if not already
  106. disabled). If it was already disabled before entering the context manager
  107. (i.e. `turn_off_internet` was called previously) then this is a no-op and
  108. leaves internet access disabled until a manual call to `turn_on_internet`.
  109. """
  110. already_disabled = INTERNET_OFF
  111. turn_off_internet(verbose=verbose)
  112. try:
  113. yield
  114. finally:
  115. if not already_disabled:
  116. turn_on_internet(verbose=verbose)