[WAN_IP] Fall back to HTTP properly on DNS connection errors (#134)

HorlogeSkynet · Jun 18, 2023 · c622948 · c622948
1 parent c025676
commit c622948
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 15 deletions.
diff --git a/archey/entries/wan_ip.py b/archey/entries/wan_ip.py
@@ -46,17 +46,14 @@ def _retrieve_ip_address(self, ip_version: int) -> Optional[str]:
         dns_query = options.get("dns_query", "myip.opendns.com")
         if dns_query:
             # Run the DNS query.
-            try:
-                ip_address = self._run_dns_query(
-                    dns_query,
-                    options.get("dns_resolver", "resolver1.opendns.com"),
-                    ip_version,
-                    options.get("dns_timeout", 1),
-                )
-            except FileNotFoundError:
-                # DNS lookup tool does not seem to be available.
-                pass
-            else:
+            ip_address = self._run_dns_query(
+                dns_query,
+                options.get("dns_resolver", "resolver1.opendns.com"),
+                ip_version,
+                options.get("dns_timeout", 1),
+            )
+            # Return IP only if the query was successful
+            if ip_address is not None:
                 return ip_address
 
         # Is retrieval via HTTP(S) request enabled ?
@@ -84,7 +81,7 @@ def _run_dns_query(query: str, resolver: str, ip_version: int, timeout: float) -
                 stderr=DEVNULL,
                 universal_newlines=True,
             ).rstrip()
-        except (TimeoutExpired, CalledProcessError):
+        except (FileNotFoundError, TimeoutExpired, CalledProcessError):
             return None
 
         # `ip_address` might be empty here.

diff --git a/archey/test/entries/test_archey_wan_ip.py b/archey/test/entries/test_archey_wan_ip.py
@@ -2,7 +2,7 @@
 
 import unittest
 from socket import timeout as SocketTimeoutError
-from subprocess import TimeoutExpired
+from subprocess import CalledProcessError, TimeoutExpired
 from unittest.mock import MagicMock, Mock, call, patch
 
 from archey.configuration import DEFAULT_CONFIG
@@ -53,7 +53,10 @@ def test_ipv4_ko_and_ipv6_ok(self, urlopen_mock, _):
         "archey.entries.wan_ip.check_output",
         side_effect=[
             "\n",  # `check_output` call will soft-fail.
-            FileNotFoundError("dig"),  # `check_output` call will hard-fail.
+            FileNotFoundError("dig"),  # `check_output` call will hard-fail (missing `dig`)
+            CalledProcessError(
+                cmd="dig", returncode=9
+            ),  # `check_output` call will hard-fail (DNS connection error)
         ],
     )
     @patch("archey.entries.wan_ip.urlopen")
@@ -66,7 +69,13 @@ def test_proper_http_fallback(self, urlopen_mock, _):
             WanIP._retrieve_ip_address(self.wan_ip_mock, 4),  # pylint: disable=protected-access
         )
 
-        # New try: HTTP method has been called !
+        # HTTP method has been called with `dig` missing.
+        self.assertEqual(
+            WanIP._retrieve_ip_address(self.wan_ip_mock, 4),  # pylint: disable=protected-access
+            "XXX.YY.ZZ.TTT",
+        )
+
+        # HTTP method has been called with `dig` unable to connect!
         self.assertEqual(
             WanIP._retrieve_ip_address(self.wan_ip_mock, 4),  # pylint: disable=protected-access
             "XXX.YY.ZZ.TTT",