From f8aa8f22f1e7e38d448475e57f4b81a88aeff7b1 Mon Sep 17 00:00:00 2001 From: Thomas Steen Rasmussen Date: Mon, 19 Feb 2024 17:34:26 +0100 Subject: [PATCH] Add socks support for tcp lookups (#85) * add initial socks support, only socks5 and only works for tcp, doh, no tests and no dot yet * export ttl metrics for all rrs in all rrsets for #81 * add PySocks to pyproject.toml * parse socks proxy with urllib.parse.urlsplit and require scheme, support socks4+5+http proxies * fix json dump when socks_proxy is empty * improve unit tests for socks proxy code * rename socks_proxy to proxy, catch errors when starting listener, improve tests * add port busy test * fix a few comments and debug log entries * http proxies default to port 8080, proxy only works for plain tcp queries, include proxy in labels, fix unit test --- pyproject.toml | 1 + src/conftest.py | 3 +++ src/dns_exporter/collector.py | 17 ++++++++++++++++ src/dns_exporter/config.py | 23 +++++++++++++++++---- src/dns_exporter/entrypoint.py | 8 +++++++- src/dns_exporter/exporter.py | 37 ++++++++++++++++++++++++++++++++++ src/dns_exporter/metrics.py | 2 ++ src/tests/test_config.py | 9 +++++++++ src/tests/test_entrypoint.py | 22 ++++++++++++++++++++ src/tests/test_exporter.py | 9 +++++++-- 10 files changed, 124 insertions(+), 7 deletions(-) create mode 100644 src/tests/test_entrypoint.py diff --git a/pyproject.toml b/pyproject.toml index cadde18..31d498d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,6 +16,7 @@ dependencies = [ "dnspython[doh,dnssec,idna,doq] >= 2.2.1", "PyYAML >= 6.0", "prometheus-client >= 0.15.0", + "PySocks >= 1.7.1", ] description = "Prometheus exporter for blackbox-style DNS monitoring" dynamic = ["version"] diff --git a/src/conftest.py b/src/conftest.py index 2009127..028178c 100644 --- a/src/conftest.py +++ b/src/conftest.py @@ -66,6 +66,9 @@ def dns_exporter_param_config(request): proc = subprocess.Popen( args=["dns_exporter", "-c", str(conf), "-d"], ) + if proc.poll(): + # process didn't start properly, bail out + return time.sleep(1) yield print(f"Stopping dns_exporter with config {request.param} on 127.0.0.1:15353 ...") diff --git a/src/dns_exporter/collector.py b/src/dns_exporter/collector.py index a29fbad..290dd48 100644 --- a/src/dns_exporter/collector.py +++ b/src/dns_exporter/collector.py @@ -1,5 +1,6 @@ import logging import re +import socket import time import typing as t import urllib.parse @@ -15,6 +16,7 @@ import dns.rdatatype import dns.resolver import httpx # type: ignore +import socks # type: ignore from dns.message import Message, QueryMessage from prometheus_client.core import CounterMetricFamily, GaugeMetricFamily from prometheus_client.registry import Collector @@ -50,6 +52,19 @@ def __init__( self.query = query self.labels = labels + # set proxy? + if self.config.proxy: + socks.set_default_proxy( + proxy_type=getattr(socks, self.config.proxy.scheme.upper()), + addr=self.config.proxy.hostname, + port=self.config.proxy.port, + ) + dns.query.socket_factory = socks.socksocket + logger.debug(f"Using proxy {self.config.proxy.geturl()}") + else: + dns.query.socket_factory = socket.socket + logger.debug("Not using a proxy for this request") + def describe(self) -> Iterator[Union[CounterMetricFamily, GaugeMetricFamily]]: """Describe the metrics that are to be returned by this collector.""" yield get_dns_qtime_metric() @@ -66,6 +81,7 @@ def collect( yield from self.collect_up() def collect_up(self) -> Iterator[GaugeMetricFamily]: + """Yield the up metric.""" yield GaugeMetricFamily( "up", "The value of this Gauge is always 1 when the dns_exporter is up", @@ -73,6 +89,7 @@ def collect_up(self) -> Iterator[GaugeMetricFamily]: ) def collect_dns(self) -> Iterator[Union[CounterMetricFamily, GaugeMetricFamily]]: + """Collect and yield DNS metrics.""" assert isinstance(self.config.ip, (IPv4Address, IPv6Address)) # mypy assert isinstance(self.config.server, urllib.parse.SplitResult) # mypy assert isinstance(self.config.server.port, int) # mypy diff --git a/src/dns_exporter/config.py b/src/dns_exporter/config.py index bbe7548..4bd4b6e 100644 --- a/src/dns_exporter/config.py +++ b/src/dns_exporter/config.py @@ -175,6 +175,9 @@ class Config: query_type: str """str: The query type used for this DNS query, like ``A`` or ``MX``. Default is ``A``""" + proxy: t.Optional[urllib.parse.SplitResult] + """str: The proxy to use for this DNS query, for example ``socks5://127.0.0.1:5000``. Supported proxy types are SOCKS4, SOCKS5, and HTTP. Leave empty to use no proxy. Default is no proxy.""" + recursion_desired: bool """bool: Set this bool to ``True`` to set the ``RD`` flag in the DNS query. Default is ``True``""" @@ -203,9 +206,7 @@ class Config: ) """IPv4Address | IPv6Address | None: The IP to use instead of using IP or hostname from server. Default is ``None``""" - server: t.Union[urllib.parse.SplitResult, None] = field( - default_factory=lambda: None - ) + server: t.Optional[urllib.parse.SplitResult] = field(default_factory=lambda: None) """urllib.parse.SplitResult | None: The DNS server to use in parsed form. Default is ``None``""" query_name: t.Optional[str] = field(default_factory=lambda: None) @@ -263,6 +264,15 @@ def __post_init__(self) -> None: "invalid_request_config", ) + # validate proxy + if self.proxy: + # proxy support only works for plain tcp for now + if self.protocol not in ["tcp"]: + logger.error(f"proxy not valid for protocol {self.protocol}") + raise ConfigError( + "invalid_request_config", + ) + @classmethod def create( cls: t.Type["Config"], @@ -277,6 +287,7 @@ def create( query_class: str = "IN", query_type: str = "A", recursion_desired: bool = True, + proxy: t.Optional[urllib.parse.SplitResult] = None, timeout: float = 5.0, validate_answer_rrs: RRValidator = RRValidator.create(), validate_authority_rrs: RRValidator = RRValidator.create(), @@ -319,6 +330,7 @@ def create( query_class=query_class.upper(), query_type=query_type.upper(), recursion_desired=recursion_desired, + proxy=proxy, timeout=float(timeout), validate_answer_rrs=validate_answer_rrs, validate_authority_rrs=validate_authority_rrs, @@ -336,6 +348,8 @@ def json(self) -> str: conf: dict[str, t.Any] = asdict(self) conf["ip"] = str(conf["ip"]) conf["server"] = conf["server"].geturl() + if conf["proxy"]: + conf["proxy"] = conf["proxy"].geturl() return json.dumps(conf) @@ -365,5 +379,6 @@ class ConfigDict(t.TypedDict, total=False): validate_response_flags: RFValidator valid_rcodes: list[str] ip: t.Union[IPv4Address, IPv6Address, None] - server: t.Union[urllib.parse.SplitResult, None] + server: t.Optional[urllib.parse.SplitResult] query_name: t.Optional[str] + proxy: t.Optional[urllib.parse.SplitResult] diff --git a/src/dns_exporter/entrypoint.py b/src/dns_exporter/entrypoint.py index de54cae..f26a8cf 100644 --- a/src/dns_exporter/entrypoint.py +++ b/src/dns_exporter/entrypoint.py @@ -166,7 +166,13 @@ def main(mockargs: Optional[list[str]] = None) -> None: logger.debug( f"Ready to serve requests. Starting listener on {args.listen_ip} port {args.port}..." ) - HTTPServer((args.listen_ip, args.port), handler).serve_forever() + try: + HTTPServer((args.listen_ip, args.port), handler).serve_forever() + except OSError: + logger.error( + f"Unable to start listener, maybe port {args.port} is in use? bailing out" + ) + sys.exit(1) if __name__ == "__main__": diff --git a/src/dns_exporter/exporter.py b/src/dns_exporter/exporter.py index 71f210c..97b4a05 100644 --- a/src/dns_exporter/exporter.py +++ b/src/dns_exporter/exporter.py @@ -29,6 +29,7 @@ import dns.rcode import dns.rdatatype import dns.resolver +import socks # type: ignore from prometheus_client import CollectorRegistry, MetricsHandler, exposition from prometheus_client.registry import RestrictedRegistry @@ -170,6 +171,39 @@ def prepare_config(cls, config: ConfigDict) -> ConfigDict: server=config["server"], protocol=config["protocol"] ) + # parse proxy? + if ( + "proxy" in config.keys() + and config["proxy"] + and not isinstance(config["proxy"], urllib.parse.SplitResult) + ): + if "://" not in config["proxy"]: + logger.error("No scheme in proxy") + raise ConfigError("invalid_request_proxy") + + # parse proxy into a SplitResult + splitresult = urllib.parse.urlsplit(config["proxy"]) + if ( + not splitresult.scheme + or splitresult.scheme.upper() not in socks.PROXY_TYPES.keys() + ): + logger.error(f"Invalid proxy scheme {splitresult}") + raise ConfigError("invalid_request_proxy") + + # make port explicit + if splitresult.port is None: + # SOCKS4 and SOCKS5 default to port 1080 + port = 8080 if splitresult.scheme == "http" else 1080 + splitresult = splitresult._replace( + netloc=f"{splitresult.netloc}:{port}" + ) + + # keep only scheme and netloc + config["proxy"] = urllib.parse.urlsplit( + splitresult.scheme + "://" + splitresult.netloc + ) + logger.debug(f"Using proxy {str(splitresult.geturl())}") + return config def validate_config(self) -> None: @@ -444,6 +478,9 @@ def do_GET(self) -> None: "port": str(self.config.server.port), "protocol": str(self.config.protocol), "family": str(self.config.family), + "proxy": str(self.config.proxy.geturl()) + if self.config.proxy + else "none", "query_name": str(self.config.query_name), "query_type": str(self.config.query_type), } diff --git a/src/dns_exporter/metrics.py b/src/dns_exporter/metrics.py index e222d08..fd1d2be 100644 --- a/src/dns_exporter/metrics.py +++ b/src/dns_exporter/metrics.py @@ -25,6 +25,7 @@ "port", "protocol", "family", + "proxy", "query_name", "query_type", "transport", @@ -41,6 +42,7 @@ "invalid_request_module", "invalid_request_config", "invalid_request_server", + "invalid_request_proxy", "invalid_request_family", "invalid_request_ip", "invalid_request_port", diff --git a/src/tests/test_config.py b/src/tests/test_config.py index c9ec66a..945aa4e 100644 --- a/src/tests/test_config.py +++ b/src/tests/test_config.py @@ -99,3 +99,12 @@ def test_rd_false(): prepared = DNSExporter.prepare_config(ConfigDict(recursion_desired="false")) c = Config.create(name="test", **prepared) assert c.recursion_desired is False + + +def test_proxy_for_unsupported_protocol(): + """Test proxy with a protocol not supported.""" + prepared = DNSExporter.prepare_config( + ConfigDict(protocol="udp", proxy="socks5://127.0.0.1") + ) + with pytest.raises(ConfigError): + Config.create(name="test", **prepared) diff --git a/src/tests/test_entrypoint.py b/src/tests/test_entrypoint.py new file mode 100644 index 0000000..cedd84e --- /dev/null +++ b/src/tests/test_entrypoint.py @@ -0,0 +1,22 @@ +# type: ignore +import time + +import pytest + +import dns_exporter.entrypoint + +mockargs = [ + "-c", + "dns_exporter/dns_exporter_example.yml", + "-d", + "-p", + "25353", +] + + +def test_listen_port_busy(dns_exporter_example_config, caplog): + """Test calling main() on a port which is already busy.""" + with pytest.raises(SystemExit): + dns_exporter.entrypoint.main(mockargs) + time.sleep(2) + assert "is in use?" in caplog.text diff --git a/src/tests/test_exporter.py b/src/tests/test_exporter.py index c2d1db7..9db4f29 100644 --- a/src/tests/test_exporter.py +++ b/src/tests/test_exporter.py @@ -82,12 +82,17 @@ def test_config_endpoint(dns_exporter_example_config): params={ "server": "dns.google", "query_name": "example.com", + "protocol": "tcp", + "proxy": "socks5://127.0.0.1:1081", }, ) config = r.json() - assert config["server"] == "udp://dns.google:53" + assert config["server"] == "tcp://dns.google:53" assert config["query_name"] == "example.com" + +def test_config_endpoint_2(dns_exporter_example_config): + """Test the /config endpoint some more.""" r = requests.get( "http://127.0.0.1:25353/config", params={ @@ -282,7 +287,7 @@ def test_internal_metrics(dns_exporter_example_config, caplog): dnsexp_http_responses_total{path="/query",response_code="200"} 40.0 dnsexp_http_responses_total{path="/",response_code="200"} 1.0 dnsexp_dns_queries_total 29.0 -dnsexp_dns_responsetime_seconds_bucket{additional="0",answer="1",authority="0",family="ipv4",flags="QR RA RD",ip="8.8.4.4",le="0.005",nsid="no_nsid",opcode="QUERY",port="53",protocol="udp",query_name="example.com",query_type="A",rcode="NOERROR",server="udp://dns.google:53",transport="UDP"} +dnsexp_dns_responsetime_seconds_bucket{additional="0",answer="1",authority="0",family="ipv4",flags="QR RA RD",ip="8.8.4.4",le="0.005",nsid="no_nsid",opcode="QUERY",port="53",protocol="udp",proxy="none",query_name="example.com",query_type="A",rcode="NOERROR",server="udp://dns.google:53",transport="UDP"} dnsexp_scrape_failures_total{reason="timeout"} 1.0 dnsexp_scrape_failures_total{reason="invalid_response_flags"} 6.0 dnsexp_scrape_failures_total{reason="invalid_response_answer_rrs"} 3.0