1#!/usr/bin/python2
2
3# Copyright 2017 The Chromium OS Authors. All rights reserved.
4# Use of this source code is governed by a BSD-style license that can be
5# found in the LICENSE file.
6
7import argparse
8import logging
9import os
10import pipes
11import re
12import signal
13import sys
14import time
15
16import common
17from autotest_lib.client.bin import utils
18from autotest_lib.client.common_lib import error
19from autotest_lib.client.common_lib import logging_config
20
21_ADB_POLLING_INTERVAL_SECONDS = 10
22_ADB_CONNECT_INTERVAL_SECONDS = 1
23_ADB_COMMAND_TIMEOUT_SECONDS = 5
24
25_signum_to_name = {}
26
27
28def _signal_handler(signum, frame):
29    logging.info('Received %s, shutting down', _signum_to_name[signum])
30    sys.stdout.flush()
31    sys.stderr.flush()
32    os._exit(0)
33
34
35def _get_adb_options(target, socket):
36    """Get adb global options."""
37    # ADB 1.0.36 does not support -L adb socket option. Parse the host and port
38    # part from the socket instead.
39    # https://developer.android.com/studio/command-line/adb.html#issuingcommands
40    pattern = r'^[^:]+:([^:]+):(\d+)$'
41    match = re.match(pattern, socket)
42    if not match:
43        raise ValueError('Unrecognized socket format: %s' % socket)
44    server_host, server_port = match.groups()
45    return '-s %s -H %s -P %s' % (
46        pipes.quote(target), pipes.quote(server_host), pipes.quote(server_port))
47
48
49def _run_adb_cmd(cmd, adb_option="", **kwargs):
50    """Run adb command.
51
52    @param cmd: command to issue with adb. (Ex: connect, devices)
53    @param target: Device to connect to.
54    @param adb_option: adb global option configuration.
55
56    @return: the stdout of the command.
57    """
58    adb_cmd = 'adb %s %s' % (adb_option, cmd)
59    while True:
60        try:
61            output = utils.system_output(adb_cmd, **kwargs)
62            break
63        except error.CmdTimeoutError as e:
64            logging.warning(e)
65            logging.info('Retrying command %s', adb_cmd)
66    logging.debug('%s: %s', adb_cmd, output)
67    return output
68
69
70def _is_adb_connected(target, adb_option=""):
71    """Return true if adb is connected to the container.
72
73    @param target: Device to connect to.
74    @param adb_option: adb global option configuration.
75    """
76    output = _run_adb_cmd('get-state', adb_option=adb_option,
77                          timeout=_ADB_COMMAND_TIMEOUT_SECONDS,
78                          ignore_status=True)
79    return output.strip() == 'device'
80
81
82def _ensure_adb_connected(target, adb_option=""):
83    """Ensures adb is connected to the container, reconnects otherwise.
84
85    @param target: Device to connect to.
86    @param adb_option: adb global options configuration.
87    """
88    did_reconnect = False
89    while not _is_adb_connected(target, adb_option):
90        if not did_reconnect:
91            logging.info('adb not connected. attempting to reconnect')
92            did_reconnect = True
93        _run_adb_cmd('connect %s' % pipes.quote(target),
94                     adb_option=adb_option,
95                     timeout=_ADB_COMMAND_TIMEOUT_SECONDS, ignore_status=True)
96        time.sleep(_ADB_CONNECT_INTERVAL_SECONDS)
97    if did_reconnect:
98        logging.info('Reconnection succeeded')
99
100
101if __name__ == '__main__':
102    logging_config.LoggingConfig().configure_logging()
103    parser = argparse.ArgumentParser(description='ensure adb is connected')
104    parser.add_argument('target', help='Device to connect to')
105    parser.add_argument('--socket', help='ADB server socket.',
106                        default='tcp:localhost:5037')
107    args = parser.parse_args()
108    adb_option = _get_adb_options(args.target, args.socket)
109
110    # Setup signal handler for logging on exit
111    for attr in dir(signal):
112        if not attr.startswith('SIG') or attr.startswith('SIG_'):
113            continue
114        if attr in ('SIGCHLD', 'SIGCLD', 'SIGKILL', 'SIGSTOP'):
115            continue
116        signum = getattr(signal, attr)
117        _signum_to_name[signum] = attr
118        signal.signal(signum, _signal_handler)
119
120    logging.info('Starting adb_keepalive for target %s on socket %s',
121                 args.target, args.socket)
122    while True:
123        time.sleep(_ADB_POLLING_INTERVAL_SECONDS)
124        _ensure_adb_connected(args.target, adb_option=adb_option)
125