Overhaul Android support and make Swarming bot use python-adb

appengine/swarming/swarming_bot/os_utilities.py

Index: appengine/swarming/swarming_bot/os_utilities.py
diff --git a/appengine/swarming/swarming_bot/os_utilities.py b/appengine/swarming/swarming_bot/os_utilities.py
index a6bbbcd2be218202e7ddc1f510d880a2123c4d96..02b3760ab3e9cd0ae9e098578500bae6ab0194d9 100755
--- a/appengine/swarming/swarming_bot/os_utilities.py
+++ b/appengine/swarming/swarming_bot/os_utilities.py
@@ -18,6 +18,7 @@ import cgi
 import ctypes
 import getpass
 import glob
+import hashlib
 import json
 import logging
 import multiprocessing
@@ -112,6 +113,19 @@ _CACHED_OAUTH2_TOKEN_GCE = {}
 _CACHED_OAUTH2_TOKEN_GCE_LOCK = threading.Lock()
 
 
+# Set this to package adb and M2Crypto.RSA when Android support is initialized.
+adb = None
+RSA = None
+
+# Set when ADB is initialized. It contains one or multiple key used to
+# authenticate to Android debug protocol (adb).
+_ADB_KEYS = []
+
+
+# Cache of /system/build.prop on Android devices connected to this host.
+_BUILD_PROP_ANDROID = {}
+
+
 def _write(filepath, content):
   """Writes out a file and returns True on success."""
   logging.info('Writing in %s:\n%s', filepath, content)
@@ -1206,6 +1220,346 @@ def get_integrity_level_win():
 ### Android.
 
 
+def initialize_android(pub_key, priv_key):

[Vadim Sh., 2015/08/20 22:40:09]

how do you feel about moving it to "android_utilities.py"? os_utilities is
already huge.

[M-A Ruel, 2015/08/21 23:37:39]

Done.

+  """Initialize Android support through adb.
+
+  You can steal pub_key, priv_key pair from ~/.android/adbkey and
+  ~/.android/adbkey.pub.

[Vadim Sh., 2015/08/20 22:40:08]

is it true only for chrome-infra bots?

[M-A Ruel, 2015/08/21 23:37:39]

It's created by adb.

+  """
+  global adb
+  global RSA
+  assert bool(pub_key) == bool(priv_key)
+
+  if adb is None:
+    try:
+      # Workaround pylint lint errors.

[Vadim Sh., 2015/08/20 22:40:09]

just disable them, it will look less weird

[M-A Ruel, 2015/08/21 23:37:39]

Not necessary anymore.

+      import adb as adb2
+      adb = adb2
+      # This is needed otherwise adb_commands is not defined.
+      from adb import adb_commands as _
+      from M2Crypto import RSA as RSA2
+      RSA = RSA2
+      logging.info('adb fully loaded')
+    except ImportError as e:
+      adb = False
+      logging.error('adb is unavailable: %s' % e)
+      return False
+
+    if pub_key:
+      _ADB_KEYS.append(M2CryptoSigner(pub_key, priv_key))
+
+    # Try to add local adb keys if available.
+    path = os.path.expanduser('~/.android/adbkey')

[Vadim Sh., 2015/08/20 22:40:08]

is it intentionally different from priv_key? What is priv_key then? Add a doc
string for priv_key and pub_key.

[M-A Ruel, 2015/08/21 23:37:39]

The private and public keys for adb are stored in separate files.
The private key is stored in a standard format but the public key has its own
encoding.

The adb package allows multiple keys. We'd have one fleet wise and a backup key
of adb, so users that authenticated with their key will still work just fine.

