Command getter refactor

Author: mzbroch

nautobot_device_onboarding/nornir_plays/command_getter.py

@@ -91,154 +91,181 @@ def _get_commands_to_run(yaml_parsed_info, skip_list=None):
     return deduplicate_command_list(all_commands)
 
 
-def netmiko_send_commands(task: Task,
-                          command_getter_yaml_data: Dict,
-                          command_getter_job: str,
-                          logger,
-                          nautobot_job
-                          ):
-    """Run commands specified in PLATFORM_COMMAND_MAP."""
-    if not task.host.platform:
-        return Result(host=task.host, result=f"{task.host.name} has no platform set.", failed=True)
-    if task.host.platform not in get_all_network_driver_mappings().keys() or not "cisco_wlc_ssh":
-        return Result(host=task.host, result=f"{task.host.name} has a unsupported platform set.", failed=True)
-    if not command_getter_yaml_data[task.host.platform].get(command_getter_job):
+def netmiko_send_commands(
+    task: Task,
+    command_getter_yaml_data: Dict,
+    command_getter_job: str,
+    logger,
+    nautobot_job,
+    **kwargs,
+):
+    """Run platform-specific commands with optional parsing and logging."""
+
+    command_exclusions = kwargs.get("command_exclusions", [])
+    connectivity_test = kwargs.get("connectivity_test", False)
+    # sync_cables = kwargs.get("sync_cables", False)
+
+    # ---- 1. Validation -------------------------------------------------------
+    validation_result = _validate_task(task, command_getter_yaml_data, command_getter_job)
+    if validation_result:
+        return validation_result
+
+    if connectivity_test and not _check_connectivity(task):
         return Result(
-            host=task.host, result=f"{task.host.name} has missing definitions in command_mapper YAML file.", failed=True
+            host=task.host,
+            result=f"{task.host.name} failed connectivity check via tcp_ping.",
+            failed=True,
         )
-    if nautobot_job.connectivity_test:
-        if not tcp_ping(task.host.hostname, task.host.port):
-            return Result(
-                host=task.host, result=f"{task.host.name} failed connectivity check via tcp_ping.", failed=True
-            )
-    task.host.data["platform_parsing_info"] = command_getter_yaml_data[task.host.platform]
 
-    skip_conditions = {
-        "interfaces__tagged_vlans": not sync_vlans,
-        "vlan_map": not sync_vlans,
-        "interfaces__untagged_vlan": not sync_vlans,
-        "interfaces__vrf": not sync_vrfs,
-        "cables": not sync_cables,
-        "software_version": not sync_software_version,
-    }
+    task.host.data["platform_parsing_info"] = command_getter_yaml_data[task.host.platform]
 
+    # ---- 2. Command Preparation ---------------------------------------------
     commands = _get_commands_to_run(
-        command_getter_yaml_data[task.host.platform][command_getter_job],
-        # getattr(nautobot_job, "sync_vlans", False),
-        # getattr(nautobot_job, "sync_vrfs", False),
-        # getattr(nautobot_job, "sync_cables", False),
-        # getattr(nautobot_job, "sync_software_version", False),
+        yaml_parsed_info=command_getter_yaml_data[task.host.platform][command_getter_job],
+        skip_list=command_exclusions,
     )
+    logger.debug(f"Commands to run: {[cmd['command'] for cmd in commands]}")
+
+    # ---- 3. Command Execution & Parsing -------------------------------------
+    for idx, command in enumerate(commands):
+        try:
+            current_result = _run_command(task, command)
+            _handle_command_result(task, idx, command, current_result, nautobot_job, logger)
+        except NornirSubTaskError as e:
+            return _handle_subtask_error(task, idx, e)
+
+    return Result(host=task.host, result="Commands executed successfully.", failed=False)
+
+
+# -----------------------------------------------------------------------------
+# Helper functions
+# -----------------------------------------------------------------------------
+
+def _validate_task(task, yaml_data, job):
+    """Check platform and YAML definition validity."""
+    platform = task.host.platform
+    if not platform:
+        return Result(host=task.host, result=f"{task.host.name} has no platform set.", failed=True)
+
+    supported_platforms = get_all_network_driver_mappings().keys()
+    if platform not in supported_platforms and platform != "cisco_wlc_ssh":
+        return Result(host=task.host, result=f"{task.host.name} has an unsupported platform.", failed=True)
 
