1#
2# Copyright (C) 2015 The Android Open Source Project
3#
4# Licensed under the Apache License, Version 2.0 (the "License");
5# you may not use this file except in compliance with the License.
6# You may obtain a copy of the License at
7#
8#      http://www.apache.org/licenses/LICENSE-2.0
9#
10# Unless required by applicable law or agreed to in writing, software
11# distributed under the License is distributed on an "AS IS" BASIS,
12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13# See the License for the specific language governing permissions and
14# limitations under the License.
15#
16import atexit
17import base64
18import logging
19import os
20import re
21import subprocess
22
23
24class FindDeviceError(RuntimeError):
25    pass
26
27
28class DeviceNotFoundError(FindDeviceError):
29    def __init__(self, serial):
30        self.serial = serial
31        super(DeviceNotFoundError, self).__init__(
32            'No device with serial {}'.format(serial))
33
34
35class NoUniqueDeviceError(FindDeviceError):
36    def __init__(self):
37        super(NoUniqueDeviceError, self).__init__('No unique device')
38
39
40class ShellError(RuntimeError):
41    def __init__(self, cmd, stdout, stderr, exit_code):
42        super(ShellError, self).__init__(
43                '`{0}` exited with code {1}'.format(cmd, exit_code))
44        self.cmd = cmd
45        self.stdout = stdout
46        self.stderr = stderr
47        self.exit_code = exit_code
48
49
50def get_devices(adb_path='adb'):
51    with open(os.devnull, 'wb') as devnull:
52        subprocess.check_call([adb_path, 'start-server'], stdout=devnull,
53                              stderr=devnull)
54    out = subprocess.check_output([adb_path, 'devices']).splitlines()
55
56    # The first line of `adb devices` just says "List of attached devices", so
57    # skip that.
58    devices = []
59    for line in out[1:]:
60        if not line.strip():
61            continue
62        if 'offline' in line:
63            continue
64
65        serial, _ = re.split(r'\s+', line, maxsplit=1)
66        devices.append(serial)
67    return devices
68
69
70def _get_unique_device(product=None, adb_path='adb'):
71    devices = get_devices(adb_path=adb_path)
72    if len(devices) != 1:
73        raise NoUniqueDeviceError()
74    return AndroidDevice(devices[0], product, adb_path)
75
76
77def _get_device_by_serial(serial, product=None, adb_path='adb'):
78    for device in get_devices(adb_path=adb_path):
79        if device == serial:
80            return AndroidDevice(serial, product, adb_path)
81    raise DeviceNotFoundError(serial)
82
83
84def get_device(serial=None, product=None, adb_path='adb'):
85    """Get a uniquely identified AndroidDevice if one is available.
86
87    Raises:
88        DeviceNotFoundError:
89            The serial specified by `serial` or $ANDROID_SERIAL is not
90            connected.
91
92        NoUniqueDeviceError:
93            Neither `serial` nor $ANDROID_SERIAL was set, and the number of
94            devices connected to the system is not 1. Having 0 connected
95            devices will also result in this error.
96
97    Returns:
98        An AndroidDevice associated with the first non-None identifier in the
99        following order of preference:
100
101        1) The `serial` argument.
102        2) The environment variable $ANDROID_SERIAL.
103        3) The single device connnected to the system.
104    """
105    if serial is not None:
106        return _get_device_by_serial(serial, product, adb_path)
107
108    android_serial = os.getenv('ANDROID_SERIAL')
109    if android_serial is not None:
110        return _get_device_by_serial(android_serial, product, adb_path)
111
112    return _get_unique_device(product, adb_path=adb_path)
113
114
115def _get_device_by_type(flag, adb_path):
116    with open(os.devnull, 'wb') as devnull:
117        subprocess.check_call([adb_path, 'start-server'], stdout=devnull,
118                              stderr=devnull)
119    try:
120        serial = subprocess.check_output([adb_path, flag, 'get-serialno']).strip()
121    except subprocess.CalledProcessError:
122        raise RuntimeError('adb unexpectedly returned nonzero')
123    if serial == 'unknown':
124        raise NoUniqueDeviceError()
125    return _get_device_by_serial(serial, adb_path=adb_path)
126
127
128def get_usb_device(adb_path='adb'):
129    """Get the unique USB-connected AndroidDevice if it is available.
130
131    Raises:
132        NoUniqueDeviceError:
133            0 or multiple devices are connected via USB.
134
135    Returns:
136        An AndroidDevice associated with the unique USB-connected device.
137    """
138    return _get_device_by_type('-d', adb_path=adb_path)
139
140
141def get_emulator_device(adb_path='adb'):
142    """Get the unique emulator AndroidDevice if it is available.
143
144    Raises:
145        NoUniqueDeviceError:
146            0 or multiple emulators are running.
147
148    Returns:
149        An AndroidDevice associated with the unique running emulator.
150    """
151    return _get_device_by_type('-e', adb_path=adb_path)
152
153
154# If necessary, modifies subprocess.check_output() or subprocess.Popen() args to run the subprocess
155# via Windows PowerShell to work-around an issue in Python 2's subprocess class on Windows where it
156# doesn't support Unicode.
157def _get_subprocess_args(args):
158    # Only do this slow work-around if Unicode is in the cmd line on Windows. PowerShell takes
159    # 600-700ms to startup on a 2013-2014 machine, which is very slow.
160    if (os.name != 'nt' or all(not isinstance(arg, unicode) for arg in args[0])):
161        return args
162
163    def escape_arg(arg):
164        # Escape for the parsing that the C Runtime does in Windows apps. In particular, this will
165        # take care of double-quotes.
166        arg = subprocess.list2cmdline([arg])
167        # Escape single-quote with another single-quote because we're about to...
168        arg = arg.replace(u"'", u"''")
169        # ...put the arg in a single-quoted string for PowerShell to parse.
170        arg = u"'" + arg + u"'"
171        return arg
172
173    # Escape command line args.
174    argv = map(escape_arg, args[0])
175    # Cause script errors (such as adb not found) to stop script immediately with an error.
176    ps_code = u'$ErrorActionPreference = "Stop"\r\n';
177    # Add current directory to the PATH var, to match cmd.exe/CreateProcess() behavior.
178    ps_code += u'$env:Path = ".;" + $env:Path\r\n';
179    # Precede by &, the PowerShell call operator, and separate args by space.
180    ps_code += u'& ' + u' '.join(argv)
181    # Make the PowerShell exit code the exit code of the subprocess.
182    ps_code += u'\r\nExit $LastExitCode'
183    # Encode as UTF-16LE (without Byte-Order-Mark) which Windows natively understands.
184    ps_code = ps_code.encode('utf-16le')
185
186    # Encode the PowerShell command as base64 and use the special -EncodedCommand option that base64
187    # decodes. Base64 is just plain ASCII, so it should have no problem passing through Win32
188    # CreateProcessA() (which python erroneously calls instead of CreateProcessW()).
189    return (['powershell.exe', '-NoProfile', '-NonInteractive', '-EncodedCommand',
190             base64.b64encode(ps_code)],) + args[1:]
191
192
193# Call this instead of subprocess.check_output() to work-around issue in Python
194# 2's subprocess class on Windows where it doesn't support Unicode.
195def _subprocess_check_output(*args, **kwargs):
196    try:
197        return subprocess.check_output(*_get_subprocess_args(args), **kwargs)
198    except subprocess.CalledProcessError as e:
199        # Show real command line instead of the powershell.exe command line.
200        raise subprocess.CalledProcessError(e.returncode, args[0],
201                                            output=e.output)
202
203
204# Call this instead of subprocess.Popen(). Like _subprocess_check_output().
205def _subprocess_Popen(*args, **kwargs):
206    return subprocess.Popen(*_get_subprocess_args(args), **kwargs)
207
208
209class AndroidDevice(object):
210    # Delimiter string to indicate the start of the exit code.
211    _RETURN_CODE_DELIMITER = 'x'
212
213    # Follow any shell command with this string to get the exit
214    # status of a program since this isn't propagated by adb.
215    #
216    # The delimiter is needed because `printf 1; echo $?` would print
217    # "10", and we wouldn't be able to distinguish the exit code.
218    _RETURN_CODE_PROBE = [';', 'echo', '{0}$?'.format(_RETURN_CODE_DELIMITER)]
219
220    # Maximum search distance from the output end to find the delimiter.
221    # adb on Windows returns \r\n even if adbd returns \n.
222    _RETURN_CODE_SEARCH_LENGTH = len('{0}255\r\n'.format(_RETURN_CODE_DELIMITER))
223
224    # Feature name strings.
225    SHELL_PROTOCOL_FEATURE = 'shell_v2'
226
227    def __init__(self, serial, product=None, adb_path='adb'):
228        self.serial = serial
229        self.product = product
230        self.adb_cmd = [adb_path]
231
232        if self.serial is not None:
233            self.adb_cmd.extend(['-s', serial])
234        if self.product is not None:
235            self.adb_cmd.extend(['-p', product])
236        self._linesep = None
237        self._features = None
238
239    @property
240    def linesep(self):
241        if self._linesep is None:
242            self._linesep = subprocess.check_output(self.adb_cmd +
243                                                    ['shell', 'echo'])
244        return self._linesep
245
246    @property
247    def features(self):
248        if self._features is None:
249            try:
250                self._features = self._simple_call(['features']).splitlines()
251            except subprocess.CalledProcessError:
252                self._features = []
253        return self._features
254
255    def _make_shell_cmd(self, user_cmd):
256        command = self.adb_cmd + ['shell'] + user_cmd
257        if self.SHELL_PROTOCOL_FEATURE not in self.features:
258            command += self._RETURN_CODE_PROBE
259        return command
260
261    def _parse_shell_output(self, out):
262        """Finds the exit code string from shell output.
263
264        Args:
265            out: Shell output string.
266
267        Returns:
268            An (exit_code, output_string) tuple. The output string is
269            cleaned of any additional stuff we appended to find the
270            exit code.
271
272        Raises:
273            RuntimeError: Could not find the exit code in |out|.
274        """
275        search_text = out
276        if len(search_text) > self._RETURN_CODE_SEARCH_LENGTH:
277            # We don't want to search over massive amounts of data when we know
278            # the part we want is right at the end.
279            search_text = search_text[-self._RETURN_CODE_SEARCH_LENGTH:]
280        partition = search_text.rpartition(self._RETURN_CODE_DELIMITER)
281        if partition[1] == '':
282            raise RuntimeError('Could not find exit status in shell output.')
283        result = int(partition[2])
284        # partition[0] won't contain the full text if search_text was truncated,
285        # pull from the original string instead.
286        out = out[:-len(partition[1]) - len(partition[2])]
287        return result, out
288
289    def _simple_call(self, cmd):
290        logging.info(' '.join(self.adb_cmd + cmd))
291        return _subprocess_check_output(
292            self.adb_cmd + cmd, stderr=subprocess.STDOUT)
293
294    def shell(self, cmd):
295        """Calls `adb shell`
296
297        Args:
298            cmd: command to execute as a list of strings.
299
300        Returns:
301            A (stdout, stderr) tuple. Stderr may be combined into stdout
302            if the device doesn't support separate streams.
303
304        Raises:
305            ShellError: the exit code was non-zero.
306        """
307        exit_code, stdout, stderr = self.shell_nocheck(cmd)
308        if exit_code != 0:
309            raise ShellError(cmd, stdout, stderr, exit_code)
310        return stdout, stderr
311
312    def shell_nocheck(self, cmd):
313        """Calls `adb shell`
314
315        Args:
316            cmd: command to execute as a list of strings.
317
318        Returns:
319            An (exit_code, stdout, stderr) tuple. Stderr may be combined
320            into stdout if the device doesn't support separate streams.
321        """
322        cmd = self._make_shell_cmd(cmd)
323        logging.info(' '.join(cmd))
324        p = _subprocess_Popen(
325            cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
326        stdout, stderr = p.communicate()
327        if self.SHELL_PROTOCOL_FEATURE in self.features:
328            exit_code = p.returncode
329        else:
330            exit_code, stdout = self._parse_shell_output(stdout)
331        return exit_code, stdout, stderr
332
333    def shell_popen(self, cmd, kill_atexit=True, preexec_fn=None,
334                    creationflags=0, **kwargs):
335        """Calls `adb shell` and returns a handle to the adb process.
336
337        This function provides direct access to the subprocess used to run the
338        command, without special return code handling. Users that need the
339        return value must retrieve it themselves.
340
341        Args:
342            cmd: Array of command arguments to execute.
343            kill_atexit: Whether to kill the process upon exiting.
344            preexec_fn: Argument forwarded to subprocess.Popen.
345            creationflags: Argument forwarded to subprocess.Popen.
346            **kwargs: Arguments forwarded to subprocess.Popen.
347
348        Returns:
349            subprocess.Popen handle to the adb shell instance
350        """
351
352        command = self.adb_cmd + ['shell'] + cmd
353
354        # Make sure a ctrl-c in the parent script doesn't kill gdbserver.
355        if os.name == 'nt':
356            creationflags |= subprocess.CREATE_NEW_PROCESS_GROUP
357        else:
358            if preexec_fn is None:
359                preexec_fn = os.setpgrp
360            elif preexec_fn is not os.setpgrp:
361                fn = preexec_fn
362                def _wrapper():
363                    fn()
364                    os.setpgrp()
365                preexec_fn = _wrapper
366
367        p = _subprocess_Popen(command, creationflags=creationflags,
368                              preexec_fn=preexec_fn, **kwargs)
369
370        if kill_atexit:
371            atexit.register(p.kill)
372
373        return p
374
375    def install(self, filename, replace=False):
376        cmd = ['install']
377        if replace:
378            cmd.append('-r')
379        cmd.append(filename)
380        return self._simple_call(cmd)
381
382    def push(self, local, remote):
383        return self._simple_call(['push', local, remote])
384
385    def pull(self, remote, local):
386        return self._simple_call(['pull', remote, local])
387
388    def sync(self, directory=None):
389        cmd = ['sync']
390        if directory is not None:
391            cmd.append(directory)
392        return self._simple_call(cmd)
393
394    def tcpip(self, port):
395        return self._simple_call(['tcpip', port])
396
397    def usb(self):
398        return self._simple_call(['usb'])
399
400    def reboot(self):
401        return self._simple_call(['reboot'])
402
403    def remount(self):
404        return self._simple_call(['remount'])
405
406    def root(self):
407        return self._simple_call(['root'])
408
409    def unroot(self):
410        return self._simple_call(['unroot'])
411
412    def connect(self, host):
413        return self._simple_call(['connect', host])
414
415    def disconnect(self, host):
416        return self._simple_call(['disconnect', host])
417
418    def forward(self, local, remote):
419        return self._simple_call(['forward', local, remote])
420
421    def forward_list(self):
422        return self._simple_call(['forward', '--list'])
423
424    def forward_no_rebind(self, local, remote):
425        return self._simple_call(['forward', '--no-rebind', local, remote])
426
427    def forward_remove(self, local):
428        return self._simple_call(['forward', '--remove', local])
429
430    def forward_remove_all(self):
431        return self._simple_call(['forward', '--remove-all'])
432
433    def reverse(self, remote, local):
434        return self._simple_call(['reverse', remote, local])
435
436    def reverse_list(self):
437        return self._simple_call(['reverse', '--list'])
438
439    def reverse_no_rebind(self, local, remote):
440        return self._simple_call(['reverse', '--no-rebind', local, remote])
441
442    def reverse_remove_all(self):
443        return self._simple_call(['reverse', '--remove-all'])
444
445    def reverse_remove(self, remote):
446        return self._simple_call(['reverse', '--remove', remote])
447
448    def wait(self):
449        return self._simple_call(['wait-for-device'])
450
451    def get_props(self):
452        result = {}
453        output, _ = self.shell(['getprop'])
454        output = output.splitlines()
455        pattern = re.compile(r'^\[([^]]+)\]: \[(.*)\]')
456        for line in output:
457            match = pattern.match(line)
458            if match is None:
459                raise RuntimeError('invalid getprop line: "{}"'.format(line))
460            key = match.group(1)
461            value = match.group(2)
462            if key in result:
463                raise RuntimeError('duplicate getprop key: "{}"'.format(key))
464            result[key] = value
465        return result
466
467    def get_prop(self, prop_name):
468        output = self.shell(['getprop', prop_name])[0].splitlines()
469        if len(output) != 1:
470            raise RuntimeError('Too many lines in getprop output:\n' +
471                               '\n'.join(output))
472        value = output[0]
473        if not value.strip():
474            return None
475        return value
476
477    def set_prop(self, prop_name, value):
478        self.shell(['setprop', prop_name, value])
479