Initial work on an Android ADB connection plugin for Ansible

adb.py

# Copyright (c) 2012, Michael DeHaan <[email protected]>
# Copyright 2015 Abhijit Menon-Sen <[email protected]>
# Copyright 2017 Toshio Kuratomi <[email protected]>
# Copyright (c) 2017 Ansible Project
# GNU General Public License v3.0+ (see COPYING or https://www.gnu.org/licenses/gpl-3.0.txt)

from __future__ import absolute_import, division, print_function

__metaclass__ = type

DOCUMENTATION = """
    name: adb
    short_description: connect via adb client binary
    description:
        - This connection plugin allows Ansible to communicate to the target Android devices through adb shell.
    author: Rylie Pavlik (@rpavlik)
    options:
      serial:
          description: Serial number to connect to.
          vars:
               - name: inventory_hostname
               - name: ansible_adb_serial_num
               - name: delegated_vars['ansible_adb_serial_num']
      host:
          description: Hostname/IP to connect to.
          vars:
               - name: ansible_adb_host
               - name: delegated_vars['ansible_adb_host']
      adb_executable:
          default: adb
          description:
            - This defines the location of the ADB binary. It defaults to C(adb) which will use the first ADB binary available in $PATH.
            - This option is usually not required, it might be useful when access to system ADB is restricted,
              or when using SSH wrappers to connect to remote hosts.
          env: [{name: ANSIBLE_ADB_EXECUTABLE}]
          ini:
          - {key: adb_executable, section: adb_connection}
          #const: ANSIBLE_ADB_EXECUTABLE
          vars:
              - name: ansible_ADB_executable
                version_added: '2.10'
      adb_args:
          description: Arguments to pass to all ADB CLI tools.
          default:
          ini:
              - section: 'adb_connection'
                key: 'adb_args'
          env:
              - name: ANSIBLE_ADB_ARGS
          vars:
              - name: ansible_adb_args
      port:
          description: Remote port to connect to.
          type: int
          ini:
            - section: defaults
              key: remote_port
          env:
            - name: ANSIBLE_ADB_REMOTE_PORT
          vars:
            - name: ansible_adb_port
          keyword:
            - name: port
      timeout:
          default: 10
          description:
              - This is the default ammount of time we will wait while establishing an ADB connection.
              - It also controls how long we can wait to access reading the connection once established (select on the socket).
          env:
              - name: ANSIBLE_TIMEOUT
              - name: ANSIBLE_ADB_TIMEOUT
                version_added: '2.11'
          ini:
              - key: timeout
                section: defaults
              - key: timeout
                section: adb_connection
          vars:
            - name: ansible_adb_timeout
          cli:
            - name: timeout
          type: integer
      use_tty:
          default: 'yes'
          description: add -tt to adb shell commands to force tty allocation.
          env: [{name: ANSIBLE_ADB_USETTY}]
          ini:
          - {key: usetty, section: adb_connection}
          type: bool
          vars:
            - name: ansible_adb_use_tty
"""

import errno
import fcntl
import hashlib
import os
import pty
import re
import shlex
import subprocess
import time

from functools import wraps
from ansible.errors import (
    AnsibleAuthenticationFailure,
    AnsibleConnectionFailure,
    AnsibleError,
    AnsibleFileNotFound,
)
from ansible.errors import AnsibleOptionsError
from ansible.module_utils.compat import selectors
from ansible.module_utils.six import PY3, text_type, binary_type
from ansible.module_utils._text import to_bytes, to_native, to_text
from ansible.module_utils.parsing.convert_bool import BOOLEANS, boolean
from ansible.plugins.connection import ConnectionBase, BUFSIZE
from ansible.plugins.shell.powershell import _parse_clixml
from ansible.utils.display import Display
from ansible.utils.path import unfrackpath, makedirs_safe

display = Display()


b_NOT_SSH_ERRORS = (
    b"Traceback (most recent call last):",  # Python-2.6 when there's an exception
    # while invoking a script via -m
    b"PHP Parse error:",  # Php always returns error 255
)

SSHPASS_AVAILABLE = None
SSH_DEBUG = re.compile(r"^debug\d+: .*")


