#!/usr/bin/python
# -*- coding: utf-8 -*-
#---------------------------------------------------------------------
#
# Copyright © 2011 Canonical Ltd.
#
# Author: James Hunt <james.hunt@canonical.com>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2, as
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License along
# with this program; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
#---------------------------------------------------------------------

#---------------------------------------------------------------------
# Script to take output of "initctl show-config -e" and convert it into
# a Graphviz DOT language (".dot") file for procesing with dot(1), etc.
#
# Notes:
#
# - Slightly laborious logic used to satisfy graphviz requirement that
#   all nodes be defined before being referenced.
#
# Usage:
#
#   initctl show-config -e > initctl.out
#   initctl2dot -f initctl.out -o upstart.dot
#   dot -Tpng -o upstart.png upstart.dot
#
# Or more simply:
#
#  initctl2dot -o - | dot -Tpng -o upstart.png
#
# See also:
#
# - dot(1).
# - initctl(8).
# - http://www.graphviz.org.
#---------------------------------------------------------------------

import sys
import re
import fnmatch
import os
from string import split
import datetime
from subprocess import (Popen, PIPE)
from optparse import OptionParser

jobs   = {}
events = {}
cmd = "initctl --system show-config -e"
script_name =  os.path.basename(sys.argv[0])

job_events = [ 'starting', 'started', 'stopping', 'stopped' ]

# list of jobs to restict output to
restrictions_list = []

default_color_emits    = 'green'
default_color_start_on = 'blue'
default_color_stop_on  = 'red'
default_color_event    = 'thistle'
default_color_job      = '#DCDCDC' # "Gainsboro"
default_color_text     = 'black'
default_color_bg       = 'white'

default_outfile        = 'upstart.dot'


def header(ofh):
  global options

  str  = "digraph upstart {\n"

  # make the default node an event to simplify glob code
  str += "  node [shape=\"diamond\", fontcolor=\"%s\", fillcolor=\"%s\", style=\"filled\"];\n" \
    % (options.color_event_text, options.color_event)
  str += "  rankdir=LR;\n"
  str += "  overlap=false;\n"
  str += "  bgcolor=\"%s\";\n" % options.color_bg
  str += "  fontcolor=\"%s\";\n" % options.color_text

  ofh.write(str)


def footer(ofh):
  global options

  epilog = "overlap=false;\n"
  epilog += "label=\"Generated on %s by %s\\n" % \
    (str(datetime.datetime.now()), script_name)

  if options.restrictions:
    epilog += "(subset, "
  else:
    epilog += "("

  if options.infile:
    epilog += "from file data).\\n"
  else:
    epilog += "from '%s' on host %s).\\n" % \
      (cmd, os.uname()[1])

  epilog += "Boxes of color %s denote jobs.\\n" % options.color_job
  epilog += "Solid diamonds of color %s denote events.\\n" % options.color_event
  epilog += "Dotted diamonds denote 'glob' events.\\n"
  epilog += "Emits denoted by %s lines.\\n" % options.color_emits
  epilog += "Start on denoted by %s lines.\\n" % options.color_start_on
  epilog += "Stop on denoted by %s lines.\\n" % options.color_stop_on
  epilog += "\";\n"
  epilog += "}\n"
  ofh.write(epilog)


# Map dash to underscore since graphviz node names cannot
# contain dashes. Also remove dollars and colons
def sanitise(s):
  return s.replace('-', '_').replace('$', 'dollar_').replace('[', \
  'lbracket').replace(']', 'rbracket').replace('!', \
  'bang').replace(':', '_').replace('*', 'star').replace('?', 'question')


# Convert a dollar in @name to a unique-ish new name, based on @job and
# return it. Used for very rudimentary instance handling.
def encode_dollar(job, name):
  if name[0] == '$':
    name = job + ':' + name
  return name


def mk_node_name(name):
  return sanitise(name)


# Jobs and events can have identical names, so prefix them to namespace
# them off.
def mk_job_node_name(name):
  return mk_node_name('job_' + name)


