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