#!/usr/bin/env ruby
#
# Copyright (C) 2008-2009  Kouhei Sutou <kou@cozmixng.org>
#
# 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'
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)
      @output_directory = output_directory
      @update_time = update_time
      @data = []
      @items = nil
      @title = nil
      @vertical_label = nil
      @step = 60 # seconds
      @width = 600
      @points_per_sample = 3
    end

    def rows
      @width / @points_per_sample
    end

    def rrd_file
      build_path(self.class.base_rrd_file_name)
    end

    def collect_data(last_update_time)
      data = @data.collect {|datum| [datum.key, datum.time]}
      GraphData::Data.new(count_data({}, data, last_update_time))
    end

    def update
      if File.exist?(rrd_file)
        last_update_time = RRD.last_update_time(rrd_file)
      else
        last_update_time = nil
      end

      data = collect_data(last_update_time)
      return if data.empty?

      end_time = data.last_time
      if last_update_time
        start_time = last_update_time + @step
      else
        start_time = data.first_time
      end

      create_rrd(start_time, *@items) unless File.exist?(rrd_file)

      start_time.to_i.step(end_time, @step) do |time|
        counts = []
        @items.each do |item|
          counts << (data[item] || {})[time] || 0
        end

        # puts("#{Time.at(time).iso8601} #{counts.join(':')}: #{rrd_file}")
        RRD.update(rrd_file, "#{time}:#{counts.join(':')}")
      end
    end

    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 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 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

    def count_data(output, data, last_update_time)
      data.each do |key, time|
        time = normalize_time(time)

        # ignore sessions which has been already registerd
        # due to RRD.update fails on past time stamp
        next if last_update_time and time <= last_update_time

        # ignore recent sessions due to RRD.update fails on
        # past time stamp
        next if time >= normalize_time(@update_time)

        time = time.to_i

        output[key] ||= {}
        output[key][time] ||= 0
        output[key][time] += 1
      end
      output
    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"]
      @child_sessions = []
      @client_sessions = []
    end

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

    def collect_session(time_stamp, content, sessions, regex)
      if regex.match(content)
        elapsed = $1
        name = $2
        session = LogData::Session.new(name)
        session.end_time = time_stamp
        session.start_time = time_stamp - Float(elapsed)
        sessions << session
      end
      sessions
    end

    def feed(time_stamp, content)
      @child_sessions =
        collect_session(time_stamp, content, @child_sessions,
                        /\A\[milter\]\[end\]\[(.+)\]\(.+\): (.+)\z/)
      @client_sessions =
        collect_session(time_stamp, content, @client_sessions,
                        /\A\[session\]\[end\]\[(.+)\]\(.+\)\z/)
    end

    def collect_data(last_update_time)
      count = {}
      [["smtp", @client_sessions],
       ["child", @child_sessions]].each do |key, sessions|
        data = sessions.collect {|session| [key, session.end_time]}
        count_data(count, data, last_update_time)
      end
      GraphData::Data.new(count)
    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 feed(time_stamp, content)
      case content
      when /\A\[reply\]\[(.+)\]\[(.+)\]\z/
        state = $1
        status = $2
        return if status == "continue" and state != "end-of-message"
        status = "pass" if status == "continue"
        @data << LogData::Mail.new(status, time_stamp)
      when /\A\[abort\]\[(.+)\]\z/
        state = $1
        @data << LogData::Mail.new("abort", time_stamp)
      end
    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
  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 feed(time_stamp, content)
      if /\A\[stop\]\[(.+)\]: (.+)\z/ =~ content
        state = $1
        name = $2
        @data << LogData::Stop.new(state, name, time_stamp)
      end
    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
  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

  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]
    @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.update
    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

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

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