#!/usr/bin/env ruby
#
# Copyright (C) 2008-2009  Kouhei Sutou <kou@clear-code.com>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# 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, see <http://www.gnu.org/licenses/>.

require 'English'
require 'pathname'
require 'fileutils'
require 'time'
require 'optparse'
require 'erb'

begin
  require 'RRD' unless Object.const_defined?(:RRD)
rescue LoadError
  module RRD
    module_function
    def method_missing(method_name, *args)
      run(method_name.to_s, *args)
    end

    def last(file)
      Time.at(Integer(run("last", file)))
    end

    def run(command, *args)
      command_line = ["rrdtool", command] + args.collect {|arg| arg.to_s}

      stdin_pipe = IO.pipe
      stdout_pipe = IO.pipe
      stderr_pipe = IO.pipe

      pid = fork do
        stdin_pipe[1].close
        STDIN.reopen(stdin_pipe[0])
        stdin_pipe[0].close

        stdout_pipe[0].close
        STDOUT.reopen(stdout_pipe[1])
        stdout_pipe[1].close

        stderr_pipe[0].close
        STDERR.reopen(stderr_pipe[1])
        stderr_pipe[1].close

        exec(*command_line)
        exit!(1)
      end

      stdin_pipe[0].close
      stdout_pipe[1].close
      stderr_pipe[1].close
      _, status = Process.waitpid2(pid)

      stdout = stdout_pipe[0].read
      stderr = stderr_pipe[0].read
      [stdin_pipe[1], stdout_pipe[0], stderr_pipe[0]].each do |pipe|
        pipe.close unless pipe.closed?
      end

      unless status.success?
        raise "failed to run: #{command_line.join(' ')}: #{stderr}"
      end
      stdout
    end
  end
end

module RRD
  module_function
  def last_update_time(file)
    time = last(file)
    if time and time < Time.at(0)
      time = Time.at(Integer(`rrdtool last '#{file}'`))
    end
    time
  end
end

