blob: 00f50510537f44087ba0e926dcadc74a19c72507 [file] [log] [blame]
Patrick Williamsc124f4f2015-09-15 14:41:29 -05001# Copyright (C) 2013 Intel Corporation
2#
3# Released under the MIT license (see COPYING.MIT)
4
5# Provides a class for setting up ssh connections,
6# running commands and copying files to/from a target.
7# It's used by testimage.bbclass and tests in lib/oeqa/runtime.
8
9import subprocess
10import time
11import os
12import select
13
14
15class SSHProcess(object):
16 def __init__(self, **options):
17
18 self.defaultopts = {
19 "stdout": subprocess.PIPE,
20 "stderr": subprocess.STDOUT,
21 "stdin": None,
22 "shell": False,
23 "bufsize": -1,
24 "preexec_fn": os.setsid,
25 }
26 self.options = dict(self.defaultopts)
27 self.options.update(options)
28 self.status = None
29 self.output = None
30 self.process = None
31 self.starttime = None
32 self.logfile = None
33
34 # Unset DISPLAY which means we won't trigger SSH_ASKPASS
35 env = os.environ.copy()
36 if "DISPLAY" in env:
37 del env['DISPLAY']
38 self.options['env'] = env
39
40 def log(self, msg):
41 if self.logfile:
42 with open(self.logfile, "a") as f:
43 f.write("%s" % msg)
44
45 def _run(self, command, timeout=None, logfile=None):
46 self.logfile = logfile
47 self.starttime = time.time()
48 output = ''
49 self.process = subprocess.Popen(command, **self.options)
50 if timeout:
51 endtime = self.starttime + timeout
52 eof = False
53 while time.time() < endtime and not eof:
54 if select.select([self.process.stdout], [], [], 5)[0] != []:
55 data = os.read(self.process.stdout.fileno(), 1024)
56 if not data:
57 self.process.stdout.close()
58 eof = True
59 else:
60 output += data
61 self.log(data)
62 endtime = time.time() + timeout
63
64
65 # process hasn't returned yet
66 if not eof:
67 self.process.terminate()
68 time.sleep(5)
69 try:
70 self.process.kill()
71 except OSError:
72 pass
73 lastline = "\nProcess killed - no output for %d seconds. Total running time: %d seconds." % (timeout, time.time() - self.starttime)
74 self.log(lastline)
75 output += lastline
76 else:
77 output = self.process.communicate()[0]
78 self.log(output.rstrip())
79
80 self.status = self.process.wait()
81 self.output = output.rstrip()
82
83 def run(self, command, timeout=None, logfile=None):
84 try:
85 self._run(command, timeout, logfile)
86 except:
87 # Need to guard against a SystemExit or other exception occuring whilst running
88 # and ensure we don't leave a process behind.
89 if self.process.poll() is None:
90 self.process.kill()
91 self.status = self.process.wait()
92 raise
93 return (self.status, self.output)
94
95class SSHControl(object):
96 def __init__(self, ip, logfile=None, timeout=300, user='root', port=None):
97 self.ip = ip
98 self.defaulttimeout = timeout
99 self.ignore_status = True
100 self.logfile = logfile
101 self.user = user
102 self.ssh_options = [
103 '-o', 'UserKnownHostsFile=/dev/null',
104 '-o', 'StrictHostKeyChecking=no',
105 '-o', 'LogLevel=ERROR'
106 ]
107 self.ssh = ['ssh', '-l', self.user ] + self.ssh_options
108 self.scp = ['scp'] + self.ssh_options
109 if port:
110 self.ssh = self.ssh + [ '-p', port ]
111 self.scp = self.scp + [ '-P', port ]
112
113 def log(self, msg):
114 if self.logfile:
115 with open(self.logfile, "a") as f:
116 f.write("%s\n" % msg)
117
118 def _internal_run(self, command, timeout=None, ignore_status = True):
119 self.log("[Running]$ %s" % " ".join(command))
120
121 proc = SSHProcess()
122 status, output = proc.run(command, timeout, logfile=self.logfile)
123
124 self.log("[Command returned '%d' after %.2f seconds]" % (status, time.time() - proc.starttime))
125
126 if status and not ignore_status:
127 raise AssertionError("Command '%s' returned non-zero exit status %d:\n%s" % (command, status, output))
128
129 return (status, output)
130
131 def run(self, command, timeout=None):
132 """
133 command - ssh command to run
134 timeout=<val> - kill command if there is no output after <val> seconds
135 timeout=None - kill command if there is no output after a default value seconds
136 timeout=0 - no timeout, let command run until it returns
137 """
138
139 # We need to source /etc/profile for a proper PATH on the target
140 command = self.ssh + [self.ip, ' . /etc/profile; ' + command]
141
142 if timeout is None:
143 return self._internal_run(command, self.defaulttimeout, self.ignore_status)
144 if timeout == 0:
145 return self._internal_run(command, None, self.ignore_status)
146 return self._internal_run(command, timeout, self.ignore_status)
147
148 def copy_to(self, localpath, remotepath):
149 command = self.scp + [localpath, '%s@%s:%s' % (self.user, self.ip, remotepath)]
150 return self._internal_run(command, ignore_status=False)
151
152 def copy_from(self, remotepath, localpath):
153 command = self.scp + ['%s@%s:%s' % (self.user, self.ip, remotepath), localpath]
154 return self._internal_run(command, ignore_status=False)