#!/usr/bin/env ruby
#
# rbmc - ruby mailcounter
#
# author: mathias gumz < gumz at cs.uni-magdeburg.de >
# about : rbmc counts given mail, i use it for generating charts of
#         incoming spam.
#
# date  : 030619 06:46:35
#

require 'getoptlong'
require 'rmail'

@mails = Hash.new
@setup = {
  'ctype'       => 'date',
  'output'      => $stdout,
  'otype'       => 'csv',
  'input'       => $stdin,
  'mail'        => 'none',
  'recursive'   => false,
  'format'      => '%Y-%m-%d',
  'seperator'   => "\t",
  'gdc-style'   => "BAR3D",
  'gdc-itype'   => "PNG",
  'gdc-title'   => 'Mail-Count',
  'gdc-ytitle'  => 'Mails',
  'gdc-xtitle'  => 'Date',
  'gdc-width'   => 600,
  'gdc-height'  => 480,
  'gdc-bgcolor' => 0xffffff,
  'gdc-d1color' => 0xff9933,
  'gdc-d2color' => 0x66ff66,
  'gdc-type'    => 'single',
  'gdc-d2scale' => 0.00001
}

begin
  require 'GDChart'
  @gdc_styletags = {

# Chart types,
    "LINE" => GDChart::LINE,
    "AREA" => GDChart::AREA,
    "BAR" => GDChart::BAR,
    "FLOATINGBAR" => GDChart::FLOATINGBAR,
    "HILOCLOSE" => GDChart::HILOCLOSE,
    "COMBO_LINE_BAR" => GDChart::COMBO_LINE_BAR,
    "COMBO_HLC_BAR" => GDChart::COMBO_HLC_BAR,
    "COMBO_LINE_AREA" => GDChart::COMBO_LINE_AREA,
    "COMBO_LINE_LINE" => GDChart::COMBO_LINE_LINE,
    "COMBO_HLC_AREA" => GDChart::COMBO_HLC_AREA,
    "HILOCLOSE3D" => GDChart::HILOCLOSE3D,
    "COMBO_LINE_BAR3D" => GDChart::COMBO_LINE_BAR3D,
    "COMBO_LINE_AREA3D" => GDChart::COMBO_LINE_AREA3D,
    "COMBO_LINE_LINE3D" => GDChart::COMBO_LINE_LINE3D,
    "COMBO_HLC_BAR3D" => GDChart::COMBO_HLC_BAR3D,
    "COMBO_HLC_AREA3D" => GDChart::COMBO_HLC_AREA3D,
    "BAR3D" => GDChart::BAR3D,
    "FLOATINGBAR3D" => GDChart::FLOATINGBAR3D,
    "AREA3D" => GDChart::AREA3D,
    "LINE3D" => GDChart::LINE3D
  }

  @gdc_stacktags = {
# Stack types,
    "STACK_DEPTH" => GDChart::STACK_DEPTH,
    "STACK_SUM" => GDChart::STACK_SUM,
    "STACK_BESIDE" => GDChart::STACK_BESIDE,
    "STACK_LAYER" => GDChart::STACK_LAYER
  }

  @gdc_hltags = {
# HiLoClose chart styles,
    "HLC_DIAMOND" => GDChart::HLC_DIAMOND,
    "HLC_CLOSE_CONNECTED" => GDChart::HLC_CLOSE_CONNECTED,
    "HLC_CONNECTING" => GDChart::HLC_CONNECTING,
    "HLC_I_CAP" => GDChart::HLC_I_CAP
  }

  @gdc_ticktags = {
# Axis tick types,
    "TICK_LABELS" => GDChart::TICK_LABELS,
    "TICK_POINTS" => GDChart::TICK_POINTS,
    "TICK_NONE	" => GDChart::TICK_NONE
  }

  @gdc_bordertags =  {
# Border types,
    "BORDER_NONE" => GDChart::BORDER_NONE,
    "BORDER_ALL" => GDChart::BORDER_ALL,
    "BORDER_X" => GDChart::BORDER_X,
    "BORDER_Y" => GDChart::BORDER_Y,
    "BORDER_Y2" => GDChart::BORDER_Y2,
    "BORDER_TOP" => GDChart::BORDER_TOP
  }

  @gdc_itypetags = {
# Output types,
    "PNG" => GDChart::PNG,
    "JPEG" => GDChart::JPEG,
    "GIF" => GDChart::GIF
  }

  @gdc_fonttags = {
# Font sizes,
    "TINY" => GDChart::TINY,
    "SMALL" => GDChart::SMALL,
    "MEDBOLD" => GDChart::MEDBOLD,
    "LARGE" => GDChart::LARGE,
    "GIANT" => GDChart::GIANT
  }
