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