class MilterManagerLogAnalyzer
  module LogData
    class Session
      attr_accessor :name, :start_time, :end_time
      def initialize(name)
        @name = name
        @start_time = nil
        @end_time = nil
      end
    end

    class Mail
      attr_accessor :status, :time, :key
      def initialize(status, time)
        @key = status
        @status = status
        @time = time
      end
    end

    class Stop
      attr_accessor :state, :name, :time
      def initialize(state, name, time)
        @state = state
        @name = name
        @time = time
      end

      def key
        @state
      end
    end
  end

  module GraphData
    class TimeRange
      def initialize(step, rows)
        @step = step
        @rows = rows
      end

      def tag
        name.downcase
      end

      def steps
        base_range.steps * n_times
      end

      def start_time
        base_range.start_time * n_times
      end
    end

    class DayRange < TimeRange
      def name
        "Day"
      end

      def steps
        (3600 * 24) / (@step * @rows) + 1
      end

      def start_time
        -(60 * 60 * 24)
      end
    end

    class WeekRange < TimeRange
      def name
        "Week"
      end

      private
      def base_range
        DayRange.new(@step, @rows)
      end

      def n_times
        7
      end
    end

    class MonthRange < TimeRange
      def name
        "Month"
      end

      private
      def base_range
        WeekRange.new(@step, @rows)
      end

      def n_times
        5
      end
    end

    class YearRange < TimeRange
      def name
        "Year"
      end

      private
      def base_range
        MonthRange.new(@step, @rows)
      end

      def n_times
        12
      end
    end

    class Data
      def initialize(counts)
        @counts = counts
      end

      def empty?
        return true unless @counts
        @counts.each_value do |counting|
          return false unless counting.empty?
        end
        true
      end

      def last_time
        last_time = 0
        @counts.each_value do |count|
          next if count.empty?
          _last_time = count.keys.sort.last
          last_time = [last_time, _last_time].max
        end
        last_time || 0
      end

      def first_time
        first_time = nil
        @counts.each_value do |count|
          next if count.empty?
          _first_time = count.keys.sort.first
          first_time ||= _first_time
          first_time = [first_time, _first_time].min
        end
        first_time
      end

      def [](key)
        @counts[key]
      end

      def []=(key, value)
        @counts[key] = value
      end
    end
  end

  class GraphGenerator
    def initialize(output_directory, update_time=nil)
      @output_directory = output_directory
      @data = []
      @items = nil
      @title = nil
      @vertical_label = nil
      @step = 60 # seconds
      @width = 600
      @points_per_sample = 3

      @update_time = normalize_time(update_time)
    end

    def rows
      @width / @points_per_sample
    end

    def rrd_file
      build_path(self.class.base_rrd_file_name)
    end

    def prepare
      if File.exist?(rrd_file)
        @last_update_time = normalize_time(RRD.last_update_time(rrd_file))
        @current_time = @last_update_time.to_i + @step
      else
        @last_update_time = nil
        @current_time = nil
      end
      @data = {}
    end

    def feed(time_stamp, content)
      time_stamp = normalize_time(time_stamp)
      return if @last_update_time and time_stamp < @last_update_time
      return if time_stamp >= @update_time

      time_stamp = time_stamp.to_i
      key = find_key(content)
      return if key.nil?

      @current_time ||= time_stamp
      if time_stamp and @current_time > time_stamp
        @data = {}
      else
        flush(time_stamp)
        @data[key] ||= 0
        @data[key] += 1
      end
    end

    def flush(next_time_stamp=nil)
      return if @current_time == next_time_stamp
      create_rrd(@current_time, *@items) unless File.exist?(rrd_file)

      start_time = @current_time
      unless @data.empty?
        counts = []
        @items.each do |item|
          counts << @data[item]
        end
        RRD.update(rrd_file, "#{@current_time}:#{counts.join(':')}")
        start_time += @step
      end
      @data = {}

      if next_time_stamp
        start_time.step(next_time_stamp - @step, @step) do |time|
          counts = ":" * (@items.size - 1)
          RRD.update(rrd_file, "#{time}:#{counts}")
        end
      end
      @current_time = next_time_stamp || @current_time
    end

    def output_graph(range_class, options={}, *args)
      return nil unless File.exist?(rrd_file)
      last_update_time = RRD.last_update_time(rrd_file)

      range = range_class.new(@step, rows)
      start_time = options[:start_time] || range.start_time
      end_time = options[:end_time] || guess_end_time(start_time)
      graph_tag = options[:graph_tag] || range.tag
      width = options[:width] || @width
      height = options[:height] || 200

      vertical_label = "#{@vertical_label}/#{unit}"
      name = graph_name(graph_tag)

      items = @items.inject([]) do |_items, item|
        _items + generate_definitions(rrd_file, range, item)
      end

      data = items + args
      data << "COMMENT:[#{last_update_time.iso8601.gsub(/:/, '\:')}]\\r"
      RRD.graph(name,
                "--title", @title,
                "--vertical-label", vertical_label,
                "--start", start_time.to_s,
                "--end", end_time.to_s,
                "--width", width.to_s,
                "--height", height.to_s,
                "--alt-y-grid",
                "--units-exponent", "0",
                *data)
      name
    end

    private
    def create_rrd(start_time, *sources)
      ranges = [
                GraphData::DayRange.new(@step, rows),
                GraphData::WeekRange.new(@step, rows),
                GraphData::MonthRange.new(@step, rows),
                GraphData::YearRange.new(@step, rows),
               ]

      rra = []
      ranges.each do |range|
        rra << "RRA:MAX:0.5:#{range.steps}:#{rows}"
        rra << "RRA:AVERAGE:0.5:#{range.steps}:#{rows}"
      end

      data_sources = sources.collect do |source|
        "DS:#{source}:GAUGE:#{@step * 3}:0:U"
      end

      RRD.create(rrd_file,
                 "--start", (start_time - 1).to_i.to_s,
                 "--step", @step,
                 *(rra + data_sources))
    end

    def build_path(*paths)
      File.join(*([@output_directory, *paths].compact))
    end

    def guess_end_time(start_time)
      if start_time.is_a?(String)
        "now"
      else
        end_time = Time.now.to_i
        step = start_time.abs * @points_per_sample / @width
        end_time - (end_time % step)
      end
    end

    def normalize_time(time)
      Time.at(time.to_i - time.to_i % @step).utc
    end

    def generate_definitions(rrd_file, range, item, label=nil)
      label ||= item
      [
       "DEF:#{label}=#{rrd_file}:#{item}:AVERAGE",
       "DEF:max_#{label}=#{rrd_file}:#{item}:MAX",
       "CDEF:n_#{label}=#{label}",
       "CDEF:real_#{label}=#{label}",
       "CDEF:real_max_#{label}=max_#{label}",
       "CDEF:real_n_#{label}=n_#{label},#{range.steps},*",
       "CDEF:total_#{label}=PREV,UN,real_n_#{label},PREV,IF,real_n_#{label},+",
      ]
    end

    def unit
      case @step
      when 60
        "min"
      when 60 * 60
        "hour"
      when 60 * 60 * 24
        "day"
      when 60 * 60 * 24 * 365
        "year"
      else
        "unknown"
      end
    end
  end

  class SessionGraphGenerator < GraphGenerator
    class << self
      def base_rrd_file_name
        "milter-log.session.rrd"
      end
    end

    def initialize(output_directory, update_time)
      super(output_directory, update_time)
      @title = 'Sessions'
      @vertical_label = "sessions"
      @items = ["smtp", "child"]
    end

    def graph_name(tag)
      build_path("session.#{tag}.png")
    end

    def find_key(content)
      key = nil
      case content
      when /\A\[milter\]\[end\]\[(.+)\]\(.+\): (.+)\z/
        elapsed = $1
        name = $2
        key = "child"
      when /\A\[session\]\[end\]\[(.+)\]\(.+\)\z/
        elapsed = $1
        name = $2
        key = "smtp"
      end
      key
    end

    def output_graph(range_class, options={})
      super(range_class, options,
            "AREA:n_smtp#0000ff:SMTP  ",
            "GPRINT:total_smtp:MAX:total\\: %8.0lf sessions",
            "GPRINT:smtp:AVERAGE:avg\\: %6.2lf sessions/#{unit}",
            "GPRINT:max_smtp:MAX:max\\: %4.0lf sessions/#{unit}\\l",
            "LINE2:n_child#00ff00:milter",
            "GPRINT:total_child:MAX:total\\: %8.0lf sessions",
            "GPRINT:child:AVERAGE:avg\\: %6.2lf sessions/#{unit}",
            "GPRINT:max_child:MAX:max\\: %4.0lf sessions/#{unit}\\l")
    end
  end

  class MailStatusGraphGenerator < GraphGenerator
    class << self
      def base_rrd_file_name
        "milter-log.mail.rrd"
      end
    end

    def initialize(output_directory, update_time)
      super(output_directory, update_time)
      @vertical_label = "mails"
      @title = 'Processed mails'
      @items = ["pass",
                "accept",
                "reject",
                "discard",
                "temporary-failure",
                "quarantine",
                "abort"]
    end

    def graph_name(tag)
      build_path("mail.#{tag}.png")
    end

    def output_graph(range_class, options={})
      entries = [
                 ["AREA", "pass", "#0000ff", "Pass"],
                 ["STACK", "accept", "#00ff00", "Accept"],
                 ["STACK", "reject", "#ff0000", "Reject"],
                 ["STACK", "discard", "#ffd400", "Discard"],
                 ["STACK", "temporary-failure", "#888888", "Temp-Fail"],
                 ["STACK", "quarantine", "#a52a2a", "Quarantine"],
                 ["STACK", "abort", "#ff9999", "Abort"],
                ]
      max_label_size = entries.collect {|_, _, _, label| label.size}.max

      items = []
      entries.each do |type, name, color, label|
        items << "#{type}:#{name}#{color}:#{label.ljust(max_label_size)}"
        items << "GPRINT:total_#{name}:MAX:total\\: %11.0lf mails"
        items << "GPRINT:#{name}:AVERAGE:avg\\: %9.2lf mails/#{unit}"
        items << "GPRINT:max_#{name}:MAX:max\\: %7.0lf mails/#{unit}\\l"
      end

      range = range_class.new(@step, rows)
      path = build_path(SessionGraphGenerator.base_rrd_file_name)
      items += generate_definitions(path, range, "smtp")

      label = "SMTP".ljust(max_label_size)
      items << "LINE2:n_smtp#000000:#{label}"
      items << "GPRINT:total_smtp:MAX:total\\: %8.0lf sessions"
      items << "GPRINT:smtp:AVERAGE:avg\\: %6.2lf sessions/#{unit}"
      items << "GPRINT:max_smtp:MAX:max\\: %4.0lf sessions/#{unit}\\l"
      super(range_class, options, *items)
    end

    private
    def find_key(content)
      key = nil
      case content
      when /\A\[reply\]\[(.+)\]\[(.+)\]\z/
        state = $1
        status = $2
        return nil if status == "continue" and state != "end-of-message"
        status = "pass" if status == "continue"
        key = status
      when /\A\[abort\]\[(.+)\]\z/
        state = $1
        status = "abort"
        key = status
      end

      key
    end
  end

  class StopGraphGenerator < GraphGenerator
    class << self
      def base_rrd_file_name
        "milter-log.stop.rrd"
      end
    end

    def initialize(output_directory, update_time)
      super(output_directory, update_time)
      @title = 'Stopped milters'
      @vertical_label = "milters"
      @items = ["connect",
                "helo",
                "envelope-from",
                "envelope-recipient",
                "header",
                "body",
                "end-of-message"]
    end

    def graph_name(tag)
      build_path("stop.#{tag}.png")
    end

    def output_graph(range_class, options={})
      entries = [
                 ["AREA", "connect", "#0000ff"],
                 ["STACK", "helo", "#ff00ff"],
                 ["STACK", "envelope-from", "#00ffff"],
                 ["STACK", "envelope-recipient", "#ffff00"],
                 ["STACK", "header", "#a52a2a"],
                 ["STACK", "body", "#ff0000"],
                 ["STACK", "end-of-message", "#00ff00"],
                ]
      max_name_size = entries.collect {|_, name, _| name.size}.max

      items = []
      entries.each do |type, name, color|
        items << "#{type}:#{name}#{color}:#{name.ljust(max_name_size)}"
        items << "GPRINT:total_#{name}:MAX:total\\: %8.0lf milters"
        items << "GPRINT:#{name}:AVERAGE:avg\\: %6.2lf milters/#{unit}"
        items << "GPRINT:max_#{name}:MAX:max\\: %4.0lf milters/#{unit}\\l"
      end

      range = range_class.new(@step, rows)
      path = build_path(SessionGraphGenerator.base_rrd_file_name)
      items += generate_definitions(path, range, "child", "milter")

      label = 'total'.ljust(max_name_size)
      items << "LINE2:n_milter#000000:#{label}"
      items << "GPRINT:total_milter:MAX:total\\: %8.0lf milters"
      items << "GPRINT:milter:AVERAGE:avg\\: %6.2lf milters/#{unit}"
      items << "GPRINT:max_milter:MAX:max\\: %4.0lf milters/#{unit}\\l"
      super(range_class, options, *items)
    end

    private
    def find_key(content)
      key = nil
      case content
      when /\A\[stop\]\[(.+)\]: (.+)\z/
        state = $1
        name = $2
        key = state
      end
      key
    end
  end

  class HTMLGenerator
    include ERB::Util

    def initialize(graph_info)
      @graph_info = graph_info
    end

    def generate
      header + index + graphs + footer
    end

    private
    def header
      <<-EOH