rescue

end

INFO = <<END
 -= rbmc - version 0.1 =-

 written by mathias gumz <gumz at cs.uni-magdeburg.de>

END

HEADER = <<END
#
# settings for gnuplot
#
#  set xdata time
#  set format x "%m/%Y"
#  set terminal png small color
#  set output "out.png"
END

MINIUSAGE = <<END

NAME

  rbmc - ruby mailcounter

SYNOPSIS

  rbmc [options] mbox/maildir

END

USAGE = <<END

OPTIONS

 --help         -h       help
 --version      -v       version

 --count-type   -c       what should be counted:
                        'date' : default

 --maildir      -d       directory to mail
 --recursive    -r       search recursive in maildirs
 --mbox         -m       mbox

 --output       -o       file, default is stdout
 --output-type  -t       one of the following:
                         'csv' : comma separated values, default
                         'gdchart' chart-picture, gif-output

SPECIALOPTIONS

 options for -c 'date'

 --years        -Y       count mails per year
 --months       -M       count mails per month
 --days         -D       count mails per day
 --dayofweek    -W       count mails per day of the week
 --hours        -H       count mails per hour of day

 options for csv-output

 --separator    -s       separator for output file, default \\t

 options for gdchart-output

 --gdc-title             title of the chart
 --gdc-xtitle            title of x-axis
 --gdc-ytitle            title of y-axis
 --gdc-titlesize         size of the title, help prints types
 --gdc-style             style of the chart
 --gdc-itype             image type
 --gdc-width             width of the chart
 --gdc-height            height of the chart
 --gdc-bgcolor           background color
 --gdc-d1color           color for data1
 --gdc-d2color           color for data2
 --gdc-type              display more than just the count
                         'single' -> no combo, default,
                         'combo' -> normal combo, choose a COMBO_STYLE
                         'stack' -> stack, choose normal style
 --gdc-d2scale           scale of second data, default 0.00001

AUTHOR

  mathias gumt <gumz at cs dot uni-magdeburg dot de>

END

