A wrapper class to run commands on remote host over SSH

ssh_utils.py

import logging
import paramiko
import sys

logging.basicConfig(
                format='%(name)s,%(levelname)s,%(asctime)s, %(message)s',
                datefmt='%m/%d/%Y %I:%M:%S %p',
                encoding='utf-8')
class Node:
    """
    A wrapper class to run commands on a remote node over SSH
    using paramiko library.
    """
    def __init__(self, host:str, user:str, password:str="", key_file:str=None):
        # Cache the input parameters for later use
        self.host = host # IP or server address
        self.user = user
        self.password = password
        self.key_file = key_file # Path to the key file for authentication

        # Setup logging
        self.logger = logging.getLogger(Node.__name__)
        self.logger.setLevel(logging.DEBUG)
        
        self.logger.info("Connecting to {} as {}".format(host, user))
        self.ssh_conn = paramiko.SSHClient()
        try:
            self.ssh_conn.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            # If both password and key_file are given, key_file takes precedence
            self.ssh_conn.connect(host,
                                username=self.user,
                                password=self.password,
                                key_filename=self.key_file)
            self.logger.info("Connected to {} as {}".format(host, user))
        except(paramiko.AuthenticationException) as ae:
            self.logger.error("Failed to connect to {} due to {}".format(host, ae), exc_info=True)
            sys.exit(1)
        except(TimeoutError) as te:
            self.logger.error("Timed out while connecting to {}. Check the host health".format(host))
            sys.exit(1)
    
    def run(self, command:str, env_var={}, fh=sys.stdout):
        """
        Runs the given command over the ssh connection and prints
        the output to the given file handler (defaults to stdout)
        Commands taking inputs are not supported yet.
        """
        self.logger.debug(">>>>> {}".format(command))
        print()
        try:
            _, stdout, stderr = self.ssh_conn.exec_command(command)
            for line in stdout:
                print(line.strip(), file=fh)
            for line in stderr:
                print(line.strip(), file=fh)
        except(paramiko.SSHException) as ssh_exc:
            self.logger.error("Failed to run \"{}\" due to {}".format(command, ssh_exc))
            sys.exit(1)
    
    def run_multiple(self, commands:list, env_var={}, fh=sys.stdout):
        """
        Execute a list of commands on the remote node over ssh
        and dumps the output to the given file handle (default: stdout)
        """
        for command in commands:
            self.run(command, env_var, fh)

    def __del__(self):
        # Close the SSH connection upon exit
        self.ssh_conn.close()

# Sample usage:
# host = "10.11.12.13"
# user = "some-user"
# key_file = "/path-to-keyfile/keyfile.pem"

# commands = ["lscpu", "lsmem", "lsblk"]
# node = Node(host, user=user, key_file=key_file)
# node.run_multiple(commands)

# anotherNode = Node(host, user="mynode.example.com", password="some-password")
# node.run("uname -a")