catatester/devicemap/bot_setup.py

"""A helper script to setup and start bots on a host machine.

  This script does following:
    - Parses devicemap.yaml and finds out how many bots are needed for a host.
    - Check if bot directories exist and create bot directories as needed.
    - Create systemd configs and bot startup script.
    - Run systemctl start swarming_bot--<bot number> for any bots not currently
      running

NOTE: This script is being invoked from a puppet manifest and will not have
access to rest of the test code or python virtual environment. So keep in mind
to not use any non standard python libraries.

Example usage:
  python bot_setup.py bot_setup.py \
    --user gtvchrome --group gtvchrome --swarm-bots-home /home/bots_home \
    --gsutil-dir /home/gsutil_home --gsutil-boto /home/.boto \
    --fragglerock-key /home/fragglerock_service_key.json \
    --eurtest-binaries /home/eurtest_binaries --venv /home/bots-home/swarm_env
"""

import argparse
import errno
import grp
import logging
import os
import pwd
import shutil
import socket
import string
import subprocess
import sys
import yaml

# This code runs on catatester host and doesn't have whole code tree, hence
# do relative import.
try:
  from catatester.devicemap import devicemap_lib
except ImportError:
  import devicemap_lib  # pylint: disable=relative-import


_BOT_FILE_MODE = 0o774  # permit owner group execution
_SYSTEMD_FILE_MODE = 0o644  # permit owner group write
_ROOT_USER = 'root'
# TODO(adrexler): Remove _DEVICEMAP_DIR once the arg is being passed from
# puppet.
_DEVICEMAP_DIR = os.path.dirname(os.path.realpath(__file__))
_STARTUP_SCRIPT_TEMPLATE_PATH = os.path.join(
    _DEVICEMAP_DIR, 'swarm_bot_startup.sh.template')
_STARTUP_SCRIPTS_PATH = 'startup_scripts'
_SWARM_BOT_PREFIX = 'swarming_bot--'
_SWARM_ZIP_RELATIVE_PATH_TO_BOT_HOME = '1/swarming_bot.zip'
_SYSTEMD_TEMPLATE_PATH = os.path.join(
    _DEVICEMAP_DIR, 'swarm_bot.systemd.template')
_DEPOT_TOOLS_DIR = os.path.join(os.path.expanduser('~'), 'depot_tools')


class OSException(Exception):
  """Exception class to handle OS Exceptions."""


def _GetNumBotConfigs(hostname):
  """Gets the number of bot configs assigned to the given hostname.

  Args:
    hostname: (String) Catatester hostame.

  Returns:
    An integer, The number of bot configs for given hostname.
  """
  logging.debug('Finding bots needed for host [%s]', hostname)
  configs = devicemap_lib.Configs()
  return configs.GetNumBotConfigs(hostname)


def CreateStartupScript(
    user, group, swarm_bots_home, bot_num, bot_name, gsutil_dir,
    gsutil_boto, fragglerock_key, eurtest_binaries, venv,
    devicemap_dir, depot_tools_dir):
  """Create the bot startup script.

  This method loads tokenized swarm_bot_startup.sh.template file, replaces the
  tokens with given param values and writes it into bot folder.

  Args:
    user: (String) User name of the catatester host.
    group: (String) Group name that is user is part of.
    swarm_bots_home: (String) Root directory of all the bots.
    bot_num: (Integer) The N-th bot to be created.
    bot_name: (String) Bot service name.
    gsutil_dir: (String) Path to gsutil installation in the host.
    gsutil_boto: (String) Path to gsutil boto file in the host.
    fragglerock_key: (String) Path to auth key file in the host.
    eurtest_binaries: (String) Path to test binaries in the host.
    venv: (String) Path to python virtual environment to be used.
    devicemap_dir: (String) Path to devicemap directory on host.
    depot_tools_dir: (String) Path to depot tools directory.

  Returns:
    The path of newly created startup script.
  """
  bot_dir_path = os.path.join(swarm_bots_home, str(bot_num))
  with open(_STARTUP_SCRIPT_TEMPLATE_PATH, 'r') as bot_startup:
    bot_startup_template = bot_startup.read()
  startup_script_content = string.Template(bot_startup_template).substitute(
      bot_dir=bot_dir_path,
      devicemap_dir=devicemap_dir,
      eurtest_binaries=eurtest_binaries,
      fragglerock_key=fragglerock_key,
      gsutil_dir=gsutil_dir,
      gsutil_boto=gsutil_boto,
      venv=venv,
      depot_tools_dir=depot_tools_dir
  )
  startup_script_path = os.path.join(swarm_bots_home, _STARTUP_SCRIPTS_PATH,
                                     '%s_startup.sh' % bot_name)
  _CreateFile(
      startup_script_path, startup_script_content, user, group, _BOT_FILE_MODE)
  return startup_script_path


