tools/process_starter.py

#!/usr/bin/env python
# Copyright 2017 The Fuchsia Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""A library with functions to start each of the Cobalt processes locally."""

import os
import shutil
import subprocess

THIS_DIR = os.path.dirname(__file__)
SRC_ROOT_DIR = os.path.abspath(os.path.join(THIS_DIR, os.pardir))
OUT_DIR = os.path.join(SRC_ROOT_DIR, 'out')
SYS_ROOT_DIR = os.path.join(SRC_ROOT_DIR, 'sysroot')

DEMO_CONFIG_DIR = os.path.abspath(os.path.join(SRC_ROOT_DIR, 'config',
    'demo'))
PRODUCTION_CONFIG_DIR = os.path.abspath(os.path.join(SRC_ROOT_DIR, 'config',
    'production'))
CONFIG_BINARY_PROTO = os.path.join(OUT_DIR, 'third_party', 'config',
    'cobalt_config.binproto')
SHUFFLER_DEMO_CONFIG_FILE = os.path.abspath(os.path.join(SRC_ROOT_DIR,
    'shuffler', 'src', 'config', 'config_demo.txt'))
SHUFFLER_DB_DIR = os.path.join("/tmp/cobalt_shuffler")

SHUFFLER_CONFIG_DIR = os.path.abspath(os.path.join(SRC_ROOT_DIR, 'shuffler',
    'src', 'shuffler_config'))
SHUFFLER_CONFIG_FILE = os.path.join(SHUFFLER_CONFIG_DIR, 'config_v0.txt')
SHUFFLER_DEMO_CONFIG_FILE = os.path.join(SHUFFLER_CONFIG_DIR, 'config_demo.txt')
SHUFFLER_TMP_DB_DIR = os.path.join("/tmp/cobalt_shuffler")

DEFAULT_SHUFFLER_PORT=5001
DEFAULT_ANALYZER_SERVICE_PORT=6001
DEFAULT_REPORT_MASTER_PORT=7001

DEFAULT_ANALYZER_PUBLIC_KEY_PEM=os.path.join(SRC_ROOT_DIR,
                                             "analyzer_public.pem")
ANALYZER_PRIVATE_KEY_PEM_NAME="analyzer_private.pem"
DEFAULT_ANALYZER_PRIVATE_KEY_PEM=os.path.join(SRC_ROOT_DIR,
                                             ANALYZER_PRIVATE_KEY_PEM_NAME)
DEFAULT_SHUFFLER_PUBLIC_KEY_PEM=os.path.join(SRC_ROOT_DIR,
                                             "shuffler_public.pem")
SHUFFLER_PRIVATE_KEY_PEM_NAME="shuffler_private.pem"
DEFAULT_SHUFFLER_PRIVATE_KEY_PEM=os.path.join(SRC_ROOT_DIR,
                                              SHUFFLER_PRIVATE_KEY_PEM_NAME)

LOCALHOST_TLS_CERT_FILE=os.path.join(SRC_ROOT_DIR, "end_to_end_tests",
                                   "localhost.crt")
LOCALHOST_TLS_KEY_FILE=os.path.join(SRC_ROOT_DIR, "end_to_end_tests",
                                  "localhost.key")


def kill_process(process, name):
  """ Kills the given process if it is running and waits for it to terminate.

  Args:
      process {Popen} A representation of the process to be killed.
          May be None.

      name {String} Name of the process for use in a user-facing message.
  """
  if process and process.poll() is None:
    print "Killing %s..." % name
    process.terminate()
    process.kill()
    process.wait()


def execute_command(cmd, wait):
  """ Executes the given command and optionally waits for it to complete.

  command {list of strings} will be passed to Popen().

  wait {bool} If true we will wait for the command to complete and return
      the result code. If false we will return immediately and return an
      instance of Popen.

  Returns:
    An instance of Popen if wait is false or an integer return code if
    wait is true.
  """
  p = subprocess.Popen(cmd)
  if not wait:
    return p
  return_code = p.wait()
  if return_code < 0:
    print
    print "****** WARNING Process [%s] terminated by signal %d" % (cmd[0], - return_code)
  return return_code

