Skip to content

Commit

Permalink
fix: allow ip hosts and hostnames (#275)
Browse files Browse the repository at this point in the history
  • Loading branch information
magnasilvar committed Jan 16, 2024
1 parent 72c7a2b commit fc15fee
Show file tree
Hide file tree
Showing 2 changed files with 14 additions and 4 deletions.
10 changes: 6 additions & 4 deletions pytest_socket.py
Original file line number Diff line number Diff line change
Expand Up @@ -227,22 +227,24 @@ def socket_allow_hosts(allowed=None, allow_unix_socket=False):
if not isinstance(allowed, list):
return

allowed_hosts_by_host = normalize_allowed_hosts(allowed)
allowed_hosts = set(itertools.chain(*allowed_hosts_by_host.values()))
allowed_ip_hosts_by_host = normalize_allowed_hosts(allowed)
allowed_ip_hosts_and_hostnames = set(
itertools.chain(*allowed_ip_hosts_by_host.values())
) | set(allowed_ip_hosts_by_host.keys())
allowed_list = sorted(
[
(
host
if len(normalized) == 1 and next(iter(normalized)) == host
else f"{host} ({','.join(sorted(normalized))})"
)
for host, normalized in allowed_hosts_by_host.items()
for host, normalized in allowed_ip_hosts_by_host.items()
]
)

def guarded_connect(inst, *args):
host = host_from_connect_args(args)
if host in allowed_hosts or (
if host in allowed_ip_hosts_and_hostnames or (
_is_unix_socket(inst.family) and allow_unix_socket
):
return _true_connect(inst, *args)
Expand Down
8 changes: 8 additions & 0 deletions tests/test_restrict_hosts.py
Original file line number Diff line number Diff line change
Expand Up @@ -123,6 +123,14 @@ def test_single_cli_arg_connect_enabled_localhost_resolved(assert_connect):
assert_connect(True, cli_arg="localhost")


def test_single_cli_arg_127_0_0_1_hostname_localhost_connect_disabled(assert_connect):
assert_connect(False, cli_arg=localhost, host="localhost")


def test_single_cli_arg_localhost_hostname_localhost_connect_enabled(assert_connect):
assert_connect(True, cli_arg="localhost", host="localhost")


def test_single_cli_arg_connect_disabled_hostname_resolved(assert_connect):
result = assert_connect(
False,
Expand Down

0 comments on commit fc15fee

Please sign in to comment.