Add socks support

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"]

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 ...")

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,13 +81,15 @@ 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",
             value=1,
         )
 
     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

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]

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__":

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),
             }

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",

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)

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

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