Update Netmiko cli_tools to make them more modular and use concurrent…

… futures (#3500)

netmiko/cli_tools/__init__.py

@@ -0,0 +1,10 @@
+import os
+
+
+__version__ = "5.0.0"
+MAX_WORKERS = int(os.environ.get("NETMIKO_MAX_THREADS", 10))
+ERROR_PATTERN = "%%%failed%%%"
+
+GREP = "/bin/grep"
+if not os.path.exists(GREP):
+    GREP = "/usr/bin/grep"

netmiko/cli_tools/cli_helpers.py

@@ -0,0 +1,60 @@
+from typing import Any, Dict
+from netmiko import ConnectHandler
+from netmiko.utilities import load_devices, obtain_all_devices
+from netmiko.cli_tools import ERROR_PATTERN
+
+
+def ssh_conn(device_name, device_params, cli_command=None, cfg_command=None):
+    try:
+        output = ""
+        with ConnectHandler(**device_params) as net_connect:
+            net_connect.enable()
+            if cli_command:
+                output += net_connect.send_command(cli_command)
+            if cfg_command:
+                output += net_connect.send_config_set(cfg_command)
+            return device_name, output
+    except Exception:
+        return device_name, ERROR_PATTERN
+
+
+def obtain_devices(device_or_group: str) -> Dict[str, Dict[str, Any]]:
+    """
+    Obtain the devices from the .netmiko.yml file using either a group-name or
+    a device-name. A group-name will be a list of device-names. A device-name
+    will just be a dictionary of device parameters (ConnectHandler **kwargs).
+    """
+    my_devices = load_devices()
+    if device_or_group == "all":
+        devices = obtain_all_devices(my_devices)
+    else:
+        try:
+            singledevice_or_group = my_devices[device_or_group]
+            devices = {}
+            if isinstance(singledevice_or_group, list):
+                # Group of Devices
+                device_group = singledevice_or_group
+                for device_name in device_group:
+                    devices[device_name] = my_devices[device_name]
+            else:
+                # Single Device (dictionary)
+                device_name = device_or_group
+                device_dict = my_devices[device_name]
+                devices[device_name] = device_dict
+        except KeyError:
+            return (
+                "Error reading from netmiko devices file."
+                " Device or group not found: {0}".format(device_or_group)
+            )
+
+    return devices
+
+
+def update_device_params(params, username=None, password=None, secret=None):
+    if username:
+        params["username"] = username
+    if password:
+        params["password"] = password
+    if secret:
+        params["secret"] = secret
+    return params

netmiko/cli_tools/netmiko_cfg.py

@@ -1,44 +1,18 @@
 #!/usr/bin/env python
 """Return output from single show cmd using Netmiko."""
-from __future__ import print_function
-from __future__ import unicode_literals
-
 import argparse
 import sys
 import os
 import subprocess
-import threading
-
-try:
-    from Queue import Queue
-except ImportError:
-    from queue import Queue
 from datetime import datetime
 from getpass import getpass
+from concurrent.futures import ThreadPoolExecutor, as_completed
 
-from netmiko import ConnectHandler
 from netmiko.utilities import load_devices, display_inventory
-from netmiko.utilities import obtain_all_devices
 from netmiko.utilities import write_tmp_file, ensure_dir_exists
 from netmiko.utilities import find_netmiko_dir
-from netmiko.utilities import SHOW_RUN_MAPPER
-
-GREP = "/bin/grep"
-if not os.path.exists(GREP):
-    GREP = "/usr/bin/grep"
-NETMIKO_BASE_DIR = "~/.netmiko"
-ERROR_PATTERN = "%%%failed%%%"
-__version__ = "0.1.0"
-
-PY2 = sys.version_info.major == 2
-PY3 = sys.version_info.major == 3
-
-if sys.version_info.major == 3:
-    string_types = (str,)
-    text_type = str
-else:
-    string_types = (basestring,)  # noqa
-    text_type = unicode  # noqa
+from netmiko.cli_tools import ERROR_PATTERN, GREP, MAX_WORKERS, __version__
+from netmiko.cli_tools.cli_helpers import obtain_devices, update_device_params, ssh_conn
 
 
 def grepx(files, pattern, grep_options, use_colors=True):
@@ -61,19 +35,6 @@ def grepx(files, pattern, grep_options, use_colors=True):
     return ""
 
 
-def ssh_conn(device_name, a_device, cfg_command, output_q):
-    try:
-        net_connect = ConnectHandler(**a_device)
-        net_connect.enable()
-        if isinstance(cfg_command, string_types):
-            cfg_command = [cfg_command]
-        output = net_connect.send_config_set(cfg_command)
-        net_connect.disconnect()
-    except Exception:
-        output = ERROR_PATTERN
-    output_q.put({device_name: output})
-
-
 def parse_arguments(args):
     """Parse command-line arguments."""
     description = "Execute single config cmd using Netmiko"
@@ -139,76 +100,62 @@ def main(args):
         display_inventory(my_devices)
         return 0
 
-    cli_command = cli_args.cmd
-    cmd_arg = False
-    if cli_command:
-        cmd_arg = True
-        if r"\n" in cli_command:
-            cli_command = cli_command.strip()
-            cli_command = cli_command.split(r"\n")
+    cfg_command = cli_args.cmd
+    if cfg_command:
+        if r"\n" in cfg_command:
+            cfg_command = cfg_command.strip()
+            cfg_command = cfg_command.split(r"\n")
     elif input:
-        cmd_arg = True
         command_data = cli_args.infile.read()
         command_data = command_data.strip()
-        cli_command = command_data.splitlines()
+        cfg_command = command_data.splitlines()
     else:
         raise ValueError("No configuration commands provided.")
     device_or_group = cli_args.devices.strip()
     pattern = r"."
     hide_failed = cli_args.hide_failed
 
-    output_q = Queue()
-    my_devices = load_devices()
-    if device_or_group == "all":
-        device_group = obtain_all_devices(my_devices)
-    else:
-        try:
-            devicedict_or_group = my_devices[device_or_group]
-            device_group = {}
-            if isinstance(devicedict_or_group, list):
-                for tmp_device_name in devicedict_or_group:
-                    device_group[tmp_device_name] = my_devices[tmp_device_name]
-            else:
-                device_group[device_or_group] = devicedict_or_group
-        except KeyError:
-            return (
-                "Error reading from netmiko devices file."
-                " Device or group not found: {0}".format(device_or_group)
-            )
+    # DEVICE LOADING #####
+    devices = obtain_devices(device_or_group)
 
     # Retrieve output from devices
     my_files = []
     failed_devices = []
-    for device_name, a_device in device_group.items():
-        if cli_username:
-            a_device["username"] = cli_username
-        if cli_password:
-            a_device["password"] = cli_password
-        if cli_secret:
-            a_device["secret"] = cli_secret
-        if not cmd_arg:
-            cli_command = SHOW_RUN_MAPPER.get(a_device["device_type"], "show run")
-        my_thread = threading.Thread(
-            target=ssh_conn, args=(device_name, a_device, cli_command, output_q)
+    results = {}
+
+    # UPDATE DEVICE PARAMS (WITH CLI ARGS) / Create Task List #####
+    device_tasks = []
+    for device_name, device_params in devices.items():
+        update_device_params(
+            device_params,
+            username=cli_username,
+            password=cli_password,
+            secret=cli_secret,
         )
-        my_thread.start()
-    # Make sure all threads have finished
-    main_thread = threading.current_thread()
-    for some_thread in threading.enumerate():
-        if some_thread != main_thread:
-            some_thread.join()
-    # Write files
-    while not output_q.empty():
-        my_dict = output_q.get()
-        netmiko_base_dir, netmiko_full_dir = find_netmiko_dir()
-        ensure_dir_exists(netmiko_base_dir)
-        ensure_dir_exists(netmiko_full_dir)
-        for device_name, output in my_dict.items():
-            file_name = write_tmp_file(device_name, output)
-            if ERROR_PATTERN not in output:
-                my_files.append(file_name)
-            else:
-                failed_devices.append(device_name)
+        device_tasks.append(
+            {
+                "device_name": device_name,
+                "device_params": device_params,
+                "cfg_command": cfg_command,
+            }
+        )
+
+    # THREADING #####
+    with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
+        futures = [executor.submit(ssh_conn, **kwargs) for kwargs in device_tasks]
+        for future in as_completed(futures):
+            device_name, output = future.result()
+            results[device_name] = output
+
+    netmiko_base_dir, netmiko_full_dir = find_netmiko_dir()
+    ensure_dir_exists(netmiko_base_dir)
+    ensure_dir_exists(netmiko_full_dir)
+    for device_name, output in results.items():
+        file_name = write_tmp_file(device_name, output)
+        if ERROR_PATTERN not in output:
+            my_files.append(file_name)
+        else:
+            failed_devices.append(device_name)
 
     grep_options = []
     grepx(my_files, pattern, grep_options)