ssh.py 9.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270
  1. #
  2. # Copyright (C) 2016 Intel Corporation
  3. #
  4. # SPDX-License-Identifier: MIT
  5. #
  6. import os
  7. import time
  8. import select
  9. import logging
  10. import subprocess
  11. import codecs
  12. from . import OETarget
  13. class OESSHTarget(OETarget):
  14. def __init__(self, logger, ip, server_ip, timeout=300, user='root',
  15. port=None, **kwargs):
  16. if not logger:
  17. logger = logging.getLogger('target')
  18. logger.setLevel(logging.INFO)
  19. filePath = os.path.join(os.getcwd(), 'remoteTarget.log')
  20. fileHandler = logging.FileHandler(filePath, 'w', 'utf-8')
  21. formatter = logging.Formatter(
  22. '%(asctime)s.%(msecs)03d %(levelname)s: %(message)s',
  23. '%H:%M:%S')
  24. fileHandler.setFormatter(formatter)
  25. logger.addHandler(fileHandler)
  26. super(OESSHTarget, self).__init__(logger)
  27. self.ip = ip
  28. self.server_ip = server_ip
  29. self.timeout = timeout
  30. self.user = user
  31. ssh_options = [
  32. '-o', 'UserKnownHostsFile=/dev/null',
  33. '-o', 'StrictHostKeyChecking=no',
  34. '-o', 'LogLevel=ERROR'
  35. ]
  36. self.ssh = ['ssh', '-l', self.user ] + ssh_options
  37. self.scp = ['scp'] + ssh_options
  38. if port:
  39. self.ssh = self.ssh + [ '-p', port ]
  40. self.scp = self.scp + [ '-P', port ]
  41. def start(self, **kwargs):
  42. pass
  43. def stop(self, **kwargs):
  44. pass
  45. def _run(self, command, timeout=None, ignore_status=True):
  46. """
  47. Runs command in target using SSHProcess.
  48. """
  49. self.logger.debug("[Running]$ %s" % " ".join(command))
  50. starttime = time.time()
  51. status, output = SSHCall(command, self.logger, timeout)
  52. self.logger.debug("[Command returned '%d' after %.2f seconds]"
  53. "" % (status, time.time() - starttime))
  54. if status and not ignore_status:
  55. raise AssertionError("Command '%s' returned non-zero exit "
  56. "status %d:\n%s" % (command, status, output))
  57. return (status, output)
  58. def run(self, command, timeout=None):
  59. """
  60. Runs command in target.
  61. command: Command to run on target.
  62. timeout: <value>: Kill command after <val> seconds.
  63. None: Kill command default value seconds.
  64. 0: No timeout, runs until return.
  65. """
  66. targetCmd = 'export PATH=/usr/sbin:/sbin:/usr/bin:/bin; %s' % command
  67. sshCmd = self.ssh + [self.ip, targetCmd]
  68. if timeout:
  69. processTimeout = timeout
  70. elif timeout==0:
  71. processTimeout = None
  72. else:
  73. processTimeout = self.timeout
  74. status, output = self._run(sshCmd, processTimeout, True)
  75. self.logger.debug('Command: %s\nOutput: %s\n' % (command, output))
  76. return (status, output)
  77. def copyTo(self, localSrc, remoteDst):
  78. """
  79. Copy file to target.
  80. If local file is symlink, recreate symlink in target.
  81. """
  82. if os.path.islink(localSrc):
  83. link = os.readlink(localSrc)
  84. dstDir, dstBase = os.path.split(remoteDst)
  85. sshCmd = 'cd %s; ln -s %s %s' % (dstDir, link, dstBase)
  86. return self.run(sshCmd)
  87. else:
  88. remotePath = '%s@%s:%s' % (self.user, self.ip, remoteDst)
  89. scpCmd = self.scp + [localSrc, remotePath]
  90. return self._run(scpCmd, ignore_status=False)
  91. def copyFrom(self, remoteSrc, localDst):
  92. """
  93. Copy file from target.
  94. """
  95. remotePath = '%s@%s:%s' % (self.user, self.ip, remoteSrc)
  96. scpCmd = self.scp + [remotePath, localDst]
  97. return self._run(scpCmd, ignore_status=False)
  98. def copyDirTo(self, localSrc, remoteDst):
  99. """
  100. Copy recursively localSrc directory to remoteDst in target.
  101. """
  102. for root, dirs, files in os.walk(localSrc):
  103. # Create directories in the target as needed
  104. for d in dirs:
  105. tmpDir = os.path.join(root, d).replace(localSrc, "")
  106. newDir = os.path.join(remoteDst, tmpDir.lstrip("/"))
  107. cmd = "mkdir -p %s" % newDir
  108. self.run(cmd)
  109. # Copy files into the target
  110. for f in files:
  111. tmpFile = os.path.join(root, f).replace(localSrc, "")
  112. dstFile = os.path.join(remoteDst, tmpFile.lstrip("/"))
  113. srcFile = os.path.join(root, f)
  114. self.copyTo(srcFile, dstFile)
  115. def deleteFiles(self, remotePath, files):
  116. """
  117. Deletes files in target's remotePath.
  118. """
  119. cmd = "rm"
  120. if not isinstance(files, list):
  121. files = [files]
  122. for f in files:
  123. cmd = "%s %s" % (cmd, os.path.join(remotePath, f))
  124. self.run(cmd)
  125. def deleteDir(self, remotePath):
  126. """
  127. Deletes target's remotePath directory.
  128. """
  129. cmd = "rmdir %s" % remotePath
  130. self.run(cmd)
  131. def deleteDirStructure(self, localPath, remotePath):
  132. """
  133. Delete recursively localPath structure directory in target's remotePath.
  134. This function is very usefult to delete a package that is installed in
  135. the DUT and the host running the test has such package extracted in tmp
  136. directory.
  137. Example:
  138. pwd: /home/user/tmp
  139. tree: .
  140. └── work
  141. ├── dir1
  142. │   └── file1
  143. └── dir2
  144. localpath = "/home/user/tmp" and remotepath = "/home/user"
  145. With the above variables this function will try to delete the
  146. directory in the DUT in this order:
  147. /home/user/work/dir1/file1
  148. /home/user/work/dir1 (if dir is empty)
  149. /home/user/work/dir2 (if dir is empty)
  150. /home/user/work (if dir is empty)
  151. """
  152. for root, dirs, files in os.walk(localPath, topdown=False):
  153. # Delete files first
  154. tmpDir = os.path.join(root).replace(localPath, "")
  155. remoteDir = os.path.join(remotePath, tmpDir.lstrip("/"))
  156. self.deleteFiles(remoteDir, files)
  157. # Remove dirs if empty
  158. for d in dirs:
  159. tmpDir = os.path.join(root, d).replace(localPath, "")
  160. remoteDir = os.path.join(remotePath, tmpDir.lstrip("/"))
  161. self.deleteDir(remoteDir)
  162. def SSHCall(command, logger, timeout=None, **opts):
  163. def run():
  164. nonlocal output
  165. nonlocal process
  166. starttime = time.time()
  167. process = subprocess.Popen(command, **options)
  168. if timeout:
  169. endtime = starttime + timeout
  170. eof = False
  171. while time.time() < endtime and not eof:
  172. logger.debug('time: %s, endtime: %s' % (time.time(), endtime))
  173. try:
  174. if select.select([process.stdout], [], [], 5)[0] != []:
  175. reader = codecs.getreader('utf-8')(process.stdout, 'surrogatepass')
  176. data = reader.read(1024, 4096)
  177. if not data:
  178. process.stdout.close()
  179. eof = True
  180. else:
  181. output += data
  182. logger.debug('Partial data from SSH call: %s' % data)
  183. endtime = time.time() + timeout
  184. except InterruptedError:
  185. continue
  186. # process hasn't returned yet
  187. if not eof:
  188. process.terminate()
  189. time.sleep(5)
  190. try:
  191. process.kill()
  192. except OSError:
  193. pass
  194. endtime = time.time() - starttime
  195. lastline = ("\nProcess killed - no output for %d seconds. Total"
  196. " running time: %d seconds." % (timeout, endtime))
  197. logger.debug('Received data from SSH call %s ' % lastline)
  198. output += lastline
  199. else:
  200. output = process.communicate()[0].decode("utf-8", errors='surrogatepass')
  201. logger.debug('Data from SSH call: %s' % output.rstrip())
  202. options = {
  203. "stdout": subprocess.PIPE,
  204. "stderr": subprocess.STDOUT,
  205. "stdin": None,
  206. "shell": False,
  207. "bufsize": -1,
  208. "preexec_fn": os.setsid,
  209. }
  210. options.update(opts)
  211. output = ''
  212. process = None
  213. # Unset DISPLAY which means we won't trigger SSH_ASKPASS
  214. env = os.environ.copy()
  215. if "DISPLAY" in env:
  216. del env['DISPLAY']
  217. options['env'] = env
  218. try:
  219. run()
  220. except:
  221. # Need to guard against a SystemExit or other exception ocurring
  222. # whilst running and ensure we don't leave a process behind.
  223. if process.poll() is None:
  224. process.kill()
  225. logger.debug('Something went wrong, killing SSH process')
  226. raise
  227. return (process.wait(), output.rstrip())