parser = GetoptLong.new
parser.set_options(
  [ "--count-type",   "-c",     GetoptLong::REQUIRED_ARGUMENT ],
  [ "--maildir",      "-d",     GetoptLong::REQUIRED_ARGUMENT ],
  [ "--recursive",    "-r",     GetoptLong::NO_ARGUMENT ],
  [ "--mbox",         "-m",     GetoptLong::REQUIRED_ARGUMENT ],
  [ "--output",       "-o",     GetoptLong::REQUIRED_ARGUMENT ],
  [ "--seperator",    "-s",     GetoptLong::REQUIRED_ARGUMENT ],
  [ "--output-type",  "-t",     GetoptLong::REQUIRED_ARGUMENT ],
  [ "--days",         "-D",     GetoptLong::NO_ARGUMENT ],
  [ "--months",       "-M",     GetoptLong::NO_ARGUMENT ],
  [ "--years",        "-Y",     GetoptLong::NO_ARGUMENT ],
  [ "--dayofweek",    "-W",     GetoptLong::NO_ARGUMENT ],
  [ "--hours",        "-H",     GetoptLong::NO_ARGUMENT ],
  [ "--version",      "-v",     GetoptLong::NO_ARGUMENT ],
  [ "--help",         "-h",     GetoptLong::NO_ARGUMENT ],
  [ "--gdc-title",              GetoptLong::REQUIRED_ARGUMENT ],
  [ "--gdc-xtitle",             GetoptLong::REQUIRED_ARGUMENT ],
  [ "--gdc-ytitle",             GetoptLong::REQUIRED_ARGUMENT ],
  [ "--gdc-titlesize",          GetoptLong::REQUIRED_ARGUMENT ],
  [ "--gdc-style",              GetoptLong::REQUIRED_ARGUMENT ],
  [ "--gdc-itype",              GetoptLong::REQUIRED_ARGUMENT ],
  [ "--gdc-width",              GetoptLong::REQUIRED_ARGUMENT ],
  [ "--gdc-height",             GetoptLong::REQUIRED_ARGUMENT ],
  [ "--gdc-bgcolor",            GetoptLong::REQUIRED_ARGUMENT ],
  [ "--gdc-d1color",            GetoptLong::REQUIRED_ARGUMENT ],
  [ "--gdc-d2color",            GetoptLong::REQUIRED_ARGUMENT ],
  [ "--gdc-type",               GetoptLong::REQUIRED_ARGUMENT],
  [ "--gdc-d2scale",            GetoptLong::REQUIRED_ARGUMENT] )


parser.each do |opt, arg|
  case opt

  when "--count-type"
    case arg
    when 'date'
      @setup['ctype']= 'date'
    else
      puts "error: no such count-type \"#{arg}\" implemented."
      exit 1
    end

  when "--days"
    @setup['format']= '%Y-%m-%d'
  when "--months"
    @setup['format']= '%Y-%m'
  when "--years"
    @setup['format']= '%Y'
  when "--dayofweek"
    @setup['format']= '%w'
  when "--hours"
    @setup['format']= '%H'



  when "--gdc-title"
    @setup['gdc-title']= arg
  when "--gdc-xtitle"
      @setup['gdc-xtitle']= arg
  when "--gdc-ytitle"
      @setup['gdc-ytitle']= arg

  when "--gdc-width"
    @setup['gdc-width']= arg.to_i
  when "--gdc-height"
    @setup['gdc-height']= arg.to_i
  when "--gdc-titlesize"
    if arg == "help"
      puts @gdc_fonttags.keys.sort
      exit 0
    end
    if @gdc_fonttags.key?(arg)
      @setup['gdc-tsize']= arg
    else
      puts "no such gdchart-fontsize \"#{arg}\". abort."
      exit 1
    end

  when "--gdc-style"
    if arg == "help"
      puts @gdc_styletags.keys.sort
      exit 0
    end
    if @gdc_styletags.key?(arg.upcase)
      @setup['gdc-style']= arg.upcase
    else
      puts "no such gdc-style \"#{arg}\". abort."
      exit 1
    end
  when "--gdc-itype"
    if arg == "help"
      puts @gdc_itypetags.keys.sort
      exit 0
    end
    if @gdc_itypetags.key?(arg.upcase)
      @setup['gdc-itype'] = arg.upcase
    else
      puts "no such gdc-itype \"#{arg}\". abort."
    end

  when "--gdc-bgcolor"
    @setup['gdc-bgcolor']= arg.hex
  when "--gdc-d1color"
    @setup['gdc-d1color']= arg.hex
   when "--gdc-d2color"
    @setup['gdc-d2color']= arg.hex

  when "--gdc-type"
    if arg =~ /single|combo|stack/
      puts "yeah"
      @setup['gdc-type']= arg
    else
      puts "no such gdc-type, either single, combo or stack is valid"
      exit 1
    end

  when "--gdc-d2scale"
    @setup['gdc-d2scale']= arg.to_f

  when "--seperator" && @setup['otype'] == 'csv'
    @setup['seperator']= arg

  when "--maildir"
    if !File.directory?(arg) && !File.readable?(arg)
      puts "error: couldnt access #{arg}, exit"
      exit 1
    end
    @setup['mail']= 'maildir'
    @setup['input']= arg
  when "--recursive"
    @setup['recursive'] = true
  when "--mbox"
    puts "not yet implemented"
    exit 0
  when "--output"
    @setup['output'] = File.open(arg, "w")


  when "--output-type"
    case arg
    when 'csv'
      @setup['otype']= 'csv'
    when 'gdchart'
      @setup['otype']= 'gdchart'
    else
      puts "error, not supported output-type \"#{arg}\"."
      exit 1
    end
  when "--help"
    puts MINIUSAGE
    puts USAGE
    exit 0
  when "--version"
    puts INFO
    exit 0
  end
