#!/usr/bin/env ruby
#
# Copyright (C) 2008  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 'pathname'
require 'time'
require 'optparse'
require 'net/smtp'
require 'digest/sha2'

class MilterPerformanceTool
  def initialize
    @smtp_server = "localhost"
    @smtp_port = 25
    @helo_fqdn = "localhost.localdomain"
    @from = "from@example.com"
    @recipients = []
    @default_recipients = ["to@example.com"]
    @n_mails = 100
  end

  def parse_options(argv)
    opts = OptionParser.new do |opts|
      opts.separator ""
      opts.separator "Help options:"

      opts.on("-h", "--help", "Show this message") do
        puts opts
        exit(0)
      end

      opts.separator ""
      opts.separator "Application options:"

      opts.on("--smtp-server=SERVER",
              "Use SERVER as SMTP server",
              "(#{@smtp_server})") do |smtp_server|
        @smtp_server = smtp_server
      end

      opts.on("--smtp-port=PORT", Integer,
              "Use PORT as SMTP port",
              "(#{@smtp_port})") do |smtp_port|
        @smtp_port = smtp_port
      end

      opts.on("--helo-fqdn=FQDN",
              "Use FQDN for SMTP HELO command",
              "(#{@helo_fqdn})") do |helo_fqdn|
        @helo_fqdn = helo_fqdn
      end

      opts.on("--from=FROM",
              "Use FROM as envelope from address on SMTP MAIL command",
              "(#{@from})") do |from|
        @from = from
      end

      opts.on("--recipient=RECIPIENT",
              "Use RECIPIENT as envelope recipient address on SMTP RCPT command",
              "This option can be used n-times to set multi recipients.",
              "(#{@default_recipients.inspect})") do |recipient|
        @recipients << recipient
      end

      opts.on("--n-mails=N", Integer,
              "Send a test mail N times",
              "(#{@n_mails})") do |n_mails|
        @n_mails = n_mails
      end
    end
    opts.parse!(argv)
    @recipients = @default_recipients if @recipients.empty?
  end

  def run
    Thread.abort_on_exception = true
    @elapsed_time = 0
    threads = []
    @n_mails.times do
      thread = Thread.start do
        @elapsed_time += start_smtp
      end
      threads.push << thread
    end
    threads.each do |thread|
      thread.join
    end
  end

  def report
    puts "Processed #{@n_mails} in #{@elapsed_time} sec."
    puts "Average #{@n_mails / @elapsed_time} mails/sec."
  end

  private
  def start_smtp
    mail_source = generate_mail_source
    start_time = Time.now
    begin
      Net::SMTP.start(@smtp_server, @port, @helo_fqdn) do |smtp|
        smtp.send_mail(mail_source, @from, *@recipients)
      end
    rescue Net::ProtoFatalError
      puts "SMTP error: #{$!}"
    end
    Time.now - start_time
  end

  def generate_mail_source
    <<-EOM
MIME-Version: 1.0
Content-Transfer-Encoding: 7bit
Content-Type: text/plain; charset=US-ASCII
X-Mailer: milter-performance-check
Message-Id: <#{random_tag}@mail.example.com>
Subject: test mail
From: #{@from}
To: #{@recipients.join(', ')}
Date: Tue, 30 Dec 2008 22:26:06 +0900

Hello,

This is a test mail.
EOM
  end

  def random_tag
    Digest::SHA2.hexdigest("#{Time.now.to_i.to_s}.#{rand(Time.now.to_i)}")
  end
end

performance_tool = MilterPerformanceTool.new
performance_tool.parse_options(ARGV)
performance_tool.run
performance_tool.report

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