From 7a8e1216662acf27124421dfbd2f170fa62f2057 Mon Sep 17 00:00:00 2001
From: Isaac Yang <47034756+seankingyang@users.noreply.github.com>
Date: Mon, 14 Oct 2024 20:10:00 +0800
Subject: [PATCH] Fix wifi netplan gateway parser (Bugfix) (#1536)

* Fix get_getway can correct parse the ip only

* Add unit test

* Fix Black format

* Add invalid ip unit test

* Fix flask8 format

* Update providers/base/bin/wifi_client_test_netplan.py

Co-authored-by: Massimiliano <massimiliano.girardi@canonical.com>

* Squash the _validate_gateway_ip and __validate_ip function

---------

Co-authored-by: Massimiliano <massimiliano.girardi@canonical.com>
---
 .../base/bin/wifi_client_test_netplan.py      | 23 ++++++++--
 .../tests/test_wifi_client_test_netplan.py    | 44 ++++++++++++++++++-
 2 files changed, 62 insertions(+), 5 deletions(-)

diff --git a/providers/base/bin/wifi_client_test_netplan.py b/providers/base/bin/wifi_client_test_netplan.py
index 7ec0f5fd6..2a1861a7a 100755
--- a/providers/base/bin/wifi_client_test_netplan.py
+++ b/providers/base/bin/wifi_client_test_netplan.py
@@ -22,6 +22,7 @@
 import time
 import shutil
 import sys
+import ipaddress
 import yaml
 
 from gateway_ping_test import ping
@@ -315,12 +316,28 @@ def print_route_info():
     print()
 
 
+def _validate_gateway_ip(gateway):
+    if not gateway:
+        return ""
+    # Check if the gateway is a valid IP address
+    # Examples:
+    # 192.168.144.1 (TP-Link 123)
+    # 192.168.144.1
+    ip_address = gateway.split()[0]  # Get the first part before any space
+    try:
+        ipaddress.ip_address(ip_address)
+        return ip_address
+    except ValueError:
+        return ""
+
+
 def get_gateway(interface, renderer):
-    gateway = None
     info = get_interface_info(interface, renderer)
-    gateway = info.get("gateway", None)
+    gateway = info.get("gateway") or ""
+    validated_gateway = _validate_gateway_ip(gateway)
     print("Got gateway address: {}".format(gateway))
-    return gateway
+    print("Validated gateway address: {}".format(validated_gateway))
+    return validated_gateway
 
 
 def perform_ping_test(interface, renderer):
diff --git a/providers/base/tests/test_wifi_client_test_netplan.py b/providers/base/tests/test_wifi_client_test_netplan.py
index d492c627a..60d28b35f 100644
--- a/providers/base/tests/test_wifi_client_test_netplan.py
+++ b/providers/base/tests/test_wifi_client_test_netplan.py
@@ -444,6 +444,18 @@ def test_get_interface_info_networkd(self, mock_check_output):
         self.assertEqual(info["state"], "routable")
         self.assertEqual(info["gateway"], "192.168.1.1")
 
+    @patch("subprocess.check_output")
+    def test_get_interface_info_networkd_any_name(self, mock_check_output):
+        mock_check_output.return_value = (
+            b"State: routable\nGateway: 192.168.1.1 (TP-Link 123)\n"
+            b"Path: pci-0000:02:00.0"
+        )
+        interface = "wlan0"
+        renderer = "networkd"
+        info = get_interface_info(interface, renderer)
+        self.assertEqual(info["state"], "routable")
+        self.assertEqual(info["gateway"], "192.168.1.1 (TP-Link 123)")
+
     @patch("subprocess.check_output")
     def test_get_interface_info_networkd_no_state(self, mock_check_output):
         mock_check_output.return_value = (
@@ -665,6 +677,21 @@ def test_gateway_found(self, mock_get_interface_info):
             "Got gateway address: 192.168.1.1", captured_output.getvalue()
         )
 
+    @patch(
+        "wifi_client_test_netplan.get_interface_info",
+        return_value={"gateway": "192.168.1.1 (TP-Link 123)"},
+    )
+    def test_gateway_found_any_name(self, mock_get_interface_info):
+        captured_output = io.StringIO()
+        sys.stdout = captured_output
+        result = get_gateway("wlan0", "networkd")
+        sys.stdout = sys.__stdout__
+        self.assertEqual(result, "192.168.1.1")
+        mock_get_interface_info.assert_called_once_with("wlan0", "networkd")
+        self.assertIn(
+            "Got gateway address: 192.168.1.1", captured_output.getvalue()
+        )
+
     @patch(
         "wifi_client_test_netplan.get_interface_info",
         return_value={},
@@ -674,9 +701,22 @@ def test_gateway_not_found(self, mock_get_interface_info):
         sys.stdout = captured_output
         result = get_gateway("wlan0", "networkd")
         sys.stdout = sys.__stdout__
-        self.assertIsNone(result)
+        self.assertEqual(result, "")
         mock_get_interface_info.assert_called_once_with("wlan0", "networkd")
-        self.assertIn("Got gateway address: None", captured_output.getvalue())
+        self.assertIn("Got gateway address: ", captured_output.getvalue())
+
+    @patch(
+        "wifi_client_test_netplan.get_interface_info",
+        return_value={"gateway": "192.168.1.1.100 (TP-Link 123)"},
+    )
+    def test_gateway_ip_not_valid(self, mock_get_interface_info):
+        captured_output = io.StringIO()
+        sys.stdout = captured_output
+        result = get_gateway("wlan0", "networkd")
+        sys.stdout = sys.__stdout__
+        self.assertEqual(result, "")
+        mock_get_interface_info.assert_called_once_with("wlan0", "networkd")
+        self.assertIn("Got gateway address: ", captured_output.getvalue())
 
     @patch(
         "wifi_client_test_netplan.get_interface_info",