Skip to content

Commit

Permalink
FIREWALL: Allow to block hosts by IP address
Browse files Browse the repository at this point in the history
When blocking a host, its hostname is resolved using the `dig`
command. If an IP address is provided, `dig` returns nothing.
Check whether it is an IP address before launching `dig`.

No matter whether the correct family is requested, do not call
`dig` when an IP address is provided, in case the caller has
already blocked the DNS host.
  • Loading branch information
aplopez committed Oct 28, 2024
1 parent 2823630 commit f8f91f8
Showing 1 changed file with 13 additions and 2 deletions.
15 changes: 13 additions & 2 deletions pytest_mh/utils/firewall.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
from __future__ import annotations

from abc import abstractmethod
from ipaddress import IPv4Address, IPv6Address, ip_address
from random import randrange
from typing import Any, Literal, TypeAlias

Expand Down Expand Up @@ -605,9 +606,19 @@ def __add_host(
self.firewall.add_rich_rule(f"family=ipv6 destination address={ip} {action}")

def __resolve_hostname(self, hostname: str, type: Literal["A", "AAAA"]) -> list[str]:
result = self.firewall.host.conn.exec(["dig", "+short", "-t", type, hostname], log_level=ProcessLogLevel.Error)
addrs = []
try:
ip = ip_address(hostname)
ip_type = IPv4Address if type == "A" else IPv6Address
if isinstance(ip, ip_type):
addrs = [hostname]
except ValueError:
result = self.firewall.host.conn.exec(
["dig", "+short", "-t", type, hostname], log_level=ProcessLogLevel.Error
)
addrs = result.stdout_lines

return result.stdout_lines
return addrs


class WindowsFirewall(Firewall):
Expand Down

0 comments on commit f8f91f8

Please sign in to comment.