def start_bigtable_emulator(wait=True):
  # Note(rudominer) We can pass -port=n to cbtemulator to run on a different
  # port.
  print
  print "Starting the Cloud Bigtable Emulator..."
  print
  path = os.path.abspath(os.path.join(SYS_ROOT_DIR, 'gcloud',
      'google-cloud-sdk', 'platform', 'bigtable-emulator', 'cbtemulator'))
  cmd = [path]
  return execute_command(cmd, wait)

SHUFFLER_PATH = os.path.abspath(os.path.join(OUT_DIR, 'shuffler', 'shuffler'))
def start_shuffler(port=DEFAULT_SHUFFLER_PORT,
    analyzer_uri='localhost:%d' % DEFAULT_ANALYZER_SERVICE_PORT,
    use_memstore=False, erase_db=True, db_dir=SHUFFLER_TMP_DB_DIR,
    config_file=SHUFFLER_DEMO_CONFIG_FILE,
    private_key_pem_file=DEFAULT_SHUFFLER_PRIVATE_KEY_PEM,
    use_tls=False,
    tls_cert_file=LOCALHOST_TLS_CERT_FILE,
    tls_key_file=LOCALHOST_TLS_KEY_FILE,
    verbose_count=0, wait=True):
  """Starts the Shuffler.

  Args:
    port {int} The port on which the Shuffler should listen.
    analyzer_uri {string} The URI of the Analyzer Service
    use_memstore {bool} If false the Shuffler will use the LevelDB store
    emtpy_db {bool} When using the LevelDB store, should the store be
        erased before the shuffler starts?
    config_file {string} The path to the Shuffler's config file.
  """
  print
  cmd = [SHUFFLER_PATH,
        "-port", str(port),
        "-private_key_pem_file", private_key_pem_file,
        "-analyzer_uri", analyzer_uri,
        "-config_file", config_file,
        "-logtostderr"]
  if use_tls:
    cmd.append("-tls")
    cmd.append("-cert_file=%s"%tls_cert_file)
    cmd.append("-key_file=%s"%tls_key_file)
  if verbose_count > 0:
    cmd.append("-v=%d"%verbose_count)
  if use_memstore:
    cmd.append("-use_memstore")
  else:
    cmd = cmd + ["-db_dir", db_dir]
    if erase_db:
      print "Erasing Shuffler's LevelDB store at %s." % db_dir
      shutil.rmtree(db_dir, ignore_errors=True)

  print "Starting the shuffler..."
  print
  return execute_command(cmd, wait)

ANALYZER_SERVICE_PATH = os.path.abspath(os.path.join(OUT_DIR, 'analyzer',
    'analyzer_service', 'analyzer_service'))
def start_analyzer_service(port=DEFAULT_ANALYZER_SERVICE_PORT,
    bigtable_project_name='', bigtable_instance_id='',
    private_key_pem_file=DEFAULT_ANALYZER_PRIVATE_KEY_PEM,
    verbose_count=0, vmodule=None,
    wait=True):
  print
  print "Starting the analyzer service..."
  print
  cmd = [ANALYZER_SERVICE_PATH,
      "-port", str(port),
      "-private_key_pem_file", private_key_pem_file,
      "-logtostderr"]
  if bigtable_project_name != '' and bigtable_instance_id != '':
    cmd = cmd + [
      "-bigtable_project_name",
      bigtable_project_name,
      "-bigtable_instance_id",
      bigtable_instance_id,
    ]
  else:
    print "Will connect to a local Bigtable Emulator instance."
    cmd.append("-for_testing_only_use_bigtable_emulator")
  if verbose_count > 0:
    cmd.append("-v=%d"%verbose_count)
  if vmodule:
    cmd.append("-vmodule=%s"%vmodule)
  return execute_command(cmd, wait)

REPORT_MASTER_PATH = os.path.abspath(os.path.join(OUT_DIR, 'analyzer',
    'report_master', 'analyzer_report_master'))