-    if (
-        getattr(nautobot_job, "sync_cables", False)
-        and "cables" not in command_getter_yaml_data[task.host.platform][command_getter_job].keys()
-    ):
-        logger.error(
-            f"{task.host.platform} has missing definitions for cables in command_mapper YAML file. Cables will not be loaded."
+    if not yaml_data.get(platform, {}).get(job):
+        return Result(
+            host=task.host,
+            result=f"{task.host.name} missing definitions in command_mapper YAML.",
+            failed=True,
         )
 
-    logger.debug(f"Commands to run: {[cmd['command'] for cmd in commands]}")
-    # All commands in this for loop are running within 1 device connection.
-    for result_idx, command in enumerate(commands):
-        send_command_kwargs = {}
+    return None
+
+
+def _check_connectivity(task):
+    """Perform TCP ping check."""
+    return tcp_ping(task.host.hostname, task.host.port)
+
+
+def _run_command(task, command):
+    """Execute a single Netmiko command."""
+    return task.run(
+        task=netmiko_send_command,
+        name=command["command"],
+        command_string=command["command"],
+        read_timeout=60,
+    )
+
+
+def _handle_command_result(task, idx, command, current_result, nautobot_job, logger):
+    """Parse and store results based on parser type."""
+    raw_output = current_result.result
+    parser_type = command.get("parser")
+
+    # Debug output
+    if nautobot_job.debug:
+        log_message = format_log_message(pprint.pformat(raw_output))
+        logger.debug(f"Result of '{command['command']}' command:<br><br>{log_message}")
+
+    # Handle invalid input gracefully
+    if isinstance(raw_output, str) and "Invalid input detected" in raw_output:
+        task.results[idx].result = []
+        task.results[idx].failed = False
+        return
+
+    if parser_type in SUPPORTED_COMMAND_PARSERS:
+        parsed = _parse_command_output(task, command, raw_output, parser_type, nautobot_job, logger)
+        task.results[idx].result = parsed
+        task.results[idx].failed = False
+    else:
+        task.results[idx].result = _handle_raw_or_none(raw_output, parser_type)
+        task.results[idx].failed = False
+
+
+def _parse_command_output(task, command, raw_output, parser_type, nautobot_job, logger):
+    """Dispatch to the appropriate parser."""
+    try:
+        if parser_type == "textfsm":
+            return _parse_textfsm(task, command, raw_output, nautobot_job, logger)
+        elif parser_type == "ttp":
+            return _parse_ttp(task, command, raw_output)
+    except Exception as e:
+        logger.warning(f"Parsing failed for {command['command']}: {e}")
+        return []
+    return []
+
+
+def _parse_textfsm(task, command, data, nautobot_job, logger):
+    git_template_dir = get_git_repo_parser_path("textfsm")
+    if git_template_dir and not check_for_required_file(git_template_dir, "index"):
+        logger.debug(f"Missing index file in {git_template_dir}, falling back to defaults.")
+        git_template_dir = None
+
+    parsed_output = parse_output(
+        platform=get_all_network_driver_mappings()[task.host.platform]["ntc_templates"],
+        template_dir=git_template_dir,
+        command=command["command"],
+        data=data,
+        try_fallback=bool(git_template_dir),
+    )
+
+    if nautobot_job.debug:
+        logger.debug(format_log_message(pprint.pformat(parsed_output)))
+
+    return parsed_output
+
+
+def _parse_ttp(task, command, data):
+    ttp_template_files = load_files_with_precedence(f"{PARSER_DIR}/ttp", "ttp")
+    template_name = f"{task.host.platform}_{command['command'].replace(' ', '_')}.ttp"
+    parser = ttp(data=data, template=ttp_template_files[template_name])
+    parser.parse()
+    return json.loads(parser.result(format="json")[0])
+
+
+def _handle_raw_or_none(raw_output, parser_type):
+    if parser_type == "raw":
+        return {"raw": raw_output}
+    if parser_type == "none":
         try:
-            current_result = task.run(
-                task=netmiko_send_command,
-                name=command["command"],
-                command_string=command["command"],
-                read_timeout=60,
-                **send_command_kwargs,
-            )
-            if nautobot_job.debug:
-                log_message = format_log_message(pprint.pformat(current_result.result))
-                logger.debug(f"Result of '{command['command']}' command:<br><br>{log_message}")
-            if command.get("parser") in SUPPORTED_COMMAND_PARSERS:
-                if isinstance(current_result.result, str):
-                    if "Invalid input detected at" in current_result.result:
-                        task.results[result_idx].result = []
-                        task.results[result_idx].failed = False
-                    else:
-                        if command["parser"] == "textfsm":
-                            try:
-                                # Look for custom textfsm templates in the git repo
-                                git_template_dir = get_git_repo_parser_path(parser_type="textfsm")
-                                if git_template_dir:
-                                    if not check_for_required_file(git_template_dir, "index"):
-                                        logger.debug(
-                                            f"Unable to find required index file in {git_template_dir} for textfsm parsing. Falling back to default templates."
-                                        )
-                                        git_template_dir = None
-                                # Parsing textfsm ourselves instead of using netmikos use_<parser> function to be able to handle exceptions
-                                # ourselves. Default for netmiko is if it can't parse to return raw text which is tougher to handle.
-                                parsed_output = parse_output(
-                                    platform=get_all_network_driver_mappings()[task.host.platform]["ntc_templates"],
-                                    template_dir=git_template_dir if git_template_dir else None,
-                                    command=command["command"],
-                                    data=current_result.result,
-                                    try_fallback=bool(git_template_dir),
-                                )
-                                if nautobot_job.debug:
-                                    log_message = format_log_message(pprint.pformat(parsed_output))
-                                    logger.debug(
-                                        f"Parsed output of '{command['command']}' command:<br><br>{log_message}"
-                                    )
-                                task.results[result_idx].result = parsed_output
-                                task.results[result_idx].failed = False
-                            except Exception:  # https://github.com/networktocode/ntc-templates/issues/369
-                                try:
-                                    if nautobot_job.debug:
-                                        traceback_str = traceback.format_exc().replace("\n", "<br>")
-                                        logger.warning(
-                                            f"Parsing failed for '{command['command']}' command:<br><br>{traceback_str}"
-                                        )
-                                except:  # noqa: E722, S110
-                                    pass
-                                task.results[result_idx].result = []
-                                task.results[result_idx].failed = False
-                        if command["parser"] == "ttp":
-                            try:
-                                # Parsing ttp ourselves instead of using netmikos use_<parser> function to be able to handle exceptions
-                                # ourselves.
-                                ttp_template_files = load_files_with_precedence(
-                                    filesystem_dir=f"{PARSER_DIR}/ttp", parser_type="ttp"
-                                )
-                                template_name = f"{task.host.platform}_{command['command'].replace(' ', '_')}.ttp"
-                                parser = ttp(
-                                    data=current_result.result,
-                                    template=ttp_template_files[template_name],
-                                )
-                                parser.parse()
-                                parsed_result = parser.result(format="json")[0]
-                                # task.results[result_idx].result = json.loads(json.dumps(parsed_result))
-                                task.results[result_idx].result = json.loads(parsed_result)
-                                task.results[result_idx].failed = False
-                            except Exception:
-                                task.results[result_idx].result = []
-                                task.results[result_idx].failed = False
-            else:
-                if command["parser"] == "raw":
-                    raw = {"raw": current_result.result}
-                    task.results[result_idx].result = json.loads(json.dumps(raw))
-                    task.results[result_idx].failed = False
-                if command["parser"] == "none":
-                    try:
-                        jsonified = json.loads(current_result.result)
-                        task.results[result_idx].result = jsonified
-                        task.results[result_idx].failed = False
-                    except Exception:
-                        task.result.failed = False
-        except NornirSubTaskError:
-            # These exceptions indicate that the device is unreachable or the credentials are incorrect.
-            # We should fail the task early to avoid trying all commands on a device that is unreachable.
-            if type(task.results[result_idx].exception).__name__ == "NetmikoAuthenticationException":
-                return Result(host=task.host, result=f"{task.host.name} failed authentication.", failed=True)
-            if type(task.results[result_idx].exception).__name__ == "NetmikoTimeoutException":
-                return Result(host=task.host, result=f"{task.host.name} SSH Timeout Occured.", failed=True)
-            # We don't want to fail the entire subtask if SubTaskError is hit, set result to empty list and failed to False
-            # Handle this type or result latter in the ETL process.
-            task.results[result_idx].result = []
-            task.results[result_idx].failed = False
+            return json.loads(raw_output)
+        except Exception:
+            return []
+    return []
+
+
+def _handle_subtask_error(task, idx, exception):
+    """Handle connection/authentication errors gracefully."""
+    exc_type = type(task.results[idx].exception).__name__
+
+    if exc_type == "NetmikoAuthenticationException":
+        return Result(host=task.host, result=f"{task.host.name} failed authentication.", failed=True)
+    if exc_type == "NetmikoTimeoutException":
+        return Result(host=task.host, result=f"{task.host.name} SSH timeout occurred.", failed=True)
+
+    task.results[idx].result = []
+    task.results[idx].failed = False
+
+    return None
 
 
 @lru_cache(maxsize=None)
@@ -309,12 +336,14 @@ def sync_devices_command_getter(job, log_level):
                         )
                         continue
                 nr_with_processors.inventory.hosts.update(single_host_inventory_constructed)