def CreateSystemdServiceFile(
    bot_service_name, startup_script_path, user, group):
  """Create the systemd service script for a bot.

  This method loads tokenized swarm_bot.systemd.template file, replaces the
  tokens with given param values and writes it into systemd service folder.

  Args:
    bot_service_name: (String) The systemd service name to be used.
    startup_script_path: (String) Startup script path of the bot.
    user: (String) User name of the catatester host.
    group: (String) Group name that is user is part of.
  """
  with open(_SYSTEMD_TEMPLATE_PATH, 'r') as bot_systemd:
    bot_systemd_template = bot_systemd.read()
  systemd_service_content = string.Template(bot_systemd_template).substitute(
      bot_name=bot_service_name,
      group=group,
      user=user,
      startup_script_path=startup_script_path)
  systemd_service_path = '/etc/systemd/system/%s.service' % bot_service_name
  _CreateFile(
      systemd_service_path, systemd_service_content, _ROOT_USER, _ROOT_USER,
      _SYSTEMD_FILE_MODE)


def SetupBot(
    bot_num, user, group, swarm_bots_home, swarm_bot_prefix, swarm_zip_path,
    gsutil_dir, gsutil_boto, fragglerock_key, eurtest_binaries, venv,
    devicemap_dir, depot_tools_dir):
  """Setup the swarming bot.

  This method :
   1. Creates bot directory.
   2. Copies the swarm binary to bot directory.
   3. Creates start up scripts and systemd configs the bot.
   4. Do a systemd config reload.
   5. And finally start the bot service, if its not up.

  Args:
    bot_num: (Integer) The N-th bot to be created.
    user: (String) User name of the catatester host.
    group: (String) Group name that is user is part of.
    swarm_bots_home: (String) Root directory of the bots.
    swarm_bot_prefix: (String) Prefix to be used for bot service name.
    swarm_zip_path: (String) Swarming binary path in the host.
    gsutil_dir: (String) Path to gsutil installation in the host.
    gsutil_boto: (String) Path to gsutil boto file in the host.
    fragglerock_key: (String) Path to auth key file in the host.
    eurtest_binaries: (String) Path to test binaries in the host.
    venv: (String) Path to python virtual environment to be used.
    devicemap_dir: (String) Path to the device map directory.
    depot_tools_dir: (String) Path to depot tools directory.
  """
  logging.debug('Working on bot #%s', bot_num)
  # Create bot directory.
  bot_dir_path = os.path.join(swarm_bots_home, str(bot_num))
  if not os.path.exists(bot_dir_path):
    logging.info('Creating directory: %s', bot_dir_path)
    os.makedirs(bot_dir_path)
    _UpdateFilePermissions(bot_dir_path, user, group, _BOT_FILE_MODE)
  # Copy swarm binary
  if os.path.dirname(swarm_zip_path) != bot_dir_path:
    shutil.copy(swarm_zip_path, bot_dir_path)
    _UpdateFilePermissions(swarm_zip_path, user, group, _BOT_FILE_MODE)
  # Create startup scripts.
  bot_name = '%s%s' % (swarm_bot_prefix, bot_num)
  startup_script_path = CreateStartupScript(
      user, group, swarm_bots_home, bot_num, bot_name, gsutil_dir, gsutil_boto,
      fragglerock_key, eurtest_binaries, venv, devicemap_dir, depot_tools_dir)
  CreateSystemdServiceFile(bot_name, startup_script_path, user, group)
  # Reload the daemon, check bot status and start service if needed.
  if subprocess.call(['systemctl', 'daemon-reload']):
    sys.exit('Failed to reload systemctl daemon for bot # %s' % bot_num)
  if subprocess.call(['systemctl', 'is-active', '--quiet', bot_name]):
    logging.debug('Bot number [%s] not running. Starting now.', bot_num)
    if subprocess.call(['systemctl', 'start', bot_name]):
      sys.exit('Failed to start bot # %s' % bot_num)