def start_report_master(port=DEFAULT_REPORT_MASTER_PORT,
                        bigtable_project_name='', bigtable_instance_id='',
                        cobalt_registry_proto_path=CONFIG_BINARY_PROTO,
                        use_tls=False,
                        tls_cert_file=LOCALHOST_TLS_CERT_FILE,
                        tls_key_file=LOCALHOST_TLS_KEY_FILE,
                        verbose_count=0, vmodule=None,
                        wait=True):
  print
  print "Starting the analyzer ReportMaster service..."
  print
  cmd = [REPORT_MASTER_PATH,
      "-port", str(port),
      "-cobalt_registry_proto_path", cobalt_registry_proto_path,
      "-logtostderr"]
  if use_tls:
    cmd.append("-use_tls")
    cmd.append("-tls_cert_file=%s"%tls_cert_file)
    cmd.append("-tls_key_file=%s"%tls_key_file)
  if bigtable_project_name != '' and bigtable_instance_id != '':
    cmd = cmd + [
      "-bigtable_project_name",
      bigtable_project_name,
      "-bigtable_instance_id",
      bigtable_instance_id,
    ]
  else:
    print "Will connect to a local Bigtable Emulator instance."
    cmd.append("-for_testing_only_use_bigtable_emulator")
  if verbose_count > 0:
    cmd.append("-v=%d"%verbose_count)
  if vmodule:
    cmd.append("-vmodule=%s"%vmodule)
  return execute_command(cmd, wait)

TEST_APP_PATH = os.path.abspath(os.path.join(OUT_DIR, 'tools', 'test_app',
                                'cobalt_test_app'))
def start_test_app(shuffler_uri='', analyzer_uri='', use_tls=False,
                   root_certs_pem_file='',
                   analyzer_pk_pem_file=DEFAULT_ANALYZER_PUBLIC_KEY_PEM,
                   shuffler_pk_pem_file=DEFAULT_SHUFFLER_PUBLIC_KEY_PEM,
                   cobalt_registry_proto_path=CONFIG_BINARY_PROTO,
                   project_id=1,
                   automatic=False,
                   verbose_count=0, wait=True):
  cmd = [TEST_APP_PATH,
      "-shuffler_uri", shuffler_uri,
      "-analyzer_uri", analyzer_uri,
      "-analyzer_pk_pem_file", analyzer_pk_pem_file,
      "-shuffler_pk_pem_file", shuffler_pk_pem_file,
      "-config_bin_proto_path", cobalt_registry_proto_path,
      "-project", str(project_id),
      "-logtostderr"]
  if (use_tls):
    cmd.append("-use_tls")
    if (root_certs_pem_file):
      cmd.append("-root_certs_pem_file")
      cmd.append(root_certs_pem_file)
  if verbose_count > 0:
    cmd.append("-v=%d"%verbose_count)
  if automatic:
    cmd.append("-mode=automatic")
  return execute_command(cmd, wait)

def start_report_client(report_master_uri='',
                        use_tls=False,
                        root_certs_pem_file=LOCALHOST_TLS_CERT_FILE,
                        project_id=1,
                        verbose_count=0, wait=True):
  path = os.path.abspath(os.path.join(OUT_DIR, 'tools', 'report_client'))
  cmd = [path,
      "-report_master_uri", report_master_uri,
      "-project_id", str(project_id),
      "-logtostderr"]
  if (use_tls):
    cmd.append("-tls")
    if (root_certs_pem_file):
      cmd.append("-ca_file")
      cmd.append(root_certs_pem_file)
  if verbose_count > 0:
    cmd.append("-v=%d"%verbose_count)
  return execute_command(cmd, wait)

OBSERVATION_QUERIER_PATH = os.path.abspath(os.path.join(OUT_DIR, 'tools',
                                           'observation_querier',
                                           'query_observations'))
def start_observation_querier(bigtable_project_name='',
                              bigtable_instance_id='',
                              verbose_count=0):
  cmd = [OBSERVATION_QUERIER_PATH,
      "-logtostderr"]
  if not bigtable_project_name or not bigtable_instance_id:
    cmd.append("-for_testing_only_use_bigtable_emulator")
  else:
    cmd = cmd + ["-bigtable_project_name", bigtable_project_name,
                 "-bigtable_instance_id", bigtable_instance_id]
  if verbose_count > 0:
    cmd.append("-v=%d"%verbose_count)
  return execute_command(cmd, wait=True)