end



def count(msg,fmt)
  case @setup['ctype']
  when 'date'
    if ( msg.header.date != nil )
      date = msg.header.date.strftime(fmt)
      if !@mails.key?(date)
        @mails[date] = [ 0, 0 ]
      end
      @mails[date][0] += 1
      @mails[date][1] += msg.to_s.length()
    end
  end
end

def countmaildir(dir, fmt, rec)
  Dir.foreach(dir) do |filename|
    if !(filename =~ /^\.*$/)
      Dir.chdir(dir)
      if File.directory?(filename)
        if rec
          countmaildir(Dir.pwd+"/"+filename, fmt, rec)
        end
      else
        File.open(filename, "r") do |file|
          RMail::Mailbox::parse_mbox(file) do |input|
            msg = RMail::Parser.read(input)
            count(msg,fmt)
          end
        end
      end
    end
  end
end

def countmbox(mbox, fmt)
  File.open(mbox, "r") do |file|
    RMail::Parser::parse_mbox(file) do |input|
      msg = RMail::Parser.read(input)
      count(msg,fmt)
    end
  end
end


def output
  out = @setup['output']
  arr = @mails.sort

  case @setup['otype']
  when 'csv'
    sep = @setup['seperator']
    if out != $stdout
      out.print HEADER
      out.print "#  set timefmt \"", @setup['format'], "\"\n"
      out.print "#  plot \"", out.path, "\" using 1:2 smooth bezier\n\n"
    end

    arr.each do
      |a|
      out.print a[0], sep, a[1][0], sep, a[1][1], "\n"
    end

  when 'gdchart'
    gdc = GDChart.new
    gdc.title = @setup['gdc-title']
    gdc.xtitle = @setup['gdc-xtitle']
    gdc.ytitle = @setup['gdc-ytitle']
    gdc.title_size = @gdc_fonttags[@setup['gdc-tsize']]
    gdc.image_type= @gdc_itypetags[@setup['gdc-itype']]
    gdc.BGColor = @setup['gdc-bgcolor']
    gdc.SetColor = [ @setup['gdc-d1color'], @setup['gdc-d2color']]

    data = Array.new
    label = Array.new

    arr.each { |a| label << a[0]; data << a[1][0] }
    case @setup['gdc-type']
    when 'single'
      gdc.out_graph(@setup['gdc-width'], @setup['gdc-height'], out, @gdc_styletags[@setup['gdc-style']], label.length, label, 1, data)
    when 'stack'
      arr.each { |a| data << (a[1][1] * @setup['gdc-d2scale']) }
      gdc.out_graph(@setup['gdc-width'], @setup['gdc-height'], out, @gdc_styletags[@setup['gdc-style']], label.length, label, 2, data)
    when 'combo'
      size = []
      arr.each { |a| size << (a[1][1]) }
      gdc.out_graph(@setup['gdc-width'], @setup['gdc-height'], out, @gdc_styletags[@setup['gdc-style']], label.length, label, 1, size, data)
    end

  end
end

case @setup['mail']
when "maildir"
  countmaildir(@setup['input'], @setup['format'], @setup['recursive'])
when "mbox"
  countmbox(@setup['input'], @setup['format'])
else "none"
  puts MINIUSAGE
  exit 0
end

output