<!DOCTYPE html PUBLIC "-//W3C//DTD XHTML 1.0 Transitional//EN"
      "http://www.w3.org/TR/xhtml1/DTD/xhtml1-transitional.dtd">
<html xmlns="http://www.w3.org/1999/xhtml">
  <head>
    <meta http-equiv="content-type" content="text/html;charset=UTF-8" />
    <meta http-equiv="Refresh" content="300" />
    <meta http-equiv="Pragma" content="no-cache" />
    <title>milter-manager statistics</title>
  </head>

  <body>
    <h1>milter-manager statistics</h1>
EOH
    end

    def index
      result = "<ul>\n"
      @graph_info.each do |label, graphs|
        next if graphs.empty?
        id = label_to_id(label)
        result << %Q[  <li><a href="\##{h(id)}">#{h(label)}</a></li>\n]
      end
      result << "</ul>\n"
      result
    end

    def graphs
      result = ""
      @graph_info.each do |label, graphs|
        next if graphs.empty?
        id = label_to_id(label)
        result << %Q[<h2 id=\"#{h(id)}\">#{h(label)}</h2>\n]
        graphs.each do |graph|
          result << %Q[  <p><img src="#{h(File.basename(graph))}" /></p>\n]
        end
        result << "\n"
      end
      result
    end

    def footer
      <<-EOF
  </body>
</html>
EOF
    end

    def label_to_id(label)
      label.downcase.gsub(/ +/, '-')
    end
  end

  attr_accessor :log, :output_directory
  def initialize
    @log = ARGF
    @update_db = true
    @sessions = nil
    @mail_status = nil
    @stops = nil
    @output_directory = "."
    @output_graphs = []
  end

  def parse_options(argv)
    opts = OptionParser.new do |opts|
      opts.on("--log=LOG_FILE",
              "The log file name in which is stored Milter log",
              "(STDIN)") do |log|
        @log = File.open(log)
      end

      opts.on("--output-directory=DIRECTORY",
              "Output graph, HTML and graph data to DIRECTORY",
              "(#{@output_directory})") do |directory|
        @output_directory = directory
        unless File.exist?(@output_directory)
          FileUtils.mkdir_p(@output_directory)
        end
      end

      opts.on("--[no-]update-db",
              "Update RRD database with log file",
              "(#{@update_db})") do |boolean|
        @update_db = boolean
      end
    end
    opts.parse!(argv)
  end

  def update
    return unless @update_db

    listeners = [sessions, mail_status, stops]
    listeners.each do |listener|
      listener.prepare
    end

    @log.each_line do |line|
      case line
      when /\A(\w{3} +\d+ \d+:\d+:\d+) [\w\-]+ ([\w\-]+)\[\d+\]: /
        time_stamp = $1
        name = $2
        content = $POSTMATCH
        next unless name == "milter-manager"
        if /\A\[statistics\] / =~ content
          content = $POSTMATCH.chomp
        else
          next
        end
        time_stamp = Time.parse(time_stamp).utc
        listeners.each do |listener|
          listener.feed(time_stamp, content)
        end
      else
      end
    end

    listeners.each do |listener|
      listener.flush
    end
  end

  def output_graph(range, label, options={})
    output_graphs = []
    [sessions, mail_status, stops].each do |generator|
      file_name = generator.output_graph(range, options)
      output_graphs << file_name if file_name
    end
    @output_graphs << [label, output_graphs]
  end

  def output_all_graph
    output_graph(GraphData::DayRange, "Last Day")
    output_graph(GraphData::WeekRange, "Last Week")
    output_graph(GraphData::MonthRange, "Last Month")
    output_graph(GraphData::YearRange, "Last Year")
  end

  def output_html
    File.open(File.join(@output_directory, "index.html"), "w") do |html|
      generator = HTMLGenerator.new(@output_graphs)
      html.print(generator.generate)
    end
  end

  private
  def now
    @now ||= Time.now.utc
  end

  def sessions
    @sessions ||= SessionGraphGenerator.new(@output_directory, now)
  end

  def mail_status
    @mail_status ||= MailStatusGraphGenerator.new(@output_directory, now)
  end

  def stops
    @stops ||= StopGraphGenerator.new(@output_directory, now)
  end
end

if __FILE__ == $0
  milter_log_tool = MilterManagerLogAnalyzer.new
  milter_log_tool.parse_options(ARGV)
  milter_log_tool.update
  milter_log_tool.output_all_graph
  milter_log_tool.output_html
end

# vi:ts=2:nowrap:ai:expandtab:sw=2