+    if os.path.isfile(path) and os.path.isfile(path+'.pub'):
+      with open(path + '.pub', 'rb') as f:
+        pub = f.read()
+      with open(path, 'rb') as f:
+        priv = f.read()
+      _ADB_KEYS.append(M2CryptoSigner(pub, priv))
+
+    if not _ADB_KEYS:
+      adb = False
+      logging.error('No key is provided, android support is disabled.')
+      return False
+  else:
+    if pub_key:
+      logging.warning(
+          'initialize_android() was called repeatedly: ignoring keys')
+  return bool(adb)
+
+
+class M2CryptoSigner(object):
+  """Implements adb_protocol.AuthSigner using
+  https://github.com/martinpaljak/M2Crypto.
+  """
+  def __init__(self, pub, priv):
+    self.priv_key = RSA.load_key_string(priv)
+    self.pub_key = pub
+
+  def Sign(self, data):
+    return self.priv_key.sign(data, 'sha1')
+
+  def GetPublicKey(self):
+    return self.pub_key
+
+
+# TODO(maruel): M2Crypto is not included by default on Ubuntu.
+# rsa is included in client/third_party/rsa/rsa/ and
+# pycrypto is normally installed on Ubuntu. It would be preferable to use one of
+# these 2 but my skills failed up to now, authentication consistently fails.
+# Revisit later or delete the code.
+#
+#
+#sys.path.insert(0, os.path.join(THIS_FILE, 'third_party', 'rsa'))
+#import rsa
+#
+#class RSASigner(object):
+#  """Implements adb_protocol.AuthSigner using http://stuvel.eu/rsa."""
+#  def __init__(self):
+#    self.privkey = rsa.PrivateKey.load_pkcs1(PRIV_CONVERTED_KEY)
+#
+#  def Sign(self, data):
+#    return rsa.sign(data, self.privkey, 'SHA-1')
+#
+#  def GetPublicKey(self):
+#    return PUB_KEY
+#
+#
+#try:
+#  from Crypto.Hash import SHA
+#  from Crypto.PublicKey import RSA
+#  from Crypto.Signature import PKCS1_v1_5
+#  from Crypto.Signature import PKCS1_PSS
+#except ImportError:
+#  SHA = None
+#
+#
+#class CryptoSigner(object):
+#  """Implements adb_protocol.AuthSigner using
+#  https://www.dlitz.net/software/pycrypto/.
+#  """
+#  def __init__(self):
+#    self.private_key = RSA.importKey(PRIV_KEY, None)
+#    self._signer = PKCS1_v1_5.new(self.private_key)
+#    #self.private_key = RSA.importKey(PRIV_CONVERTED_KEY, None)
+#    #self._signer = PKCS1_PSS.new(self.private_key)
+#
+#  def Sign(self, data):
+#    h = SHA.new(data)
+#    return self._signer.sign(h)
+#
+#  def GetPublicKey(self):
+#    return PUB_KEY
+
+
+def kill_adb_android():
+  """adb sucks. Kill it with fire."""
+  if not adb:
+    return
+  try:
+    subprocess.call(['adb', 'kill-server'])
+  except OSError:
+    pass
+  subprocess.call(['killall', 'adb'])

[Vadim Sh., 2015/08/20 22:40:08]

killall --exact adb

