feat: Support Python subprocesses

README.md

@@ -16,6 +16,7 @@ tests to ensure network calls are prevented.
 ## Features
 
 - Disables all network calls flowing through Python\'s `socket` interface.
+- Python subprocesses are supported
 
 ## Requirements
 

pyproject.toml

@@ -12,6 +12,7 @@ include = [
     { path = "README.md", format = "sdist" },
     { path = "tests", format = "sdist" },
     { path = ".flake8", format = "sdist" },
+    { path = "pytest_socket.pth", format = ["sdist", "wheel"] }
 ]
 classifiers = [
     "Development Status :: 4 - Beta",

pytest_socket.embed

@@ -0,0 +1,48 @@
+"""Inject pytest-socket into Python subprocesses via a .pth file.
+
+To update the .pth file, simply run this script which will write the code
+below to a .pth file in the same directory as a single line.
+"""
+
+import os
+os.environ["_PYTEST_SOCKET_SUBPROCESS"] = ""
+
+# .PTH START
+if config := os.getenv("_PYTEST_SOCKET_SUBPROCESS", None):
+    import json
+    state = None
+    try:
+        import socket
+        from pytest_socket import disable_socket, _create_guarded_connect
+        state = json.loads(config)
+
+        if state["mode"] == "disable":
+            disable_socket(allow_unix_socket=state["allow_unix_socket"])
+        elif state["mode"] == "allow-hosts":
+            socket.socket.connect = _create_guarded_connect(
+                allowed_hosts=state["allowed_hosts"],
+                allow_unix_socket=state["allow_unix_socket"],
+                _pretty_allowed_list=state["_pretty_allowed_list"]
+            )
+
+    except Exception as exc:
+        import sys
+        sys.stderr.write(
+            "pytest-socket: Failed to set up subprocess socket patching.\n"
+            f"Configuration: {state}\n"
+            f"{exc.__class__.__name__}: {exc}\n"
+        )
+# .PTH END
+
+if __name__ == "__main__":
+    from pathlib import Path
+
+    src = Path(__file__)
+    dst = src.with_suffix(".pth")
+    lines = src.read_text().splitlines()
+    code = "\n".join(lines[lines.index("# .PTH START") + 1 : lines.index("# .PTH END")])
+
+    print(f"Writing to {dst}")
+    # Only lines beginning with an import will be executed.
+    # https://docs.python.org/3/library/site.html
+    dst.write_text(f"import os; exec({code!r})\n")

pytest_socket.pth

@@ -0,0 +1 @@
+import os; exec('if config := os.getenv("_PYTEST_SOCKET_SUBPROCESS", None):\n    import json\n    state = None\n    try:\n        import socket\n        from pytest_socket import disable_socket, _create_guarded_connect\n        state = json.loads(config)\n\n        if state["mode"] == "disable":\n            disable_socket(allow_unix_socket=state["allow_unix_socket"])\n        elif state["mode"] == "allow-hosts":\n            socket.socket.connect = _create_guarded_connect(\n                allowed_hosts=state["allowed_hosts"],\n                allow_unix_socket=state["allow_unix_socket"],\n                _pretty_allowed_list=state["_pretty_allowed_list"]\n            )\n\n    except Exception as exc:\n        import sys\n        sys.stderr.write(\n            "pytest-socket: Failed to set up subprocess socket patching.\\n"\n            f"Configuration: {state}\\n"\n            f"{exc.__class__.__name__}: {exc}\\n"\n        )')

pytest_socket.py

@@ -1,16 +1,34 @@
 import ipaddress
 import itertools
+import json
+import os
 import socket
 import typing
 from collections import defaultdict
 from dataclasses import dataclass, field
 
 import pytest
 
+_SUBPROCESS_ENVVAR = "_PYTEST_SOCKET_SUBPROCESS"
 _true_socket = socket.socket
 _true_connect = socket.socket.connect
 
 
+def update_subprocess_config(config: typing.Dict[str, object]) -> None:
+    """Enable pytest-socket in Python subprocesses.
+
+    The configuration will be read by the .pth file to mirror the
+    restrictions in the main process.
+    """
+    os.environ[_SUBPROCESS_ENVVAR] = json.dumps(config)
+
+
+def delete_subprocess_config() -> None:
+    """Disable pytest-socket in Python subprocesses."""
+    if _SUBPROCESS_ENVVAR in os.environ:
+        del os.environ[_SUBPROCESS_ENVVAR]
+
+
 class SocketBlockedError(RuntimeError):
     def __init__(self, *_args, **_kwargs):
         super().__init__("A test tried to use socket.socket.")
@@ -103,11 +121,15 @@ def __new__(cls, family=-1, type=-1, proto=-1, fileno=None):
             raise SocketBlockedError()
 
     socket.socket = GuardedSocket
+    update_subprocess_config(
+        {"mode": "disable", "allow_unix_socket": allow_unix_socket}
+    )
 
 
 def enable_socket():
     """re-enable socket.socket to enable the Internet. useful in testing."""
     socket.socket = _true_socket
+    delete_subprocess_config()
 
 
 def pytest_configure(config):
@@ -249,6 +271,25 @@ def normalize_allowed_hosts(
     return ip_hosts
 
 
+def _create_guarded_connect(
+    allowed_hosts: typing.Sequence[str],
+    allow_unix_socket: bool,
+    _pretty_allowed_list: typing.Sequence[str],
+) -> typing.Callable:
+    """Create a function to replace socket.connect."""
+
+    def guarded_connect(inst, *args):
+        host = host_from_connect_args(args)
+        if host in allowed_hosts or (
+            _is_unix_socket(inst.family) and allow_unix_socket
+        ):
+            return _true_connect(inst, *args)
+
+        raise SocketConnectBlockedError(_pretty_allowed_list, host)
+
+    return guarded_connect
+
+
 def socket_allow_hosts(
     allowed: typing.Union[str, typing.List[str], None] = None,
     allow_unix_socket: bool = False,
@@ -276,19 +317,21 @@ def socket_allow_hosts(
         ]
     )
 
-    def guarded_connect(inst, *args):
-        host = host_from_connect_args(args)
-        if host in allowed_ip_hosts_and_hostnames or (
-            _is_unix_socket(inst.family) and allow_unix_socket
-        ):
-            return _true_connect(inst, *args)
-
-        raise SocketConnectBlockedError(allowed_list, host)
-
-    socket.socket.connect = guarded_connect
+    socket.socket.connect = _create_guarded_connect(
+        allowed_ip_hosts_and_hostnames, allow_unix_socket, allowed_list
+    )
+    update_subprocess_config(
+        {
+            "mode": "allow-hosts",
+            "allowed_hosts": list(allowed_ip_hosts_and_hostnames),
+            "allow_unix_socket": allow_unix_socket,
+            "_pretty_allowed_list": allowed_list,
+        }
+    )
 
 
 def _remove_restrictions():
     """restore socket.socket.* to allow access to the Internet. useful in testing."""
     socket.socket = _true_socket
     socket.socket.connect = _true_connect
+    delete_subprocess_config()

tests/common.py

@@ -14,3 +14,9 @@ def assert_socket_blocked(result, passed=0, skipped=0, failed=1):
     result.stdout.fnmatch_lines(
         "*Socket*Blocked*Error: A test tried to use socket.socket.*"
     )
+
+
+def assert_host_blocked(result, host):
+    result.stdout.fnmatch_lines(
+        f'*A test tried to use socket.socket.connect() with host "{host}"*'
+    )

tests/test_restrict_hosts.py

@@ -4,7 +4,7 @@
 
 import pytest
 
-from pytest_socket import normalize_allowed_hosts
+from pytest_socket import assert_host_blocked, normalize_allowed_hosts
 
 localhost = "127.0.0.1"
 
@@ -46,12 +46,6 @@ def {2}():
 """
 
 
-def assert_host_blocked(result, host):
-    result.stdout.fnmatch_lines(
-        f'*A test tried to use socket.socket.connect() with host "{host}"*'
-    )
-
-
 @pytest.fixture
 def assert_connect(httpbin, testdir):
     def assert_socket_connect(should_pass, **kwargs):