catatester/scripts/generate_monitoring_configs.py

"""Generates collectd and fluentd config from template."""
import argparse
import logging
import os
import socket
import subprocess

import jinja2


_DEFAULT_PROJECT_ID = 'google.com:fragglerock'
_DEFAULT_ZONE_NAME = 'us-central1-c'
_UNINITIALIZED = 'uninitialized'


def _GetHostName():
  """Returns local machine's host name."""
  return socket.gethostname().split('.')[0]


def _GetFqdn():
  """Returns local machine's FQDN."""
  return socket.getfqdn()


def _GetIpAddress():
  """Returns IP address of the given interface."""
  output = subprocess.check_output(['hostname', '-i'])
  return output.strip()


def ParseCmdOptions():
  """Parses command inputs."""
  parser = argparse.ArgumentParser(
      description='Generates collectd and fluentd config.')
  parser.add_argument('--hostname', help='Host name')
  parser.add_argument('--fqdn', help='Host FQDN')
  parser.add_argument('--ip', help='IPv4 address')
  parser.add_argument(
      '--project_id', help='Project ID, e.g. google.com:fragglerock')
  parser.add_argument(
      '--zone_name', help='Instance zone name, e.g. asia-east1-a')
  parser.add_argument(
      '--template_dir', help='Directory containing the templates.')
  parser.add_argument(
      '--debug',
      action='store_true',
      default=False,
      help='Print extra debug information.')

  args = parser.parse_args()
  if args.debug:
    logging.basicConfig(level=logging.DEBUG)
  else:
    logging.basicConfig(level=logging.INFO)

  if not args.hostname or args.hostname == _UNINITIALIZED:
    args.hostname = _GetHostName()
  if not args.fqdn or args.fqdn == _UNINITIALIZED:
    args.fqdn = _GetFqdn()
  if not args.ip or args.ip == _UNINITIALIZED:
    args.ip = _GetIpAddress()
  if not args.project_id or args.project_id == _UNINITIALIZED:
    args.project_id = _DEFAULT_PROJECT_ID
  if not args.zone_name or args.zone_name == _UNINITIALIZED:
    args.zone_name = _DEFAULT_ZONE_NAME
  if not args.template_dir or not os.path.isdir(args.template_dir):
    args.template_dir = os.path.dirname(__file__)

  return args


def _GenerateCollectdConfig(flags):
  """Outputs collectd.conf.

  Args:
    flags: argparse.Namespace instance containing config information.
  """
  template_path = os.path.join(flags.template_dir, 'collectd.conf.template')
  output_path = os.path.join(flags.template_dir, 'collectd.conf')
  with open(template_path) as fn:
    tm = jinja2.Template(fn.read())
    # Values demand double-quotes
    conf = tm.render(
        hostname='\"' + flags.hostname + '\"',
        fqdn='\"' + flags.fqdn + '\"',
        project_name='\"' + flags.project_id + '\"',
        zone_name='\"' + flags.zone_name + '\"')
    with open(output_path, 'w') as output:
      output.write(conf)


def _GenerateFluentdConfig(flags):
  """Outputs google-fluentd.conf.

  Args:
    flags: argparse.Namespace instance containing config information.
  """
  template_path = os.path.join(flags.template_dir,
                               'google-fluentd.conf.template')
  output_path = os.path.join(flags.template_dir,
                             'google-fluentd.conf')
  with open(template_path) as fn:
    tm = jinja2.Template(fn.read())
    # Values demand double-quotes
    conf = tm.render(
        hostname=flags.hostname,
        fqdn=flags.fqdn,
        project_id=flags.project_id,
        zone_name=flags.zone_name)
    with open(output_path, 'w') as output:
      output.write(conf)


def main():
  """Generates collectd and fluentd config from template."""
  flags = ParseCmdOptions()
  logging.debug('hostname: %s', flags.hostname)
  logging.debug('fqdn: %s', flags.fqdn)
  logging.debug('project: %s', flags.project_id)
  logging.debug('zone: %s', flags.zone_name)
  _GenerateCollectdConfig(flags)
  _GenerateFluentdConfig(flags)


if __name__ == '__main__':
  main()