(In general I'm not a fan of blindly killing processes)

[M-A Ruel, 2015/08/21 23:37:39]

Added --exact. But that's adb, it needs to be killed afaik.

+
+
+def get_devices_android():
+  """Returns the list of devices available.
+
+  Caller MUST call close_devices_android(cmds) on the return value.
+
+  Returns one of:
+    - dict of {serial_number: adb.adb_commands.AdbCommands}. The value may be
+      None if there was an Auth failure.
+    - None if adb is unavailable.
+  """
+  if not adb:
+    return None
+
+  cmds = {}
+  for handle in adb.adb_commands.AdbCommands.Devices():
+    logging.info('Trying Android %s', handle.serial_number)
+    try:
+      handle.Open()
+    except adb.common.usb1.USBErrorBusy:
+      logging.warning('Got USBErrorBusy. Killing adb')
+      kill_adb_android()
+      # TODO(maruel): Handle when it throws again, which means another process
+      # holds an handle to the USB ports or group acl (plugdev) hasn't been

[Vadim Sh., 2015/08/20 22:40:09]

typo: a handle

[M-A Ruel, 2015/08/21 23:37:39]

Done.

+      # setup properly.
+      handle.Open()

[Vadim Sh., 2015/08/20 22:40:09]

maybe at least log exception here and ignore this device?

[M-A Ruel, 2015/08/21 23:37:39]

Done.

+
+    try:
+      cmd = adb.adb_commands.AdbCommands.Connect(
+            handle, banner='swarming', rsa_keys=_ADB_KEYS, auth_timeout_ms=100)

[Vadim Sh., 2015/08/20 22:40:08]

100ms? Is it really that fast?

[M-A Ruel, 2015/08/21 23:37:39]

Yes.

+    except adb.usb_exceptions.DeviceAuthError as e:
+      logging.warning('AUTH FAILURE: %s: %s', handle.serial_number, e)
+      cmd = None
+    except adb.usb_exceptions.ReadFailedError as e:
+      logging.warning('READ FAILURE: %s: %s', handle.serial_number, e)
+      cmd = None
+    except ValueError as e:
+      logging.warning(
+          'Trying unpluging and pluging it back: %s: %s',
+          handle.serial_number, e)
+      cmd = None
+    cmds[handle.serial_number] = cmd
+
+  logging.info('FOUND %d devices', len(cmds))
+  # Remove any /system/build.prop cache so if a device is disconnect, reflashed
+  # then reconnected, it will likely be refresh properly. The main concern is
+  # that the bot didn't have the time to loop once while this is being done.
+  # Restarting the bot works fine too.
+  for key in _BUILD_PROP_ANDROID.keys():
+    if key not in cmds:
+      _BUILD_PROP_ANDROID.pop(key)
+  return cmds
+
+
+def close_devices_android(devices):
+  """Closes all devices opened by get_devices_android()."""
+  for device in (devices or {}).itervalues():
+    if device:
+      device.Close()

[Vadim Sh., 2015/08/20 22:40:09]

catch and log exceptions here? (so that all devices are closed even if one
failed)

[M-A Ruel, 2015/08/21 23:37:39]

I don't know which exception type would be thrown (thanks python) and I don't
expect exceptions to be thrown here.

+
+
+def get_build_prop_android(cmd):
+  """Returns the system properties for a device.
+
+  This isn't really changing through the lifetime of a bot. One corner case is
+  when the device is flashed or disconnected.
+  """
+  if cmd.handle.serial_number not in _BUILD_PROP_ANDROID:
+    properties = {}
+    try:
+      out = cmd.Shell('cat /system/build.prop').decode('utf-8')
+    except adb.usb_exceptions.ReadFailedError:
+      # It's a bit annoying because it means timeout_ms was wasted. Blacklist
+      # the device until it is disconnected and reconnected.
+      properties = None
+    else:
+      for line in out.splitlines():
+        if line.startswith(u'#') or not line:
+          continue
+        key, value = line.split(u'=', 1)
+        properties[key] = value
+    _BUILD_PROP_ANDROID[cmd.handle.serial_number] = properties
+  return _BUILD_PROP_ANDROID[cmd.handle.serial_number]
+
+
+def get_temp_android(cmd):
+  """Returns the device's 2 temperatures."""
+  temps = []
+  for i in xrange(2):
+    try:
+      temps.append(
+          int(cmd.Shell('cat /sys/class/thermal/thermal_zone%d/temp' % i)))
+    except ValueError:
+      pass
+  return temps
+
+
+def get_battery_android(cmd):
+  """Returns details about the battery's state."""
+  props = {}
+  out = cmd.Shell('dumpsys battery').decode('utf-8', 'replace')
+  for line in out.splitlines():
+    if line.endswith(u':'):
+      continue
+    key, value = line.split(u': ', 2)
+    props[key.lstrip()] = value
+  out = {u'power': []}
+  if props[u'AC powered'] == u'true':
+    out[u'power'].append(u'AC')
+  if props[u'USB powered'] == u'true':
+    out[u'power'].append(u'USB')
+  if props[u'Wireless powered'] == u'true':
+    out[u'power'].append(u'Wireless')
+  for key in (u'health', u'level', u'status', u'temperature', u'voltage'):
+    out[key] = props[key]
+  return out
+
+
+def get_disk_android(cmd):
+  """Returns details about the battery's state."""
+  props = {}
+  out = cmd.Shell('dumpsys diskstats').decode('utf-8', 'replace')
+  for line in out.splitlines():
+    if line.endswith(u':'):
+      continue
+    key, value = line.split(u': ', 2)
+    match = re.match(u'^(\d+)K / (\d+)K.*', value)
+    if match:
+      props[key.lstrip()] = {
+        'free_mb': round(float(match.group(1)) / 1024., 1),
+        'size_mb': round(float(match.group(2)) / 1024., 1),
+      }
+  return {
+    u'cache': props[u'Cache-Free'],
+    u'data': props[u'Data-Free'],
+    u'system': props[u'System-Free'],
+  }
+
+
+def get_dimensions_all_devices_android(cmds):
+  """Returns the default dimensions for an host with multiple android devices.
+  """
+  dimensions = get_dimensions()
+  if not cmds:
+    return dimensions
+
+  # Pop a few dimensions otherwise there will be too many dimensions.
+  del dimensions[u'cpu']
+  del dimensions[u'cores']
+  del dimensions[u'gpu']
+  dimensions.pop(u'machine_type')
+
+  # Make sure all the devices use the same board.
+  keys = (u'build.id', u'product.board')
+  for key in keys:
+    dimensions[key] = set()
+  dimensions[u'android'] = []
+  for serial_number, cmd in cmds.iteritems():
+    if cmd:
+      properties = get_build_prop_android(cmd)
+      if properties:
+        for key in keys:
+          dimensions[key].add(properties[u'ro.' + key])
+    dimensions[u'android'].append(serial_number)
+  dimensions[u'android'].sort()
+  for key in keys:
+    if not dimensions[key]:
+      del dimensions[key]
+    else:
+      dimensions[key] = sorted(dimensions[key])
+  nb_android = len(dimensions[u'android'])
+  dimensions[u'android_devices'] = map(
+      str, range(nb_android, max(0, nb_android-2), -1))
+  return dimensions
+
+
+def get_state_all_devices_android(cmds):
+  """Returns state information about all the devices connected to the host.
+  """
+  state = get_state()
+  if not cmds:
+    return state
+
+  # Add a few values that were poped from dimensions.
+  cpu_type = get_cpu_type()
+  cpu_bitness = get_cpu_bitness()
+  state[u'cpu'] = [
+    cpu_type,
+    cpu_type + u'-' + cpu_bitness,
+  ]
+  state[u'cores'] = [unicode(get_num_processors())]
+  state[u'gpu'] = get_gpu()[0]
+  machine_type = get_machine_type()
+  if machine_type:
+    state[u'machine_type'] = [machine_type]
+
+  keys = (
+      u'board.platform', u'product.cpu.abi', u'product.cpu.abi2',
+      u'build.tags', u'build.type', u'build.version.sdk')
+  state['devices'] = {}
+  for serial_number, cmd in cmds.iteritems():
+    if cmd:
+      properties = get_build_prop_android(cmd)
+      if properties:
+        # TODO(maruel): uptime, diskstats, wifi, power, throttle, etc.
+        device = {
+          u'build': {key: properties[u'ro.'+key] for key in keys},
+          u'disk': get_disk_android(cmd),
+          u'battery': get_battery_android(cmd),
+          u'state': u'available',
+          u'temp': get_temp_android(cmd),
+        }
+      else:
+        device = {u'state': u'unavailable'}
+    else:
+      device = {u'state': 'unauthenticated'}
+    state[u'devices'][serial_number] = device
+  return state
 
 
 ###
@@ -1536,10 +1890,12 @@ def trim_rolled_log(name):
 def main():
   """Prints out the output of get_dimensions() and get_state()."""
   # Pass an empty tag, so pop it up since it has no significance.
+  devices = get_devices_android()
   data = {
-    u'dimensions': get_dimensions(),
-    u'state': get_state(),
+    u'dimensions': get_dimensions_all_devices_android(devices),
+    u'state': get_state_all_devices_android(devices),
   }
+  close_devices_android(devices)

[Vadim Sh., 2015/08/20 22:40:08]

put close_devices_android(...) in finally section. (Because "Caller MUST call
close_devices_android")

[M-A Ruel, 2015/08/21 23:37:39]

Done.

   json.dump(data, sys.stdout, indent=2, sort_keys=True, separators=(',', ': '))
   print('')
   return 0