diff --git a/itf/plugins/com/sftp.py b/itf/plugins/com/sftp.py index fc191a5..9b7a631 100644 --- a/itf/plugins/com/sftp.py +++ b/itf/plugins/com/sftp.py @@ -17,22 +17,44 @@ # Reduce the logging level of paramiko, from DEBUG to INFO logging.getLogger("paramiko").setLevel(logging.INFO) - logger = logging.getLogger(__name__) class Sftp: - def __init__(self, ssh, target_ip, port=22): + def __init__( + self, + ssh: Ssh, + target_ip: str, + port: int = 22, + timeout: int = 15, + n_retries: int = 5, + retry_interval: int = 1, + pkey_path: str = None, + username: str = "root", + password: str = "", + ): + """ + Initialize SFTP connection to target with given ssh connection or parameters to create new ssh connection. + :param Ssh ssh: Existing SSH connection + :param str target_ip: The IP address of the target SSH server. + :param int port: The port number of the target SSH server. Default is 22. + :param int timeout: The timeout duration (in seconds) for the SSH connection. Default is 15 seconds. + :param int n_retries: The number of retries to attempt for the SSH connection. Default is 5 retries. + :param int retry_interval: The interval (in seconds) between retries. Default is 1 second. + :param str pkey_path: The file path to the private key for authentication. Default is None. + :param str username: The username for SSH authentication. Default is "root". + :param str password: The password for SSH authentication. Default is an empty string. + """ if not ssh: self._new_ssh = True else: self._new_ssh = False - self._ssh = ssh or Ssh(target_ip=target_ip, port=port) + self._ssh = ssh or Ssh(target_ip, port, timeout, n_retries, retry_interval, pkey_path, username, password) self._sftp = None def __enter__(self): """ - Open sftp connection to target given an ssh connection + Open sftp connection to target given an SSH connection """ if self._new_ssh: self._ssh = self._ssh.__enter__() @@ -43,7 +65,7 @@ def __exit__(self, exc_type, exc_val, exc_tb): self._sftp.close() if self._new_ssh: self._ssh.close() - logger.info("Closed ssh connection.") + logger.info("Closed SSH connection.") def walk(self, remote_path): """ diff --git a/itf/plugins/com/ssh.py b/itf/plugins/com/ssh.py index 87595fc..0c081b0 100644 --- a/itf/plugins/com/ssh.py +++ b/itf/plugins/com/ssh.py @@ -33,7 +33,7 @@ def __init__( password: str = "", ): """ - Initializes the Ssh class with connection parameters. + Initialize SSH connection to the target. :param str target_ip: The IP address of the target SSH server. :param int port: The port number of the target SSH server. Default is 22. @@ -58,6 +58,8 @@ def __enter__(self): self._ssh = paramiko.SSHClient() self._ssh.set_missing_host_key_policy(paramiko.client.AutoAddPolicy()) + logger.info(f"Connecting to {self._target_ip} ...") + for _ in range(self._retries): try: self._ssh.connect( @@ -70,17 +72,19 @@ def __enter__(self): banner_timeout=200, look_for_keys=False, ) + logger.info(f"SSH connection to {self._target_ip} established") break - except Exception: + except Exception as ex: + logger.debug(f"SSH connection to {self._target_ip} failed with error: \n{ex}") time.sleep(self._retry_interval) else: - raise Exception(f"ssh connection to {self._target_ip} failed") + raise Exception(f"SSH connection to {self._target_ip} failed") return self._ssh def __exit__(self, exc_type, exc_val, exc_tb): self._ssh.close() - logger.info("Closed ssh connection.") + logger.info("Closed SSH connection.") def command_with_etc(command):