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