def mk_event_node_name(name):
  return mk_node_name('event_' + name)


def show_event(ofh, name):
    global options
    str = "%s [label=\"%s\", shape=diamond, fontcolor=\"%s\", fillcolor=\"%s\"," % \
      (mk_event_node_name(name), name, options.color_event_text, options.color_event)

    if '*' in name:
      str += " style=\"dotted\""
    else:
      str += " style=\"filled\""

    str += "];\n"

    ofh.write(str)

def show_events(ofh):
  global events
  global options
  global restrictions_list

  events_to_show = []

  if restrictions_list:
    for job in restrictions_list:

      # We want all events emitted by the jobs in the restrictions_list.
      events_to_show += jobs[job]['emits']

      # We also want all events that jobs in restrictions_list start/stop
      # on.
      events_to_show += jobs[job]['start on']['event']
      events_to_show += jobs[job]['stop on']['event']

      # We also want all events emitted by all jobs that jobs in the
      # restrictions_list start/stop on. Finally, we want all events
      # emmitted by those jobs in the restrictions_list that we
      # start/stop on.
      for j in jobs[job]['start on']['job']:
        if jobs.has_key(j) and jobs[j].has_key('emits'):
          events_to_show += jobs[j]['emits']

      for j in jobs[job]['stop on']['job']:
        if jobs.has_key(j) and jobs[j].has_key('emits'):
          events_to_show += jobs[j]['emits']
  else:
    events_to_show = events

  for e in events_to_show:
    show_event(ofh, e)


def show_job(ofh, name):
  global options

  ofh.write("""
    %s [shape=\"record\", label=\"<job> %s | { <start> start on | <stop> stop on }\", fontcolor=\"%s\", style=\"filled\", fillcolor=\"%s\"];
    """ % (mk_job_node_name(name), name, options.color_job_text, options.color_job))


def show_jobs(ofh):
  global jobs
  global options
  global restrictions_list

  if restrictions_list:
    jobs_to_show = restrictions_list
  else:
    jobs_to_show = jobs

  for j in jobs_to_show:
    show_job(ofh, j)
    # add those jobs which are referenced by existing jobs, but which
    # might not be available as .conf files. For example, plymouth.conf
    # references gdm *or* kdm, but you are unlikely to have both
    # installed.
    for s in jobs[j]['start on']['job']:
      if s not in jobs_to_show:
        show_job(ofh, s)

    for s in jobs[j]['stop on']['job']:
      if s not in jobs_to_show:
        show_job(ofh, s)

  if not restrictions_list:
    return

  # Having displayed the jobs in restrictions_list,
  # we now need to display all jobs that *those* jobs
  # start on/stop on.
  for j in restrictions_list:
    for job in jobs[j]['start on']['job']:
      show_job(ofh, job)
    for job in jobs[j]['stop on']['job']:
      show_job(ofh, job)

  # Finally, show all jobs which emit events that jobs in the
  # restrictions_list care about.
  for j in restrictions_list:

    for e in jobs[j]['start on']['event']:
      for k in jobs:
        if e in jobs[k]['emits']:
          show_job(ofh, k)

    for e in jobs[j]['stop on']['event']:
      for k in jobs:
        if e in jobs[k]['emits']:
          show_job(ofh, k)


def show_edge(ofh, from_node, to_node, color):
  ofh.write("%s -> %s [color=\"%s\"];\n" % (from_node, to_node, color))


def show_start_on_job_edge(ofh, from_job, to_job):
  global options
  show_edge(ofh, "%s:start" % mk_job_node_name(from_job),
    "%s:job" % mk_job_node_name(to_job), options.color_start_on)


def show_start_on_event_edge(ofh, from_job, to_event):
  global options
  show_edge(ofh, "%s:start" % mk_job_node_name(from_job),
    mk_event_node_name(to_event), options.color_start_on)


def show_stop_on_job_edge(ofh, from_job, to_job):
  global options
  show_edge(ofh, "%s:stop" % mk_job_node_name(from_job),
    "%s:job" % mk_job_node_name(to_job), options.color_stop_on)


