[v5,06/10] dts: add ssh connection module
Checks
Commit Message
The module uses the pexpect python library and implements connection to
a node and two ways to interact with the node:
1. Send a string with specified prompt which will be matched after
the string has been sent to the node.
2. Send a command to be executed. No prompt is specified here.
Signed-off-by: Owen Hilyard <ohilyard@iol.unh.edu>
Signed-off-by: Juraj Linkeš <juraj.linkes@pantheon.tech>
---
dts/framework/exception.py | 48 +++++
.../remote_session/session_factory.py | 16 ++
dts/framework/remote_session/ssh_session.py | 189 ++++++++++++++++++
dts/framework/utils.py | 13 ++
4 files changed, 266 insertions(+)
create mode 100644 dts/framework/remote_session/session_factory.py
create mode 100644 dts/framework/remote_session/ssh_session.py
create mode 100644 dts/framework/utils.py
Comments
On Mon, Sep 26, 2022 at 02:17:09PM +0000, Juraj Linkeš wrote:
> The module uses the pexpect python library and implements connection to
> a node and two ways to interact with the node:
> 1. Send a string with specified prompt which will be matched after
> the string has been sent to the node.
> 2. Send a command to be executed. No prompt is specified here.
>
> Signed-off-by: Owen Hilyard <ohilyard@iol.unh.edu>
> Signed-off-by: Juraj Linkeš <juraj.linkes@pantheon.tech>
> ---
> dts/framework/exception.py | 48 +++++
> .../remote_session/session_factory.py | 16 ++
> dts/framework/remote_session/ssh_session.py | 189 ++++++++++++++++++
> dts/framework/utils.py | 13 ++
> 4 files changed, 266 insertions(+)
> create mode 100644 dts/framework/remote_session/session_factory.py
> create mode 100644 dts/framework/remote_session/ssh_session.py
> create mode 100644 dts/framework/utils.py
>
> diff --git a/dts/framework/exception.py b/dts/framework/exception.py
> index 60fd98c9ca..8466990aa5 100644
> --- a/dts/framework/exception.py
> +++ b/dts/framework/exception.py
> @@ -9,6 +9,54 @@
> """
>
>
> +class TimeoutException(Exception):
> + """
> + Command execution timeout.
> + """
> +
> + command: str
> + output: str
> +
> + def __init__(self, command: str, output: str):
> + self.command = command
> + self.output = output
> +
> + def __str__(self) -> str:
> + return f"TIMEOUT on {self.command}"
> +
> + def get_output(self) -> str:
> + return self.output
> +
> +
> +class SSHConnectionException(Exception):
> + """
> + SSH connection error.
> + """
> +
> + host: str
> +
> + def __init__(self, host: str):
> + self.host = host
> +
> + def __str__(self) -> str:
> + return f"Error trying to connect with {self.host}"
> +
> +
> +class SSHSessionDeadException(Exception):
> + """
> + SSH session is not alive.
> + It can no longer be used.
> + """
> +
> + host: str
> +
> + def __init__(self, host: str):
> + self.host = host
> +
> + def __str__(self) -> str:
> + return f"SSH session with {self.host} has died"
> +
> +
> class ConfigParseException(Exception):
> """
> Configuration file parse failure exception.
> diff --git a/dts/framework/remote_session/session_factory.py b/dts/framework/remote_session/session_factory.py
> new file mode 100644
> index 0000000000..ff05df97bf
> --- /dev/null
> +++ b/dts/framework/remote_session/session_factory.py
> @@ -0,0 +1,16 @@
> +# SPDX-License-Identifier: BSD-3-Clause
> +# Copyright(c) 2022 PANTHEON.tech s.r.o.
> +# Copyright(c) 2022 University of New Hampshire
> +#
> +
> +from framework.config import NodeConfiguration
> +from framework.logger import DTSLOG
> +
> +from .remote_session import RemoteSession
> +from .ssh_session import SSHSession
> +
> +
> +def create_remote_session(
> + node_config: NodeConfiguration, name: str, logger: DTSLOG
> +) -> RemoteSession:
> + return SSHSession(node_config, name, logger)
> diff --git a/dts/framework/remote_session/ssh_session.py b/dts/framework/remote_session/ssh_session.py
> new file mode 100644
> index 0000000000..e0614e0f90
> --- /dev/null
> +++ b/dts/framework/remote_session/ssh_session.py
> @@ -0,0 +1,189 @@
> +# SPDX-License-Identifier: BSD-3-Clause
> +# Copyright(c) 2010-2014 Intel Corporation
> +# Copyright(c) 2022 PANTHEON.tech s.r.o.
> +# Copyright(c) 2022 University of New Hampshire
> +#
> +
> +
> +import time
> +
> +from pexpect import pxssh
> +
> +from framework.config import NodeConfiguration
> +from framework.exception import (
> + SSHConnectionException,
> + SSHSessionDeadException,
> + TimeoutException,
> +)
> +from framework.logger import DTSLOG
> +from framework.utils import GREEN, RED
> +
> +from .remote_session import RemoteSession
> +
> +
> +class SSHSession(RemoteSession):
> + """
> + Module for creating Pexpect SSH sessions to a node.
> + """
> +
> + session: pxssh.pxssh
> + magic_prompt: str
> +
> + def __init__(
> + self,
> + node_config: NodeConfiguration,
> + session_name: str,
> + logger: DTSLOG,
> + ):
> + self.magic_prompt = "MAGIC PROMPT"
> + super(SSHSession, self).__init__(node_config, session_name, logger)
> +
> + def _connect(self) -> None:
> + """
> + Create connection to assigned node.
> + """
> + retry_attempts = 10
> + login_timeout = 20 if self.port else 10
> + password_regex = (
> + r"(?i)(?:password:)|(?:passphrase for key)|(?i)(password for .+:)"
> + )
> + try:
> + for retry_attempt in range(retry_attempts):
> + self.session = pxssh.pxssh(encoding="utf-8")
> + try:
> + self.session.login(
> + self.ip,
> + self.username,
> + self.password,
> + original_prompt="[$#>]",
> + port=self.port,
> + login_timeout=login_timeout,
> + password_regex=password_regex,
> + )
> + break
> + except Exception as e:
> + print(e)
> + time.sleep(2)
> + print(f"Retrying connection: retry number {retry_attempt + 1}.")
> + else:
> + raise Exception(f"Connection to {self.hostname} failed")
> +
> + self.logger.info(f"Connection to {self.hostname} succeeded")
> + self.send_expect("stty -echo", "#")
> + self.send_expect("stty columns 1000", "#")
> + except Exception as e:
> + print(RED(str(e)))
> + if getattr(self, "port", None):
> + suggestion = (
> + f"\nSuggestion: Check if the firewall on {self.hostname} is "
> + f"stopped.\n"
> + )
> + print(GREEN(suggestion))
> +
> + raise SSHConnectionException(self.hostname)
> +
> + def send_expect_base(self, command: str, prompt: str, timeout: float) -> str:
> + self.clean_session()
> + original_prompt = self.session.PROMPT
> + self.session.PROMPT = prompt
> + self.__sendline(command)
> + self.__prompt(command, timeout)
> +
> + before = self._get_output()
> + self.session.PROMPT = original_prompt
> + return before
> +
> + def send_expect(
> + self, command: str, prompt: str, timeout: float = 15, verify: bool = False
> + ) -> str | int:
> + try:
> + ret = self.send_expect_base(command, prompt, timeout)
> + if verify:
> + ret_status = self.send_expect_base("echo $?", prompt, timeout)
> + try:
> + retval = int(ret_status)
> + if not retval:
> + self.logger.error(f"Command: {command} failure!")
> + self.logger.error(ret)
> + return retval
> + else:
> + return ret
Just a minor nit. Isn't the verify logic reversed in this commit?
In V4 "if not retval" was an OK case (returning the output), now it
reports an error.
> + except ValueError:
> + return ret
> + else:
> + return ret
> + except Exception as e:
> + print(
> + f"Exception happened in [{command}] and output is "
> + f"[{self._get_output()}]"
> + )
> + raise e
> +
> + def _send_command(self, command: str, timeout: float = 1) -> str:
> + try:
> + self.clean_session()
> + self.__sendline(command)
> + except Exception as e:
> + raise e
> +
> + output = self.get_output(timeout=timeout)
> + self.session.PROMPT = self.session.UNIQUE_PROMPT
> + self.session.prompt(0.1)
> +
> + return output
> +
> + def clean_session(self) -> None:
> + self.get_output(timeout=0.01)
> +
> + def _get_output(self) -> str:
> + if not self.is_alive():
> + raise SSHSessionDeadException(self.hostname)
> + before = self.session.before.rsplit("\r\n", 1)[0]
> + if before == "[PEXPECT]":
> + return ""
> + return before
> +
> + def get_output(self, timeout: float = 15) -> str:
> + """
> + Get all output before timeout
> + """
> + self.session.PROMPT = self.magic_prompt
> + try:
> + self.session.prompt(timeout)
> + except Exception:
> + pass
> +
> + before = self._get_output()
> + self.__flush()
> +
> + self.logger.debug(before)
> + return before
> +
> + def __flush(self) -> None:
> + """
> + Clear all session buffer
> + """
> + self.session.buffer = ""
> + self.session.before = ""
> +
> + def __prompt(self, command: str, timeout: float) -> None:
> + if not self.session.prompt(timeout):
> + raise TimeoutException(command, self._get_output()) from None
> +
> + def __sendline(self, command: str) -> None:
> + if not self.is_alive():
> + raise SSHSessionDeadException(self.hostname)
> + if len(command) == 2 and command.startswith("^"):
> + self.session.sendcontrol(command[1])
> + else:
> + self.session.sendline(command)
> +
> + def _close(self, force: bool = False) -> None:
> + if force is True:
> + self.session.close()
> + else:
> + if self.is_alive():
> + self.session.logout()
> +
> + def is_alive(self) -> bool:
> + return self.session.isalive()
> diff --git a/dts/framework/utils.py b/dts/framework/utils.py
> new file mode 100644
> index 0000000000..26b784ebb5
> --- /dev/null
> +++ b/dts/framework/utils.py
> @@ -0,0 +1,13 @@
> +# SPDX-License-Identifier: BSD-3-Clause
> +# Copyright(c) 2010-2014 Intel Corporation
> +# Copyright(c) 2022 PANTHEON.tech s.r.o.
> +# Copyright(c) 2022 University of New Hampshire
> +#
> +
> +
> +def RED(text: str) -> str:
> + return f"\u001B[31;1m{str(text)}\u001B[0m"
> +
> +
> +def GREEN(text: str) -> str:
> + return f"\u001B[32;1m{str(text)}\u001B[0m"
> --
> 2.30.2
>
Reviewed-by: Stanislaw Kardach <kda@semihalf.com>
> -----Original Message-----
> From: Stanislaw Kardach <kda@semihalf.com>
> Sent: Tuesday, September 27, 2022 12:12 PM
> To: Juraj Linkeš <juraj.linkes@pantheon.tech>
> Cc: thomas@monjalon.net; david.marchand@redhat.com;
> Honnappa.Nagarahalli@arm.com; ohilyard@iol.unh.edu; lijuan.tu@intel.com;
> bruce.richardson@intel.com; dev@dpdk.org
> Subject: Re: [PATCH v5 06/10] dts: add ssh connection module
>
> On Mon, Sep 26, 2022 at 02:17:09PM +0000, Juraj Linkeš wrote:
> > The module uses the pexpect python library and implements connection
> > to a node and two ways to interact with the node:
> > 1. Send a string with specified prompt which will be matched after
> > the string has been sent to the node.
> > 2. Send a command to be executed. No prompt is specified here.
> >
> > Signed-off-by: Owen Hilyard <ohilyard@iol.unh.edu>
> > Signed-off-by: Juraj Linkeš <juraj.linkes@pantheon.tech>
> > ---
> > dts/framework/exception.py | 48 +++++
> > .../remote_session/session_factory.py | 16 ++
> > dts/framework/remote_session/ssh_session.py | 189 ++++++++++++++++++
> > dts/framework/utils.py | 13 ++
> > 4 files changed, 266 insertions(+)
> > create mode 100644 dts/framework/remote_session/session_factory.py
> > create mode 100644 dts/framework/remote_session/ssh_session.py
> > create mode 100644 dts/framework/utils.py
> >
> > diff --git a/dts/framework/exception.py b/dts/framework/exception.py
> > index 60fd98c9ca..8466990aa5 100644
> > --- a/dts/framework/exception.py
> > +++ b/dts/framework/exception.py
> > @@ -9,6 +9,54 @@
> > """
> >
> >
> > +class TimeoutException(Exception):
> > + """
> > + Command execution timeout.
> > + """
> > +
> > + command: str
> > + output: str
> > +
> > + def __init__(self, command: str, output: str):
> > + self.command = command
> > + self.output = output
> > +
> > + def __str__(self) -> str:
> > + return f"TIMEOUT on {self.command}"
> > +
> > + def get_output(self) -> str:
> > + return self.output
> > +
> > +
> > +class SSHConnectionException(Exception):
> > + """
> > + SSH connection error.
> > + """
> > +
> > + host: str
> > +
> > + def __init__(self, host: str):
> > + self.host = host
> > +
> > + def __str__(self) -> str:
> > + return f"Error trying to connect with {self.host}"
> > +
> > +
> > +class SSHSessionDeadException(Exception):
> > + """
> > + SSH session is not alive.
> > + It can no longer be used.
> > + """
> > +
> > + host: str
> > +
> > + def __init__(self, host: str):
> > + self.host = host
> > +
> > + def __str__(self) -> str:
> > + return f"SSH session with {self.host} has died"
> > +
> > +
> > class ConfigParseException(Exception):
> > """
> > Configuration file parse failure exception.
> > diff --git a/dts/framework/remote_session/session_factory.py
> > b/dts/framework/remote_session/session_factory.py
> > new file mode 100644
> > index 0000000000..ff05df97bf
> > --- /dev/null
> > +++ b/dts/framework/remote_session/session_factory.py
> > @@ -0,0 +1,16 @@
> > +# SPDX-License-Identifier: BSD-3-Clause # Copyright(c) 2022
> > +PANTHEON.tech s.r.o.
> > +# Copyright(c) 2022 University of New Hampshire #
> > +
> > +from framework.config import NodeConfiguration from framework.logger
> > +import DTSLOG
> > +
> > +from .remote_session import RemoteSession from .ssh_session import
> > +SSHSession
> > +
> > +
> > +def create_remote_session(
> > + node_config: NodeConfiguration, name: str, logger: DTSLOG
> > +) -> RemoteSession:
> > + return SSHSession(node_config, name, logger)
> > diff --git a/dts/framework/remote_session/ssh_session.py
> > b/dts/framework/remote_session/ssh_session.py
> > new file mode 100644
> > index 0000000000..e0614e0f90
> > --- /dev/null
> > +++ b/dts/framework/remote_session/ssh_session.py
> > @@ -0,0 +1,189 @@
> > +# SPDX-License-Identifier: BSD-3-Clause # Copyright(c) 2010-2014
> > +Intel Corporation # Copyright(c) 2022 PANTHEON.tech s.r.o.
> > +# Copyright(c) 2022 University of New Hampshire #
> > +
> > +
> > +import time
> > +
> > +from pexpect import pxssh
> > +
> > +from framework.config import NodeConfiguration from
> > +framework.exception import (
> > + SSHConnectionException,
> > + SSHSessionDeadException,
> > + TimeoutException,
> > +)
> > +from framework.logger import DTSLOG
> > +from framework.utils import GREEN, RED
> > +
> > +from .remote_session import RemoteSession
> > +
> > +
> > +class SSHSession(RemoteSession):
> > + """
> > + Module for creating Pexpect SSH sessions to a node.
> > + """
> > +
> > + session: pxssh.pxssh
> > + magic_prompt: str
> > +
> > + def __init__(
> > + self,
> > + node_config: NodeConfiguration,
> > + session_name: str,
> > + logger: DTSLOG,
> > + ):
> > + self.magic_prompt = "MAGIC PROMPT"
> > + super(SSHSession, self).__init__(node_config, session_name,
> > + logger)
> > +
> > + def _connect(self) -> None:
> > + """
> > + Create connection to assigned node.
> > + """
> > + retry_attempts = 10
> > + login_timeout = 20 if self.port else 10
> > + password_regex = (
> > + r"(?i)(?:password:)|(?:passphrase for key)|(?i)(password for .+:)"
> > + )
> > + try:
> > + for retry_attempt in range(retry_attempts):
> > + self.session = pxssh.pxssh(encoding="utf-8")
> > + try:
> > + self.session.login(
> > + self.ip,
> > + self.username,
> > + self.password,
> > + original_prompt="[$#>]",
> > + port=self.port,
> > + login_timeout=login_timeout,
> > + password_regex=password_regex,
> > + )
> > + break
> > + except Exception as e:
> > + print(e)
> > + time.sleep(2)
> > + print(f"Retrying connection: retry number {retry_attempt + 1}.")
> > + else:
> > + raise Exception(f"Connection to {self.hostname}
> > + failed")
> > +
> > + self.logger.info(f"Connection to {self.hostname} succeeded")
> > + self.send_expect("stty -echo", "#")
> > + self.send_expect("stty columns 1000", "#")
> > + except Exception as e:
> > + print(RED(str(e)))
> > + if getattr(self, "port", None):
> > + suggestion = (
> > + f"\nSuggestion: Check if the firewall on {self.hostname} is "
> > + f"stopped.\n"
> > + )
> > + print(GREEN(suggestion))
> > +
> > + raise SSHConnectionException(self.hostname)
> > +
> > + def send_expect_base(self, command: str, prompt: str, timeout: float) ->
> str:
> > + self.clean_session()
> > + original_prompt = self.session.PROMPT
> > + self.session.PROMPT = prompt
> > + self.__sendline(command)
> > + self.__prompt(command, timeout)
> > +
> > + before = self._get_output()
> > + self.session.PROMPT = original_prompt
> > + return before
> > +
> > + def send_expect(
> > + self, command: str, prompt: str, timeout: float = 15, verify: bool = False
> > + ) -> str | int:
> > + try:
> > + ret = self.send_expect_base(command, prompt, timeout)
> > + if verify:
> > + ret_status = self.send_expect_base("echo $?", prompt, timeout)
> > + try:
> > + retval = int(ret_status)
> > + if not retval:
> > + self.logger.error(f"Command: {command} failure!")
> > + self.logger.error(ret)
> > + return retval
> > + else:
> > + return ret
> Just a minor nit. Isn't the verify logic reversed in this commit?
> In V4 "if not retval" was an OK case (returning the output), now it reports an
> error.
Yes, this is a mistake on my part:
0 is False and any other value is True, which is (basically) the opposite of shell. That means not retval is True when retval = 0 and that shouldn't produce an error. Thanks for the catch.
> > + except ValueError:
> > + return ret
> > + else:
> > + return ret
> > + except Exception as e:
> > + print(
> > + f"Exception happened in [{command}] and output is "
> > + f"[{self._get_output()}]"
> > + )
> > + raise e
> > +
> > + def _send_command(self, command: str, timeout: float = 1) -> str:
> > + try:
> > + self.clean_session()
> > + self.__sendline(command)
> > + except Exception as e:
> > + raise e
> > +
> > + output = self.get_output(timeout=timeout)
> > + self.session.PROMPT = self.session.UNIQUE_PROMPT
> > + self.session.prompt(0.1)
> > +
> > + return output
> > +
> > + def clean_session(self) -> None:
> > + self.get_output(timeout=0.01)
> > +
> > + def _get_output(self) -> str:
> > + if not self.is_alive():
> > + raise SSHSessionDeadException(self.hostname)
> > + before = self.session.before.rsplit("\r\n", 1)[0]
> > + if before == "[PEXPECT]":
> > + return ""
> > + return before
> > +
> > + def get_output(self, timeout: float = 15) -> str:
> > + """
> > + Get all output before timeout
> > + """
> > + self.session.PROMPT = self.magic_prompt
> > + try:
> > + self.session.prompt(timeout)
> > + except Exception:
> > + pass
> > +
> > + before = self._get_output()
> > + self.__flush()
> > +
> > + self.logger.debug(before)
> > + return before
> > +
> > + def __flush(self) -> None:
> > + """
> > + Clear all session buffer
> > + """
> > + self.session.buffer = ""
> > + self.session.before = ""
> > +
> > + def __prompt(self, command: str, timeout: float) -> None:
> > + if not self.session.prompt(timeout):
> > + raise TimeoutException(command, self._get_output()) from
> > + None
> > +
> > + def __sendline(self, command: str) -> None:
> > + if not self.is_alive():
> > + raise SSHSessionDeadException(self.hostname)
> > + if len(command) == 2 and command.startswith("^"):
> > + self.session.sendcontrol(command[1])
> > + else:
> > + self.session.sendline(command)
> > +
> > + def _close(self, force: bool = False) -> None:
> > + if force is True:
> > + self.session.close()
> > + else:
> > + if self.is_alive():
> > + self.session.logout()
> > +
> > + def is_alive(self) -> bool:
> > + return self.session.isalive()
> > diff --git a/dts/framework/utils.py b/dts/framework/utils.py new file
> > mode 100644 index 0000000000..26b784ebb5
> > --- /dev/null
> > +++ b/dts/framework/utils.py
> > @@ -0,0 +1,13 @@
> > +# SPDX-License-Identifier: BSD-3-Clause # Copyright(c) 2010-2014
> > +Intel Corporation # Copyright(c) 2022 PANTHEON.tech s.r.o.
> > +# Copyright(c) 2022 University of New Hampshire #
> > +
> > +
> > +def RED(text: str) -> str:
> > + return f"\u001B[31;1m{str(text)}\u001B[0m"
> > +
> > +
> > +def GREEN(text: str) -> str:
> > + return f"\u001B[32;1m{str(text)}\u001B[0m"
> > --
> > 2.30.2
> >
>
> Reviewed-by: Stanislaw Kardach <kda@semihalf.com>
>
> --
> Best Regards,
> Stanislaw Kardach
@@ -9,6 +9,54 @@
"""
+class TimeoutException(Exception):
+ """
+ Command execution timeout.
+ """
+
+ command: str
+ output: str
+
+ def __init__(self, command: str, output: str):
+ self.command = command
+ self.output = output
+
+ def __str__(self) -> str:
+ return f"TIMEOUT on {self.command}"
+
+ def get_output(self) -> str:
+ return self.output
+
+
+class SSHConnectionException(Exception):
+ """
+ SSH connection error.
+ """
+
+ host: str
+
+ def __init__(self, host: str):
+ self.host = host
+
+ def __str__(self) -> str:
+ return f"Error trying to connect with {self.host}"
+
+
+class SSHSessionDeadException(Exception):
+ """
+ SSH session is not alive.
+ It can no longer be used.
+ """
+
+ host: str
+
+ def __init__(self, host: str):
+ self.host = host
+
+ def __str__(self) -> str:
+ return f"SSH session with {self.host} has died"
+
+
class ConfigParseException(Exception):
"""
Configuration file parse failure exception.
new file mode 100644
@@ -0,0 +1,16 @@
+# SPDX-License-Identifier: BSD-3-Clause
+# Copyright(c) 2022 PANTHEON.tech s.r.o.
+# Copyright(c) 2022 University of New Hampshire
+#
+
+from framework.config import NodeConfiguration
+from framework.logger import DTSLOG
+
+from .remote_session import RemoteSession
+from .ssh_session import SSHSession
+
+
+def create_remote_session(
+ node_config: NodeConfiguration, name: str, logger: DTSLOG
+) -> RemoteSession:
+ return SSHSession(node_config, name, logger)
new file mode 100644
@@ -0,0 +1,189 @@
+# SPDX-License-Identifier: BSD-3-Clause
+# Copyright(c) 2010-2014 Intel Corporation
+# Copyright(c) 2022 PANTHEON.tech s.r.o.
+# Copyright(c) 2022 University of New Hampshire
+#
+
+
+import time
+
+from pexpect import pxssh
+
+from framework.config import NodeConfiguration
+from framework.exception import (
+ SSHConnectionException,
+ SSHSessionDeadException,
+ TimeoutException,
+)
+from framework.logger import DTSLOG
+from framework.utils import GREEN, RED
+
+from .remote_session import RemoteSession
+
+
+class SSHSession(RemoteSession):
+ """
+ Module for creating Pexpect SSH sessions to a node.
+ """
+
+ session: pxssh.pxssh
+ magic_prompt: str
+
+ def __init__(
+ self,
+ node_config: NodeConfiguration,
+ session_name: str,
+ logger: DTSLOG,
+ ):
+ self.magic_prompt = "MAGIC PROMPT"
+ super(SSHSession, self).__init__(node_config, session_name, logger)
+
+ def _connect(self) -> None:
+ """
+ Create connection to assigned node.
+ """
+ retry_attempts = 10
+ login_timeout = 20 if self.port else 10
+ password_regex = (
+ r"(?i)(?:password:)|(?:passphrase for key)|(?i)(password for .+:)"
+ )
+ try:
+ for retry_attempt in range(retry_attempts):
+ self.session = pxssh.pxssh(encoding="utf-8")
+ try:
+ self.session.login(
+ self.ip,
+ self.username,
+ self.password,
+ original_prompt="[$#>]",
+ port=self.port,
+ login_timeout=login_timeout,
+ password_regex=password_regex,
+ )
+ break
+ except Exception as e:
+ print(e)
+ time.sleep(2)
+ print(f"Retrying connection: retry number {retry_attempt + 1}.")
+ else:
+ raise Exception(f"Connection to {self.hostname} failed")
+
+ self.logger.info(f"Connection to {self.hostname} succeeded")
+ self.send_expect("stty -echo", "#")
+ self.send_expect("stty columns 1000", "#")
+ except Exception as e:
+ print(RED(str(e)))
+ if getattr(self, "port", None):
+ suggestion = (
+ f"\nSuggestion: Check if the firewall on {self.hostname} is "
+ f"stopped.\n"
+ )
+ print(GREEN(suggestion))
+
+ raise SSHConnectionException(self.hostname)
+
+ def send_expect_base(self, command: str, prompt: str, timeout: float) -> str:
+ self.clean_session()
+ original_prompt = self.session.PROMPT
+ self.session.PROMPT = prompt
+ self.__sendline(command)
+ self.__prompt(command, timeout)
+
+ before = self._get_output()
+ self.session.PROMPT = original_prompt
+ return before
+
+ def send_expect(
+ self, command: str, prompt: str, timeout: float = 15, verify: bool = False
+ ) -> str | int:
+ try:
+ ret = self.send_expect_base(command, prompt, timeout)
+ if verify:
+ ret_status = self.send_expect_base("echo $?", prompt, timeout)
+ try:
+ retval = int(ret_status)
+ if not retval:
+ self.logger.error(f"Command: {command} failure!")
+ self.logger.error(ret)
+ return retval
+ else:
+ return ret
+ except ValueError:
+ return ret
+ else:
+ return ret
+ except Exception as e:
+ print(
+ f"Exception happened in [{command}] and output is "
+ f"[{self._get_output()}]"
+ )
+ raise e
+
+ def _send_command(self, command: str, timeout: float = 1) -> str:
+ try:
+ self.clean_session()
+ self.__sendline(command)
+ except Exception as e:
+ raise e
+
+ output = self.get_output(timeout=timeout)
+ self.session.PROMPT = self.session.UNIQUE_PROMPT
+ self.session.prompt(0.1)
+
+ return output
+
+ def clean_session(self) -> None:
+ self.get_output(timeout=0.01)
+
+ def _get_output(self) -> str:
+ if not self.is_alive():
+ raise SSHSessionDeadException(self.hostname)
+ before = self.session.before.rsplit("\r\n", 1)[0]
+ if before == "[PEXPECT]":
+ return ""
+ return before
+
+ def get_output(self, timeout: float = 15) -> str:
+ """
+ Get all output before timeout
+ """
+ self.session.PROMPT = self.magic_prompt
+ try:
+ self.session.prompt(timeout)
+ except Exception:
+ pass
+
+ before = self._get_output()
+ self.__flush()
+
+ self.logger.debug(before)
+ return before
+
+ def __flush(self) -> None:
+ """
+ Clear all session buffer
+ """
+ self.session.buffer = ""
+ self.session.before = ""
+
+ def __prompt(self, command: str, timeout: float) -> None:
+ if not self.session.prompt(timeout):
+ raise TimeoutException(command, self._get_output()) from None
+
+ def __sendline(self, command: str) -> None:
+ if not self.is_alive():
+ raise SSHSessionDeadException(self.hostname)
+ if len(command) == 2 and command.startswith("^"):
+ self.session.sendcontrol(command[1])
+ else:
+ self.session.sendline(command)
+
+ def _close(self, force: bool = False) -> None:
+ if force is True:
+ self.session.close()
+ else:
+ if self.is_alive():
+ self.session.logout()
+
+ def is_alive(self) -> bool:
+ return self.session.isalive()
new file mode 100644
@@ -0,0 +1,13 @@
+# SPDX-License-Identifier: BSD-3-Clause
+# Copyright(c) 2010-2014 Intel Corporation
+# Copyright(c) 2022 PANTHEON.tech s.r.o.
+# Copyright(c) 2022 University of New Hampshire
+#
+
+
+def RED(text: str) -> str:
+ return f"\u001B[31;1m{str(text)}\u001B[0m"
+
+
+def GREEN(text: str) -> str:
+ return f"\u001B[32;1m{str(text)}\u001B[0m"