+
             nr_with_processors.run(
                 task=netmiko_send_commands,
                 command_getter_yaml_data=nr_with_processors.inventory.defaults.data["platform_parsing_info"],
                 command_getter_job="sync_devices",
                 logger=logger,
                 nautobot_job=job,
+                # command_getter_yaml_exclusions=command_exclusions,
             )
     except Exception as err:  # pylint: disable=broad-exception-caught
         try:
@@ -348,21 +377,44 @@ def sync_network_data_command_getter(job, log_level):
                     "defaults": {
                         "platform_parsing_info": add_platform_parsing_info(),
                         "network_driver_mappings": list(get_all_network_driver_mappings().keys()),
-                        "sync_vlans": job.sync_vlans,
-                        "sync_vrfs": job.sync_vrfs,
-                        "sync_cables": job.sync_cables,
-                        "sync_software_version": job.sync_software_version,
+                        # "sync_vlans": job.sync_vlans,
+                        # "sync_vrfs": job.sync_vrfs,
+                        # "sync_cables": job.sync_cables,
+                        # "sync_software_version": job.sync_software_version,
                     },
                 },
             },
         ) as nornir_obj:
+
+            command_exclusions = {
+                "interfaces__tagged_vlans": not job.sync_vlans,
+                "vlan_map": not job.sync_vlans,
+                "interfaces__untagged_vlan": not job.sync_vlans,
+                "interfaces__vrf": not job.sync_vrfs,
+                "cables": not job.sync_cables,
+                "software_version": not job.sync_software_version,
+            }
+            exclusions = [k for k, v in command_exclusions.items() if v]
+
+            # commands = _get_commands_to_run(
+            #     command_getter_yaml_data[task.host.platform][command_getter_job],
+            #     # getattr(nautobot_job, "sync_vlans", False),
+            #     # getattr(nautobot_job, "sync_vrfs", False),
+            #     # getattr(nautobot_job, "sync_cables", False),
+            #     # getattr(nautobot_job, "sync_software_version", False),
+            # )
+
+
             nr_with_processors = nornir_obj.with_processors([CommandGetterProcessor(logger, compiled_results, job)])
             nr_with_processors.run(
                 task=netmiko_send_commands,
                 command_getter_yaml_data=nr_with_processors.inventory.defaults.data["platform_parsing_info"],
                 command_getter_job="sync_network_data",
                 logger=logger,
-                nautobot_job=job,
+                command_exclusions=exclusions,
+                connectivity_test=job.connectivity_test,
+                sync_cables=job.sync_cables,
+                # nautobot_job=job,
             )
     except Exception:  # pylint: disable=broad-exception-caught
         logger.info(f"Error During Sync Network Data Command Getter: {traceback.format_exc()}")