Improved device selection logic

services/agenthandler.py

@@ -130,24 +130,16 @@ def findAgents(api, return_dataframe):
     else:
         logging.debug("User declined to export the DataFrame.")
 
-
-def selectAgents(api: AirlockAPIWrapper) -> List[Agent]:
+def collect_device_names() -> List[str]:
     print(colorText("🔍 Device Search", "cyan"))
     print(colorText("Enter the device hostnames you'd like to search for, one per line.", "cyan"))
     print(colorText("When you're done, press Enter twice (Three times if you have a single device).\n", "cyan"))
     print(colorText("Example:", "cyan"))
-    print(colorText("H00000", "cyan"))
-    print(colorText("UTN00000", "cyan"))
-    print(colorText("i-hSuperSecretServer", "cyan"))
-    print(colorText("u-hVenderBroke\n", "cyan"))
-
+    print(colorText("H00000\nUTN00000\ni-hSuperSecretServer\nu-hVenderBroke\n", "cyan"))
     print(colorText("Paste or type your device names below:", "white"))
-    policies = [Policy(**row.to_dict()) for _, row in api.policy_find_all().iterrows()]
 
     device_input_lines = []
     empty_line_count = 0
-
-    # Regex to validate each line
     valid_line_pattern = re.compile(r'^[a-zA-Z0-9_\- ]+$')
 
     while True:
@@ -158,49 +150,97 @@ def selectAgents(api: AirlockAPIWrapper) -> List[Agent]:
             empty_line_count += 1
             if empty_line_count == 2:
                 break
-            continue  # Don't validate empty lines
+            continue
         else:
             empty_line_count = 0
 
-        # Validate only non-empty lines
         if valid_line_pattern.match(stripped_line):
             device_input_lines.append(stripped_line)
         else:
             print(colorText(f"⚠️ Invalid input: '{stripped_line}' — only letters, numbers, underscores, spaces, and hyphens are allowed.", "yellow"))
-            
-    device_names = [name for name in device_input_lines if name]
+
+    return [name for name in device_input_lines if name]
+
+
+def choose_match_type() -> bool:
+    print(colorText("Use exact match? (Y for exact, N for fuzzy):", "white"))
+    return get_sanitized_input("").strip().lower() in ["y", "yes"]
+
+
+def match_agents(device_names: List[str], agents: List['Agent'], use_exact: bool) -> List['Agent']:
+    if use_exact:
+        return [
+            agent for agent in agents
+            if agent.hostname.lower() in [name.lower() for name in device_names]
+        ]
+    else:
+        pattern = "|".join(map(re.escape, device_names))
+        regex = re.compile(pattern, re.IGNORECASE)
+        return [agent for agent in agents if regex.search(agent.hostname)]
+
+
+def show_unmatched(device_names: List[str], matched_agents: List['Agent'], use_exact: bool):
+    if use_exact:
+        unmatched = [name for name in device_names if not any(agent.hostname.lower() == name.lower() for agent in matched_agents)]
+    else:
+        unmatched = [name for name in device_names if not any(re.search(re.escape(name), agent.hostname, re.IGNORECASE) for agent in matched_agents)]
+
+    if unmatched:
+        logger.debug(f"⚠️ No matches for: {', '.join(unmatched)}")
+        print(colorText(f"⚠️ No matches for: {', '.join(unmatched)}", "yellow"))
+
+
+def enrich_agents(agents: List['Agent'], policies: List['Policy']):
+    for agent in agents:
+        agent.enrich_with_policies(policies)
+
+
+def selectAgents(api: 'AirlockAPIWrapper') -> List['Agent']:
+    device_names = collect_device_names()
     if not device_names:
         logger.debug("No device names entered")
         print(colorText("⚠️ No device names entered.", "red"))
         return []
 
-    # Build regex pattern to match hostnames
-    pattern = "|".join(map(re.escape, device_names))
-    regex = re.compile(pattern, re.IGNORECASE)
+    use_exact = choose_match_type()
 
-    # Fetch agents
+    policies = [Policy(**row.to_dict()) for _, row in api.policy_find_all().iterrows()]
     agents = [Agent(**row.to_dict()) for _, row in api.agent_find_all().iterrows()]
-    matched_agents = [agent for agent in agents if regex.search(agent.hostname)]
+    matched_agents = match_agents(device_names, agents, use_exact)
     matched_agents.sort(key=lambda agent: agent.hostname.lower())
 
-    # Show unmatched
-    unmatched = [name for name in device_names if not any(regex.search(agent.hostname) for agent in agents)]
-    if unmatched:
-        logger.debug(f"⚠️ No matches for: {', '.join(unmatched)}")
-        print(colorText(f"⚠️ No matches for: {', '.join(unmatched)}", "yellow"))
+    show_unmatched(device_names, matched_agents, use_exact)
 
     if not matched_agents:
         logger.debug("❌ No matching devices found.")
         print(colorText("❌ No matching devices found.", "red"))
-    else:
-        logger.debug(f"✅ Found {len(matched_agents)} matching device(s).")        
-        print(colorText(f"✅ Found {len(matched_agents)} matching device(s).", "green"))
+        return []
 
-    # Enrich each agent using its class method
-    for agent in matched_agents:
-        agent.enrich_with_policies(policies)
+    print(colorText(f"✅ Found {len(matched_agents)} matching device(s).", "green"))
+    logger.info("Matched agent hostnames:")
+    rows = (len(matched_agents) + 2) // 3  # 3 columns
+    for row in range(rows):
+        line = ""
+        for col in range(3):
+            idx = row + col * rows
+            if idx < len(matched_agents):
+                line += f"{matched_agents[idx].hostname:<30}"
+        logger.info(line)
+        
+    matched_agents = Selector.select_with_mode(
+        matched_agents,
+        label_func=lambda agent: agent.hostname,
+        header="Matched Devices:"
+    )
 
+    if not matched_agents:
+        logger.debug("❌ No matching devices remain after refinement.")
+        print(colorText("❌ No matching devices remain after refinement.", "red"))
+        return []
+
+    enrich_agents(matched_agents, policies)
     return matched_agents
+
 def moveAgentToRelatedPolicy(
     api: AirlockAPIWrapper,
     agent: Agent,