def show_stop_on_event_edge(ofh, from_job, to_event):
  global options
  show_edge(ofh, "%s:stop" % mk_job_node_name(from_job),
    mk_event_node_name(to_event), options.color_stop_on)


def show_job_emits_edge(ofh, from_job, to_event):
  global options
  show_edge(ofh, "%s:job" % mk_job_node_name(from_job),
    mk_event_node_name(to_event), options.color_emits)


def show_edges(ofh):
  global events
  global jobs
  global options
  global restrictions_list

  glob_jobs = {}

  if restrictions_list:
    jobs_list = restrictions_list
  else:
    jobs_list = jobs

  for job in jobs_list:

    for s in jobs[job]['start on']['job']:
      show_start_on_job_edge(ofh, job, s)

    for s in jobs[job]['start on']['event']:
      show_start_on_event_edge(ofh, job, s)

    for s in jobs[job]['stop on']['job']:
      show_stop_on_job_edge(ofh, job, s)

    for s in jobs[job]['stop on']['event']:
      show_stop_on_event_edge(ofh, job, s)

    for e in jobs[job]['emits']:
      if '*' in e:
        # handle glob patterns in 'emits'
        glob_events = []
        for _e in events:
          if e != _e and fnmatch.fnmatch(_e, e):
            glob_events.append(_e)
        glob_jobs[job] = glob_events

      show_job_emits_edge(ofh, job, e)

    if not restrictions_list:
      continue

    # Add links to events emitted by all jobs which current job
    # start/stops on
    for j in jobs[job]['start on']['job']:
      if not jobs.has_key(j):
        continue
      for e in jobs[j]['emits']:
        show_job_emits_edge(ofh, j, e)

    for j in jobs[job]['stop on']['job']:
      for e in jobs[j]['emits']:
        show_job_emits_edge(ofh, j, e)

  # Create links from jobs (which advertise they emits a class of
  # events, via the glob syntax) to all the events they create.
  for g in glob_jobs:
    for ge in glob_jobs[g]:
      show_job_emits_edge(ofh, g, ge)

  if not restrictions_list:
    return

  # Add jobs->event links to jobs which emit events that current job
  # start/stops on.
  for j in restrictions_list:

    for e in jobs[j]['start on']['event']:
      for k in jobs:
        if e in jobs[k]['emits'] and e not in restrictions_list:
          show_job_emits_edge(ofh, k, e)

    for e in jobs[j]['stop on']['event']:
      for k in jobs:
        if e in jobs[k]['emits'] and e not in restrictions_list:
          show_job_emits_edge(ofh, k, e)


def read_data():
  global jobs
  global events
  global options
  global cmd
  global job_events

  if options.infile:
    try:
      ifh = open(options.infile, 'r')
    except:
      sys.exit("ERROR: cannot read file '%s'" % options.infile)
  else:
    try:
      ifh = Popen(split(cmd), stdout=PIPE).stdout
    except:
      sys.exit("ERROR: cannot run '%s'" % cmd)

  for line in ifh.readlines():
      record = {}
      line = line.rstrip()

      result = re.match('^\s+start on ([^,]+) \(job:\s*([^,]*), env:', line)
      if result:
        _event = encode_dollar(job, result.group(1))
        _job   = result.group(2)
        if _job:
          jobs[job]['start on']['job'][_job] = 1
        else:
          jobs[job]['start on']['event'][_event] = 1
          events[_event] = 1
        continue

      result = re.match('^\s+stop on ([^,]+) \(job:\s*([^,]*), env:', line)
      if result:
        _event = encode_dollar(job, result.group(1))
        _job   = result.group(2)
        if _job:
          jobs[job]['stop on']['job'][_job] = 1
        else:
          jobs[job]['stop on']['event'][_event] = 1
          events[_event] = 1
        continue

      if re.match('^\s+emits', line):
        event = (line.lstrip().split())[1]
        event = encode_dollar(job, event)
        events[event] = 1
        jobs[job]['emits'][event] = 1
      else:
        tokens = (line.lstrip().split())

        if len(tokens) != 1:
          sys.exit("ERROR: invalid line: %s" % line.lstrip())

        job_record      = {}

        start_on        = {}
        start_on_jobs   = {}
        start_on_events = {}

        stop_on         = {}
        stop_on_jobs    = {}
        stop_on_events  = {}

        emits           = {}

        start_on['job']    = start_on_jobs
        start_on['event']  = start_on_events

        stop_on['job']     = stop_on_jobs
        stop_on['event']   = stop_on_events

        job_record['start on'] = start_on
        job_record['stop on']  = stop_on
        job_record['emits']    = emits

        job = (tokens)[0]
        jobs[job] = job_record