class Connection(ConnectionBase):
    """ adb based connections """

    transport = "adb"

    def __init__(self, *args, **kwargs):
        super(Connection, self).__init__(*args, **kwargs)

        # TODO: all should come from get_option(), but not might be set at this point yet
        self.host = self._play_context.remote_addr
        self.port = self._play_context.port

    def _connect(self):
        self.serial = self.get_option("serial")
        self.port = self.get_option("port")
        self.host = self.get_option("host")
        self.adb = to_bytes(self.get_option("adb_executable"))
        if self.port is not None and self.host is not None:
            cmd = self._build_command(
                self.adb, "connect", "{}:{}".format(self.host, self.port)
            )
            p = subprocess.Popen(
                cmd,
                stdin=subprocess.PIPE,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
            )
            stdout, stderr = p.communicate()
            status_code = p.wait()
            if status_code != 0:
                display.error(u"Failed to bring up connection: %s" % to_text(stderr))
        cmd = self._build_command(self.adb, "wait-for-device")
        p = subprocess.Popen(
            cmd,
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )
        stdout, stderr = p.communicate()
        status_code = p.wait()
        if status_code != 0:
            display.error(u"Failed to wait for device: %s" % to_text(stderr))

        self._connected = True
        return self

    def close(self):
        if self.port is not None and self.host is not None:
            cmd = self._build_command(
                self.adb, "disconnect", "{}:{}".format(self.host, self.port)
            )
            p = subprocess.Popen(
                cmd,
                stdin=subprocess.PIPE,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
            )
            stdout, stderr = p.communicate()
            status_code = p.wait()
            if status_code != 0:
                display.warning(u"Failed to disconnect: %s" % to_text(stderr))

            self._connected = False

    def _add_args(self, b_command, b_args, explanation):
        """
        Adds arguments to the adb command and displays a caller-supplied explanation of why.

        :arg b_command: A list containing the command to add the new arguments to.
            This list will be modified by this method.
        :arg b_args: An iterable of new arguments to add.  This iterable is used
            more than once so it must be persistent (ie: a list is okay but a
            StringIO would not)
        :arg explanation: A text string containing explaining why the arguments
            were added.  It will be displayed with a high enough verbosity.
        .. note:: This function does its work via side-effect.  The b_command list has the new arguments appended.
        """
        display.vvvvv(
            u"ADB: %s: (%s)" % (explanation, ")(".join(to_text(a) for a in b_args)),
            host=(self.host if self.host else self.serial),
        )
        b_command += b_args

    def _build_command(self, binary, *other_args):
        """
        Takes a executable (adb or wrapper) and optional extra arguments and returns the remote command
        wrapped in local ssh shell commands and ready for execution.

        :arg binary: actual executable to use to execute command.
        :arg other_args: dict of, value pairs passed as arguments to the adb binary

        """

        b_command = []

        #
        # First, the command to invoke
        #

        b_command += [to_bytes(binary, errors="surrogate_or_strict")]

        #
        # Next, additional arguments based on the configuration.
        #

        if self._play_context.verbosity > 3:
            b_command.append(b"-vvv")

        # Next, we add adb_args
        adb_args = self.get_option("adb_args")
        if adb_args:
            b_args = [
                to_bytes(a, errors="surrogate_or_strict")
                for a in self._split_ssh_args(adb_args)
            ]
            self._add_args(b_command, b_args, u"ansible.cfg set adb_args")


        # Finally, we add any caller-supplied extras.
        if other_args:
            b_command += [to_bytes(a) for a in other_args]

        return b_command

    # Used by _run() to kill processes on failures
    @staticmethod
    def _terminate_process(p):
        """ Terminate a process, ignoring errors """
        try:
            p.terminate()
        except (OSError, IOError):
            pass

    def _bare_run(self, cmd, in_data, sudoable=True, checkrc=True):
        """
        Starts the command and communicates with it until it ends.
        """

        # We don't use _shell.quote as this is run on the controller and independent from the shell plugin chosen
        display_cmd = u" ".join(shlex.quote(to_text(c)) for c in cmd)
        display.vvv(u"ADB: EXEC {0}".format(display_cmd), host=self.serial)

        # Start the given command. If we don't need to pipeline data, we can try
        # to use a pseudo-tty (ssh will have been invoked with -tt). If we are
        # pipelining data, or can't create a pty, we fall back to using plain
        # old pipes.

        p = None

        if isinstance(cmd, (text_type, binary_type)):
            cmd = to_bytes(cmd)
        else:
            cmd = list(map(to_bytes, cmd))

        if not in_data:
            try:
                # Make sure stdin is a proper pty to avoid tcgetattr errors
                master, slave = pty.openpty()
                p = subprocess.Popen(
                    cmd, stdin=slave, stdout=subprocess.PIPE, stderr=subprocess.PIPE
                )
                stdin = os.fdopen(master, "wb", 0)
                os.close(slave)
            except (OSError, IOError):
                p = None

        if not p:
            try:
                p = subprocess.Popen(
                    cmd,
                    stdin=subprocess.PIPE,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.PIPE,
                )
                stdin = p.stdin
            except (OSError, IOError) as e:
                raise AnsibleError(
                    "Unable to execute adb command line on a controller due to: %s"
                    % to_native(e)
                )

        #
        # ADB shell state machine
        #

        # Now we read and accumulate output from the running process until it
        # exits. Depending on the circumstances, we may also need to write
        # pipelined input to the process.

        states = [
            "ready_to_send",
            "awaiting_exit",
        ]

        state = states.index("ready_to_send")

        # We store accumulated stdout and stderr output from the process here,
        # but strip any privilege escalation prompt/confirmation lines first.
        # Output is accumulated into tmp_*, complete lines are extracted into
        # an array, then checked and removed or copied to stdout or stderr. We
        # set any flags based on examining the output in self._flags.

        b_stdout = b_stderr = b""
        b_tmp_stdout = b_tmp_stderr = b""

        self._flags = dict(
            become_prompt=False,
            become_success=False,
            become_error=False,
            become_nopasswd_error=False,
        )

        # select timeout should be longer than the connect timeout, otherwise
        # they will race each other when we can't connect, and the connect
        # timeout usually fails
        timeout = 2 + self.get_option("timeout")
        for fd in (p.stdout, p.stderr):
            fcntl.fcntl(
                fd, fcntl.F_SETFL, fcntl.fcntl(fd, fcntl.F_GETFL) | os.O_NONBLOCK
            )

        # TODO: bcoca would like to use SelectSelector() when open
        # select is faster when filehandles is low and we only ever handle 1.
        selector = selectors.DefaultSelector()
        selector.register(p.stdout, selectors.EVENT_READ)
        selector.register(p.stderr, selectors.EVENT_READ)

        # If we can send initial data without waiting for anything, we do so
        # before we start polling
        if states[state] == "ready_to_send" and in_data:
            self._send_initial_data(stdin, in_data, p)
            state += 1

        try:
            while True:
                poll = p.poll()
                events = selector.select(timeout)

                # We pay attention to timeouts only while negotiating a prompt.

                if not events:
                    # We timed out
                    if state <= states.index("awaiting_escalation"):
                        # If the process has already exited, then it's not really a
                        # timeout; we'll let the normal error handling deal with it.
                        if poll is not None:
                            break
                        self._terminate_process(p)
                        raise AnsibleError(
                            "Timeout (%ds) waiting for privilege escalation prompt: %s"
                            % (timeout, to_native(b_stdout))
                        )

                # Read whatever output is available on stdout and stderr, and stop
                # listening to the pipe if it's been closed.

                for key, event in events:
                    if key.fileobj == p.stdout:
                        b_chunk = p.stdout.read()
                        if b_chunk == b"":
                            # stdout has been closed, stop watching it
                            selector.unregister(p.stdout)
                            # When ssh has ControlMaster (+ControlPath/Persist) enabled, the
                            # first connection goes into the background and we never see EOF
                            # on stderr. If we see EOF on stdout, lower the select timeout
                            # to reduce the time wasted selecting on stderr if we observe
                            # that the process has not yet existed after this EOF. Otherwise
                            # we may spend a long timeout period waiting for an EOF that is
                            # not going to arrive until the persisted connection closes.
                            timeout = 1
                        b_tmp_stdout += b_chunk
                        display.debug(
                            u"stdout chunk (state=%s):\n>>>%s<<<\n"
                            % (state, to_text(b_chunk))
                        )
                    elif key.fileobj == p.stderr:
                        b_chunk = p.stderr.read()
                        if b_chunk == b"":
                            # stderr has been closed, stop watching it
                            selector.unregister(p.stderr)
                        b_tmp_stderr += b_chunk
                        display.debug(
                            "stderr chunk (state=%s):\n>>>%s<<<\n"
                            % (state, to_text(b_chunk))
                        )

                # We examine the output line-by-line until we have negotiated any
                # privilege escalation prompt and subsequent success/error message.
                # Afterwards, we can accumulate output without looking at it.

                b_stdout += b_tmp_stdout
                b_stderr += b_tmp_stderr
                b_tmp_stdout = b_tmp_stderr = b""

                # Once we're sure that the privilege escalation prompt, if any, has
                # been dealt with, we can send any initial data and start waiting
                # for output.

                if states[state] == "ready_to_send":
                    if in_data:
                        self._send_initial_data(stdin, in_data, p)
                    state += 1

                # Now we're awaiting_exit: has the child process exited? If it has,
                # and we've read all available output from it, we're done.

                if poll is not None:
                    if not selector.get_map() or not events:
                        break
                    # We should not see further writes to the stdout/stderr file
                    # descriptors after the process has closed, set the select
                    # timeout to gather any last writes we may have missed.
                    timeout = 0
                    continue

                # If the process has not yet exited, but we've already read EOF from
                # its stdout and stderr (and thus no longer watching any file
                # descriptors), we can just wait for it to exit.

                elif not selector.get_map():
                    p.wait()
                    break

                # Otherwise there may still be outstanding data to read.
        finally:
            selector.close()
            # close stdin, stdout, and stderr after process is terminated and
            # stdout/stderr are read completely (see also issues #848, #64768).
            stdin.close()
            p.stdout.close()
            p.stderr.close()

        if p.returncode == 255:

            additional = to_native(b_stderr)
            if in_data and checkrc:
                raise AnsibleConnectionFailure(
                    'Data could not be sent to remote host "%s". Make sure this host can be reached over adb: %s'
                    % (self.serial, additional)
                )

        return (p.returncode, b_stdout, b_stderr)

    #
    # Main public methods
    #
    def exec_command(self, cmd, in_data=None, sudoable=None):
        """ run a command on the remote host """

        super(Connection, self).exec_command(cmd, in_data=in_data, sudoable=sudoable)

        display.vvv(u"ESTABLISH ADB CONNECTION", host=self.serial)

        # we can only use tty when we are not pipelining the modules. piping
        # data into /usr/bin/python inside a tty automatically invokes the
        # python interactive-mode but the modules are not compatible with the
        # interactive-mode ("unexpected indent" mainly because of empty lines)

        adb_executable = self.get_option("adb_executable")

        # -tt can cause various issues in some environments so allow the user
        # to disable it as a troubleshooting method.
        use_tty = self.get_option("use_tty")

        if not in_data and use_tty:
            args = ("-tt", "-s", self.serial, "shell", cmd)
        else:
            args = ("-s", self.serial, "shell", cmd)

        cmd = self._build_command(adb_executable, *args)
        (returncode, stdout, stderr) = self._bare_run(cmd, in_data)

        return (returncode, stdout, stderr)

    def put_file(self, in_path, out_path):
        """ transfer a file from local to remote """

        super(Connection, self).put_file(in_path, out_path)

        display.vvv(u"PUT {0} TO {1}".format(in_path, out_path), host=self.serial)
        if not os.path.exists(to_bytes(in_path, errors="surrogate_or_strict")):
            raise AnsibleFileNotFound(
                "file or module does not exist: {0}".format(to_native(in_path))
            )

        if getattr(self._shell, "_IS_WINDOWS", False):
            out_path = self._escape_win_path(out_path)

        cmd = self._build_command(
            self.adb, "-s", self.serial, "push", in_path, out_path
        )
        return self._bare_run(cmd, None, None)

    def fetch_file(self, in_path, out_path):
        """ fetch a file from remote to local """

        super(Connection, self).fetch_file(in_path, out_path)

        display.vvv(u"FETCH {0} TO {1}".format(in_path, out_path), host=self.serial)

        # need to add / if path is rooted
        if getattr(self._shell, "_IS_WINDOWS", False):
            in_path = self._escape_win_path(in_path)

        cmd = self._build_command(
            self.adb, "-s", self.serial, "pull", in_path, out_path
        )
        return self._bare_run(cmd, None, None)

    def reset(self):

        self.close()