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