From c6229484fca6423eaa188b37d7d7d2c5447c63e2 Mon Sep 17 00:00:00 2001
From: Michael Bromilow <developer@bromilow.uk>
Date: Sun, 18 Jun 2023 08:02:09 +0100
Subject: [PATCH] [WAN_IP] Fall back to HTTP properly on DNS connection errors
 (#134)

---
 archey/entries/wan_ip.py                  | 21 +++++++++------------
 archey/test/entries/test_archey_wan_ip.py | 15 ++++++++++++---
 2 files changed, 21 insertions(+), 15 deletions(-)

diff --git a/archey/entries/wan_ip.py b/archey/entries/wan_ip.py
index 223da61e..4b031686 100644
--- 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
index 927f939f..a479bcf9 100644
--- 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",