def main():
  global jobs
  global options
  global cmd
  global default_color_emits
  global default_color_start_on
  global default_color_stop_on
  global default_color_event
  global default_color_job
  global default_color_text
  global default_color_bg
  global restrictions_list

  description = "Convert initctl(8) output to GraphViz dot(1) format."
  epilog = \
    "See http://www.graphviz.org/doc/info/colors.html for available colours."

  parser = OptionParser(description=description, epilog=epilog)

  parser.add_option("-r", "--restrict-to-jobs",
      dest="restrictions",
      help="Limit display of 'start on' and 'stop on' conditions to " +
      "specified jobs (comma-separated list).")

  parser.add_option("-f", "--infile",
      dest="infile",
      help="File to read '%s' output from. If not specified, " \
      "initctl will be run automatically." % cmd)

  parser.add_option("-o", "--outfile",
      dest="outfile",
      help="File to write output to (default=%s)" % default_outfile)

  parser.add_option("--color-emits",
      dest="color_emits",
      help="Specify color for 'emits' lines (default=%s)." %
      default_color_emits)

  parser.add_option("--color-start-on",
      dest="color_start_on",
      help="Specify color for 'start on' lines (default=%s)." %
      default_color_start_on)

  parser.add_option("--color-stop-on",
      dest="color_stop_on",
      help="Specify color for 'stop on' lines (default=%s)." %
      default_color_stop_on)

  parser.add_option("--color-event",
      dest="color_event",
      help="Specify color for event boxes (default=%s)." %
      default_color_event)

  parser.add_option("--color-text",
      dest="color_text",
      help="Specify color for summary text (default=%s)." %
      default_color_text)

  parser.add_option("--color-bg",
      dest="color_bg",
      help="Specify background color for diagram (default=%s)." %
      default_color_bg)

  parser.add_option("--color-event-text",
      dest="color_event_text",
      help="Specify color for text in event boxes (default=%s)." %
      default_color_text)

  parser.add_option("--color-job-text",
      dest="color_job_text",
      help="Specify color for text in job boxes (default=%s)." %
      default_color_text)

  parser.add_option("--color-job",
      dest="color_job",
      help="Specify color for job boxes (default=%s)." %
      default_color_job)

  parser.set_defaults(color_emits=default_color_emits,
  color_start_on=default_color_start_on,
  color_stop_on=default_color_stop_on,
  color_event=default_color_event,
  color_job=default_color_job,
  color_job_text=default_color_text,
  color_event_text=default_color_text,
  color_text=default_color_text,
  color_bg=default_color_bg,
  outfile=default_outfile)

  (options, args) = parser.parse_args()

  if options.outfile == '-':
    ofh = sys.stdout
  else:
    try:
      ofh = open(options.outfile, "w")
    except:
      sys.exit("ERROR: cannot open file %s for writing" % options.outfile)

  if options.restrictions:
    restrictions_list = options.restrictions.split(",")

  read_data()

  for job in restrictions_list:
    if not job in jobs:
      sys.exit("ERROR: unknown job %s" % job)

  header(ofh)
  show_events(ofh)
  show_jobs(ofh)
  show_edges(ofh)
  footer(ofh)


if __name__ == "__main__":
  main()