def _UpdateFilePermissions(file_path, owner, group, mode):
  """Update file with specified mode and permissions."""
  uid = pwd.getpwnam(owner).pw_uid
  gid = grp.getgrnam(group).gr_gid
  os.chown(file_path, uid, gid)
  os.chmod(file_path, mode)


def _CreateFile(file_path, content, owner, group, mode):
  """Create file with specified content, mode and permissions."""
  if not os.path.exists(os.path.dirname(file_path)):
    try:
      os.makedirs(os.path.dirname(file_path))
      _UpdateFilePermissions(os.path.dirname(file_path), owner, group, mode)
    except OSError as exc:
      if exc.errno != errno.EEXIST:
        raise OSException('Unable to create dir %s due to %s.'
                          % (os.path.dirname(file_path), exc))

  with open(file_path, 'w') as f:
    f.write(content)
  _UpdateFilePermissions(file_path, owner, group, mode)


def _ParseArgs():
  """Parses command-line input."""
  parser = argparse.ArgumentParser()
  parser.add_argument(
      '--user', help='User name of the catatester host.', required=True)
  parser.add_argument('--group', help='Group name to be used.', required=True)
  parser.add_argument(
      '--swarm_bots_home', help='Root directory of the bots.', required=True)
  parser.add_argument(
      '--gsutil_dir', help='Path to gsutil installation.', required=True)
  parser.add_argument(
      '--gsutil_boto', help='Path to gs util boto file.', required=True)
  parser.add_argument(
      '--fragglerock_key', help='Path to fragglerock auth key file.',
      required=True)
  parser.add_argument(
      '--eurtest_binaries', help='Test binary path.', required=True)
  parser.add_argument(
      '--venv', help='Path to python to virtual environment.', required=True)
  parser.add_argument(
      '--devicemap_dir', help='Path to devicemap directory on the host.',
      default=_DEVICEMAP_DIR)
  parser.add_argument(
      '--depot_tools_dir', help='Path to depot_tools in the host.',
      default=_DEPOT_TOOLS_DIR)
  args = parser.parse_args()
  logging.debug('Executing with flags: %s', args)
  return args


def main():
  flags = _ParseArgs()
  bot_count = _GetNumBotConfigs(socket.gethostname().split('.')[0])
  swarm_bots_home = flags.swarm_bots_home
  swarm_zip_path = os.path.join(
      swarm_bots_home, _SWARM_ZIP_RELATIVE_PATH_TO_BOT_HOME)
  if not os.path.exists(swarm_zip_path):
    raise ValueError('Swarm binary does not exists at %s' % swarm_zip_path)
  for bot_num in range(1, bot_count +1):
    SetupBot(
        bot_num=bot_num, user=flags.user, group=flags.group,
        swarm_bots_home=swarm_bots_home, swarm_bot_prefix=_SWARM_BOT_PREFIX,
        swarm_zip_path=swarm_zip_path, gsutil_dir=flags.gsutil_dir,
        gsutil_boto=flags.gsutil_boto, fragglerock_key=flags.fragglerock_key,
        eurtest_binaries=flags.eurtest_binaries, venv=flags.venv,
        devicemap_dir=flags.devicemap_dir,
        depot_tools_dir=flags.depot_tools_dir)


if __name__ == '__main__':
  main()