Home | History | Annotate | Download | only in cros 
      1 #!/usr/bin/python
      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 
      7 import argparse
      8 import logging
      9 import os
     10 import pipes
     11 import re
     12 import signal
     13 import sys
     14 import time
     15 
     16 import common
     17 from autotest_lib.client.bin import utils
     18 from autotest_lib.client.common_lib import error
     19 from 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 
     28 def _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 
     35 def _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 
     49 def _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 
     70 def _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 
     82 def _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 
    101 if __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