Netmiko

Packs/Netmiko/Integrations/Netmiko/Netmiko.py

@@ -14,12 +14,18 @@
 ''' HELPER FUNCTIONS '''
 
 
+# Return only specific keys from dictionary
+def include_keys(dictionary, keys):
+    key_set = set(keys) & set(dictionary.keys())
+    return {key: dictionary[key] for key in key_set}
+
+
 def return_file(keys):
     return_file.readlines = lambda: keys.split("\n")  # type: ignore
     return return_file
 
 
-class Client:
+class Client:  # pragma: no cover
     def __init__(self, platform, hostname, username, password, port, keys):
         self.platform = platform
         self.hostname = hostname
@@ -39,7 +45,7 @@ def connect(self):
         else:
             try:
                 self.net_connect = Netmiko(device_type=self.platform, host=self.hostname, port=self.port,
-                                           use_keys=True, username=self.username, password=self.password)
+                                           use_keys=False, username=self.username, password=self.password)
             except Exception as err:
                 return_error(err)
 
@@ -73,7 +79,7 @@ def cmds(self, require_exit, exit_argument, commands, enable, isConfig):
         return output
 
 
-def test_command(client):
+def test_command(client):  # pragma: no cover
     client.connect()
     client.disconnect()
     demisto.results('ok')
@@ -84,7 +90,7 @@ def cmds_command(client, args):
 
     # Parse the commands
     cmds = args.get('cmds')
-    if type(cmds) != list:
+    if type(cmds) != list:  # pragma: no cover
         try:
             cmds = cmds.split('\n')
         except Exception as err:
@@ -126,25 +132,38 @@ def cmds_command(client, args):
             demisto.error(f"Error with raw print output - {err}")
 
     else:
-        md = tableToMarkdown(f'Command(s) against {client.hostname} ({client.platform}):', output.get('Commands', []))
+        hdrs = ["Hostname", "DateTimeUTC", "Command", "Output"]
+        data = []
+
+        # Single command
+        if len(cmds) == 1:
+            data.append(output["Commands"][0])
+
+        # Multiple commands
+        else:
+            for item in output["Commands"]:
+                data.append(include_keys(item, hdrs))
+
+        md = tableToMarkdown(f'Command(s) against {client.hostname} ({client.platform}):', data, headers=hdrs)
     outputs_key_field = None
     outputs_prefix = None
     outputs = None
     if not disable_context:
         outputs_prefix = "Netmiko"
         outputs_key_field = 'DateTimeUTC'
-        outputs = output.get('Commands')
+        outputs = output
 
     command_results = CommandResults(
         outputs_prefix=outputs_prefix,
         outputs_key_field=outputs_key_field,
         outputs=outputs,
         readable_output=md
     )
-    return_results(command_results)
+
+    return command_results
 
 
-def main():
+def main():  # pragma: no cover
 
     params = demisto.params()
     args = demisto.args()
@@ -175,8 +194,9 @@ def main():
     if command == 'test-module':
         test_command(client)
     elif command == 'netmiko-cmds':
-        cmds_command(client, args)
+        results = cmds_command(client, args)
+        return_results(results)
 
 
-if __name__ in ['__main__', 'builtin', 'builtins']:
+if __name__ in ['__main__', 'builtin', 'builtins']:  # pragma: no cover
     main()