Refactor WiFi password validation and add Linux support

Scripts/wifi_profile_extractor.py

@@ -10,19 +10,43 @@ class WifiProfileExtractor:
         self.run = run.Run().run
         self.utils = utils.Utils()
 
-    def validate_wifi_password(self, password):
-        if not password:
-            return False
+    def get_authentication_type(self, authentication_type):
+        authentication_type = authentication_type.lower()
+
+        open_types = ("none", "owe", "open")
+        for open_type in open_types:
+            if open_type in authentication_type:
+                return "open"
+
+        if "wep" in authentication_type or "shared" in authentication_type:
+            return "wep"
         
+        if "wpa" in authentication_type or "sae" in authentication_type:
+            return "wpa"
+        
+        return None
+
+    def validate_wifi_password(self, authentication_type=None, password=None):
+        print("Validating password with authentication type: {}".format(authentication_type))
+
+        if password is None:
+            return None
+
+        if authentication_type is None:
+            return password
+
+        if authentication_type == "open":
+            return ""
+
         try:
             password.encode('ascii')
         except UnicodeEncodeError:
-            return False
-        
-        if 8 <= len(password) <= 63 and all(32 <= ord(c) <= 126 for c in password):
-            return True
+            return None
             
-        return False
+        if 8 <= len(password) <= 63 and all(32 <= ord(c) <= 126 for c in password):
+            return password
+                
+        return None
 
     def get_wifi_password_macos(self, ssid):
         output = self.run({
@@ -38,10 +62,7 @@ class WifiProfileExtractor:
         except:
             password = output[0].strip() if output[0].strip() else None
             
-        if password and self.validate_wifi_password(password):
-            return password
-        
-        return None
+        return self.validate_wifi_password("wpa", password)
         
     def get_wifi_password_windows(self, ssid):
         output = self.run({
@@ -51,13 +72,35 @@ class WifiProfileExtractor:
         if output[-1] != 0:
             return None
 
+        authentication_type = None
+        password = None
+
         for line in output[0].splitlines():
-            if "Key Content" in line:
+            if authentication_type is None and "Authentication" in line:
+                authentication_type = self.get_authentication_type(line.split(":")[1].strip())
+            elif "Key Content" in line:
                 password = line.split(":")[1].strip()
-                if self.validate_wifi_password(password):
-                    return password
+
+        return self.validate_wifi_password(authentication_type, password)
+
+    def get_wifi_password_linux(self, ssid):
+        output = self.run({
+            "args": ["nmcli", "--show-secrets", "connection", "show", ssid]
+        })
+
+        if output[-1] != 0:
+            return None
         
-        return None
+        authentication_type = None
+        password = None
+
+        for line in output[0].splitlines():
+            if "802-11-wireless-security.key-mgmt:" in line:
+                authentication_type = self.get_authentication_type(line.split(":")[1].strip())
+            elif "802-11-wireless-security.psk:" in line:
+                password = line.split(":")[1].strip()
+
+        return self.validate_wifi_password(authentication_type, password)
 
     def ask_network_count(self, total_networks):
         self.utils.head("WiFi Network Retrieval")
@@ -65,7 +108,7 @@ class WifiProfileExtractor:
         print("Found {} WiFi networks on this device.".format(total_networks))
         print("")
         print("How many networks would you like to process?")
-        print("  1-5 - Specific number (default: 5)")
+        print("  1-{} - Specific number (default: 5)".format(total_networks))
         print("  A   - All available networks")
         print("")
         
@@ -74,15 +117,15 @@ class WifiProfileExtractor:
         if num_choice == "a":
             print("Will process all available networks.")
             return total_networks
-        else:
-            try:
-                max_networks = min(int(num_choice), total_networks)
-                print("Will process up to {} networks.".format(max_networks))
-                return max_networks
-            except:
-                max_networks = min(5, total_networks)
-                print("Invalid choice. Will process up to {} networks.".format(max_networks))
-                return max_networks
+
+        try:
+            max_networks = min(int(num_choice), total_networks)
+            print("Will process up to {} networks.".format(max_networks))
+            return max_networks
+        except:
+            max_networks = min(5, total_networks)
+            print("Invalid choice. Will process up to {} networks.".format(max_networks))
+            return max_networks
             
     def process_networks(self, ssid_list, max_networks, get_password_func):
         networks = []
@@ -97,25 +140,25 @@ class WifiProfileExtractor:
                 print("")
                 print("Processing {}/{}: {}".format(processed_count + 1, len(ssid_list), ssid))
                 if os_name == "Darwin":
-                    print("Please enter your administrator name and password or click 'Deny' to skip this network...")
+                    print("Please enter your administrator name and password or click 'Deny' to skip this network.")
                 
                 password = get_password_func(ssid)
-                if password:
+                if password is not None:
                     if (ssid, password) not in networks:
                         consecutive_failures = 0
                         networks.append((ssid, password))
-                        print("Successfully retrieved password for {}".format(ssid))
+                        print("Successfully retrieved password.")
                         
                         if len(networks) == max_networks:
                             break
                 else:
                     consecutive_failures += 1 if os_name == "Darwin" else 0
-                    print("Skipped or could not retrieve password for {}".format(ssid))
+                    print("Could not retrieve password for this network.")
 
                     if consecutive_failures >= max_consecutive_failures:
-                        continue_input = self.utils.request_input("\nUnable to retrieve passwords. Continue trying? (Y/n): ").strip().lower() or "y"
+                        continue_input = self.utils.request_input("\nUnable to retrieve passwords. Continue trying? (Yes/no): ").strip().lower() or "yes"
                         
-                        if continue_input != "y":
+                        if continue_input != "yes":
                             break
 
                         consecutive_failures = 0
@@ -124,9 +167,9 @@ class WifiProfileExtractor:
                 print("Error processing network '{}': {}".format(ssid, str(e)))
 
                 if consecutive_failures >= max_consecutive_failures:
-                    continue_input = self.utils.request_input("\nUnable to retrieve passwords. Continue trying? (Y/n): ").strip().lower() or "y"
+                    continue_input = self.utils.request_input("\nUnable to retrieve passwords. Continue trying? (Yes/no): ").strip().lower() or "yes"
 
-                    if continue_input != "y":
+                    if continue_input != "yes":
                         break
 
                     consecutive_failures = 0
@@ -134,9 +177,9 @@ class WifiProfileExtractor:
                 processed_count += 1
             
             if processed_count >= max_networks and len(networks) < max_networks and processed_count < len(ssid_list):
-                continue_input = self.utils.request_input("\nOnly retrieved {}/{} networks. Try more to reach your target? (Y/n): ".format(len(networks), max_networks)).strip().lower() or "y"
+                continue_input = self.utils.request_input("\nOnly retrieved {}/{} networks. Try more to reach your target? (Yes/no): ".format(len(networks), max_networks)).strip().lower() or "yes"
                 
-                if continue_input != "y":
+                if continue_input != "yes":
                     break
 
                 consecutive_failures = 0
@@ -186,15 +229,36 @@ class WifiProfileExtractor:
         
         if not ssid_list:
             return []
-            
-        max_networks = self.ask_network_count(len(ssid_list))
+
+        max_networks = len(ssid_list)
     
         self.utils.head("WiFi Profile Extractor")
         print("")
-        print("Ready to retrieve passwords for networks.")
+        print("Retrieving passwords for {} network(s)...".format(len(ssid_list)))
         
         return self.process_networks(ssid_list, max_networks, self.get_wifi_password_windows)
 
+    def get_preferred_networks_linux(self):
+        output = self.run({
+            "args": ["nmcli", "-t", "-f", "NAME", "connection", "show"]
+        })
+
+        if output[-1] != 0:
+            return []
+        
+        ssid_list = [network.strip() for network in output[0].splitlines() if network.strip()]
+        
+        if not ssid_list:
+            return []
+            
+        max_networks = len(ssid_list)
+    
+        self.utils.head("WiFi Profile Extractor")
+        print("")
+        print("Retrieving passwords for {} network(s)...".format(len(ssid_list)))
+        
+        return self.process_networks(ssid_list, max_networks, self.get_wifi_password_linux)
+
     def get_wifi_interfaces(self):
         output = self.run({
             "args": ["networksetup", "-listallhardwareports"]
@@ -250,6 +314,8 @@ class WifiProfileExtractor:
         
         if os_name == "Windows":
             profiles = self.get_preferred_networks_windows()
+        elif os_name == "Linux":
+            profiles = self.get_preferred_networks_linux()
         elif os_name == "Darwin":
             wifi_interfaces = self.get_wifi_interfaces()
 
@@ -273,10 +339,10 @@ class WifiProfileExtractor:
         print("")
         print("Found the following WiFi profiles with saved passwords:")
         print("")
-        print("Index  SSID                    Password")
+        print("Index  SSID                             Password")
         print("-------------------------------------------------------")
         for index, (ssid, password) in enumerate(profiles, start=1):
-            print("{:<6} {:<23} {}".format(index, ssid[:23], password))
+            print("{:<6} {:<32} {:<8}".format(index, ssid[:31] + "..." if len(ssid) > 31 else ssid, password[:12] + "..." if len(password) > 12 else password))
         print("")
         print("Successfully applied {} WiFi profiles.".format(len(profiles)))
         print("")