@@ -8,6 +8,7 @@
from typing import Any, Optional
from .logger import DTSLOG
+from .settings import SETTINGS
from .ssh_pexpect import SSHPexpect
@@ -68,3 +69,21 @@ def close(self, force: bool = False) -> None:
self.logger.logger_exit()
self.session.close(force)
+
+ def check_available(self) -> bool:
+ MAGIC_STR = "DTS_CHECK_SESSION"
+ out = self.session.send_command("echo %s" % MAGIC_STR, timeout=0.1)
+ # if not available, try to send ^C and check again
+ if MAGIC_STR not in out:
+ self.logger.info("Try to recover session...")
+ self.session.send_command("^C", timeout=SETTINGS.timeout)
+ out = self.session.send_command("echo %s" % MAGIC_STR, timeout=0.1)
+ if MAGIC_STR not in out:
+ return False
+
+ return True
+
+ def copy_file_to(
+ self, src: str, dst: str = "~/", password: str = "", node_session: Any = None
+ ) -> None:
+ self.session.copy_file_to(src, dst, password, node_session)
@@ -5,11 +5,16 @@
#
import time
-from typing import Optional
+from typing import Any, Optional
+import pexpect
from pexpect import pxssh
-from .exception import SSHConnectionException, SSHSessionDeadException, TimeoutException
+from .exception import (
+ SSHConnectionException,
+ SSHSessionDeadException,
+ TimeoutException,
+)
from .logger import DTSLOG
from .utils import GREEN, RED
@@ -203,3 +208,55 @@ def close(self, force: bool = False) -> None:
def isalive(self) -> bool:
return self.session.isalive()
+
+ def copy_file_to(
+ self, src: str, dst: str = "~/", password: str = "", node_session: Any = None
+ ) -> None:
+ """
+ Sends a local file to a remote place.
+ """
+ command: str
+ if ":" in self.node:
+ command = "scp -v -P {0} -o NoHostAuthenticationForLocalhost=yes {1} {2}@{3}:{4}".format(
+ str(self.port), src, self.username, self.ip, dst
+ )
+ else:
+ command = "scp -v {0} {1}@{2}:{3}".format(
+ src, self.username, self.node, dst
+ )
+ if password == "":
+ self._spawn_scp(command, self.password, node_session)
+ else:
+ self._spawn_scp(command, password, node_session)
+
+ def _spawn_scp(self, scp_cmd: str, password: str, node_session: Any) -> None:
+ """
+ Transfer a file with SCP
+ """
+ self.logger.info(scp_cmd)
+ # if node_session is not None, copy file from/to node env
+ # if node_session is None, copy file from/to current dts env
+ p: pexpect.spawn
+ if node_session is not None:
+ node_session.session.clean_session()
+ node_session.session.__sendline(scp_cmd)
+ p = node_session.session.session
+ else:
+ p = pexpect.spawn(scp_cmd)
+ time.sleep(0.5)
+ ssh_newkey: str = "Are you sure you want to continue connecting"
+ i: int = p.expect(
+ [ssh_newkey, "[pP]assword", "# ", pexpect.EOF, pexpect.TIMEOUT], 120
+ )
+ if i == 0: # add once in trust list
+ p.sendline("yes")
+ i = p.expect([ssh_newkey, "[pP]assword", pexpect.EOF], 2)
+
+ if i == 1:
+ time.sleep(0.5)
+ p.sendline(password)
+ p.expect("Exit status 0", 60)
+ if i == 4:
+ self.logger.error("SCP TIMEOUT error %d" % i)
+ if node_session is None:
+ p.close()