#!/usr/local/bin/ruby
# Unofficial TxD Disk & HTTP Bandwidth Usage Meter (rsrcmeter)
# Version 2.0.4
# Written by AJ Zmudosky
# rsrcmeter.textjoy.com
#
# THIS SOFTWARE IS PROVIDED BY THE AUTHOR "AS IS" AND WITHOUT WARRANTY OF
# ANY KIND. IN NO EVENT SHALL THE AUTHOR OR CONTRIBUTORS BE LIABLE FOR ANY
# DAMAGES HOWEVER CAUSED, AND ON ANY THEORY OF LIABILITY, ARISING IN ANY
# WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY
# OF SUCH DAMAGE.

require 'date'
require 'optparse'

# E-mail regexp
module RFC2822
  EmailAddress = begin
    alpha = "a-zA-Z"
    digit = "0-9"
    atext = "[#{alpha}#{digit}\!\#\$\%\&\'\*+\/\=\?\^\_\`\{\|\}\~\-]"
    dot_atom_text = "#{atext}+([.]#{atext}*)*"
    dot_atom = "#{dot_atom_text}"
    qtext = '[^\\x0d\\x22\\x5c\\x80-\\xff]'
    text = "[\\x01-\\x09\\x11\\x12\\x14-\\x7f]"
    quoted_pair = "(\\x5c#{text})"
    qcontent = "(?:#{qtext}|#{quoted_pair})"
    quoted_string = "[\"]#{qcontent}+[\"]"
    atom = "#{atext}+"
    word = "(?:#{atom}|#{quoted_string})"
    obs_local_part = "#{word}([.]#{word})*"
    local_part = "(?:#{dot_atom}|#{quoted_string}|#{obs_local_part})"
    no_ws_ctl = "\\x01-\\x08\\x11\\x12\\x14-\\x1f\\x7f"   
    dtext = "[#{no_ws_ctl}\\x21-\\x5a\\x5e-\\x7e]"
    dcontent = "(?:#{dtext}|#{quoted_pair})"
    domain_literal = "\\[#{dcontent}+\\]"
    obs_domain = "#{atom}([.]#{atom})*"
    domain = "(?:#{dot_atom}|#{domain_literal}|#{obs_domain})"
    addr_spec = "#{local_part}\@#{domain}"
    pattern = /^#{addr_spec}$/
  end
end

# Default options
opt = {
  :base_dir => ENV['HOME'],
  :stat_file => 'histstat',
  :temp_file => 'temp-rsrcmeter-bwcalc',
  :bw_months => 5
}

# Our option flag parsing
OptionParser.new do |opts|
  opts.banner = "Unofficial TxD Disk & HTTP Bandwidth Usage Meter\nUsage: rsrcmeter [options]"
  
  opts.on("-b", "--base BASE", "Change the base directory, defaults to user's $HOME", String) {|base|
    if File.exists?(base)
      opt[:base_dir] = base
    else
      raise(ArgumentError, "Directory (#{base}) specified with -b does not exist")
    end
  }
  
  opts.on("--stat-file FILE", "Change default historical data file from '#{opt[:stat_file]}'", String) {|file|
    opt[:stat_file] = file
  }
  
  opts.on("--temp-file FILE", "Change default temporary file from '#{opt[:temp_file]}'", String) {|file|
    opt[:temp_file] = file
  }
  
  opts.on("-m NUM", "Number of months' bandwidth to show prior to current month (default: #{opt[:bw_months]})", Integer) {|mons|
    opt[:bw_months] = mons if mons > 0
  }
  
  opts.on("-e", "--email ADDRESS", "E-mail results to ADDRESS *instead* of outputting to the screen", String) {|addr|
    if RFC2822::EmailAddress =~ addr
      opt[:email_to] = addr
    else
      raise(ArgumentError, "Invalid e-mail address provided to -e")
    end
  }
  
  opts.on_tail("-h", "--help", "Show this message") do
    puts opts
    exit
  end
end.parse!

Dir.chdir(opt[:base_dir])
File.delete(opt[:temp_file]) if File.exists?(opt[:temp_file])

# Look for existing monthly statistics file
if File.exists?(opt[:stat_file])
  # Load and process existing stats
  @stats = Array.new
  lines = File.readlines(opt[:stat_file])
  lines.each {|l| (@stats << [l.chomp.split(',')[0], l.chomp.split(',')[1].to_f]) unless (l.strip)[0,1] == '#'}
  @stats.sort!
  # Determine if any months need to be added
  if Date.today.month == 1
    pm_year = Date.today.year - 1
    pm_month = 12
  else
    pm_year = Date.today.year
    pm_month = Date.today.month - 1
  end
  stop_month = pm_year.to_s + (pm_month < 10 ? '0' : '') + pm_month.to_s
  proc_year = @stats.last[0][0,4].to_i
  proc_month = @stats.last[0][4,2].to_i
  proc_year += 1 if proc_month == 12
  proc_month = (proc_month == 12) ? 1 : proc_month + 1
  proc_str = proc_year.to_s + (proc_month < 10 ? '0' : '') + proc_month.to_s
  # Add any months not already calculated
  while proc_str <= stop_month
    datestr = proc_str + "??"
    File.delete(opt[:temp_file]) if File.exists?(opt[:temp_file])
    `cat logs/access_log.#{datestr} 2>/dev/null > #{opt[:temp_file]}`
    `cat domains/*/logs/access_log.#{datestr} 2>/dev/null >> #{opt[:temp_file]}`
    `zcat logs/access_log.#{datestr}.gz 2>/dev/null >> #{opt[:temp_file]}`
    `zcat domains/*/logs/access_log.#{datestr}.gz 2>/dev/null >> #{opt[:temp_file]}`
    month_usage = `cat #{opt[:temp_file]} | awk '{sum += $10} END {print sum}'`.chomp.to_i
    @stats << [proc_str, month_usage]
    # advance counting
    proc_year += 1 if proc_month == 12
    proc_month = (proc_month == 12) ? 1 : proc_month + 1
    proc_str = proc_year.to_s + (proc_month < 10 ? '0' : '') + proc_month.to_s
  end
# Statistics files doesn't exist
else
  # Look for earliest log file
  file_list = Array.new
  Dir.foreach(opt[:base_dir] + "/logs") {|i| file_list << i if /^access_log.\d{8}/ =~ i}
  if File.exists?(opt[:base_dir] + "/domains")
    domains = Array.new
    Dir.foreach(opt[:base_dir] + "/domains") {|dom| domains << dom if /^\w+\.\w+/ =~ dom}
    domains.each {|dom| Dir.foreach(opt[:base_dir] + "/domains/" + dom + "/logs") {|i| file_list << i if /^access_log.\d{8}/  =~ i}}
  end
  file_list.sort!
  unless Date.today.year.to_s + Date.today.month.to_s == file_list.first[11,6]
    # Process from earliest_log year/month forward
    earliest_log_year = file_list.first[11,4]
    earliest_log_month = file_list.first[15,2]
    @stats = Array.new
    earliest_log_year.to_i.upto(Date.today.year) do |year|
      start_month = (earliest_log_year.to_i == year) ? earliest_log_month.to_i : 1;
      start_month.upto(12) do |month|
        unless (year == Date.today.year and month >= Date.today.month)
          month = "0" + month.to_s if month < 10
          datestr = year.to_s + month.to_s + '??'
          File.delete(opt[:temp_file]) if File.exists?(opt[:temp_file])
          `cat logs/access_log.#{datestr} 2>/dev/null > #{opt[:temp_file]}`
          `cat domains/*/logs/access_log.#{datestr} 2>/dev/null >> #{opt[:temp_file]}`
          `zcat logs/access_log.#{datestr}.gz 2>/dev/null >> #{opt[:temp_file]}`
          `zcat domains/*/logs/access_log.#{datestr}.gz 2>/dev/null >> #{opt[:temp_file]}`
          month_usage = `cat #{opt[:temp_file]} | awk '{sum += $10} END {print sum}'`.chomp.to_i
          @stats << [year.to_s + month.to_s, month_usage]
        end
      end
    end
  end
end

# Write @stats back to opt[:stat_file] to record any changes
if (@stats.size > 0) then
  File.open(opt[:stat_file], 'w') {|f|
    f.puts "# rsrcmeter historical statistics file"
    f.puts "# Contains only completed months"
    @stats.each {|i| f.puts i[0] + ',' + i[1].to_s}
  }
end

# Calculate bandwidth consumption for current month
datestr = Date.today.year.to_s + Date.today.month.to_s + "??"
`cat logs/access_log 2>/dev/null > #{opt[:temp_file]}`
`cat domains/*/logs/access_log 2>/dev/null >> #{opt[:temp_file]}`
`cat logs/access_log.#{datestr} 2>/dev/null >> #{opt[:temp_file]}`
`cat domains/*/logs/access_log.#{datestr} 2>/dev/null >> #{opt[:temp_file]}`
`zcat logs/access_log.#{datestr}.gz 2>/dev/null >> #{opt[:temp_file]}`
`zcat domains/*/logs/access_log.#{datestr}.gz 2>/dev/null >> #{opt[:temp_file]}`
usage = `cat #{opt[:temp_file]} | awk '{sum += $10} END {print sum}'`.chomp.to_f / 1024 / 1024
File.delete(opt[:temp_file]) if File.exists?(opt[:temp_file])
# Determine disk usage
quotaline = `quota -g | tail -n 1`
disk_usage = `echo -n "#{quotaline}" | awk '{print $2}'`.to_f
disk_quota = `echo -n "#{quotaline}" | awk '{print $3}'`.to_f
disk_percent_used = (disk_usage / disk_quota) * 100

# Output results
results = ""
results << "Disk usage: " + sprintf("%.4f", disk_usage / 1024) + " MiB (Quota: " + sprintf("%.4f", disk_quota / 1024 / 1024) +" GiB | " + sprintf("%.1f", disk_percent_used) + "% used)\n"
results << "Bandwidth:\n" + Date::ABBR_MONTHNAMES[Date.today.month] + " " + Date.today.year.to_s + ": " + sprintf("%.4f", usage) + " MiB (Month to Date)\n"
hist_usage = @stats.last(opt[:bw_months]).reverse
hist_usage.each {|h| results << Date::ABBR_MONTHNAMES[Date.parse(h[0] + "01").month] + ' ' + h[0][0,4] + ": " + sprintf("%.4f", (h[1].to_f / 1024 / 1024)) + " MiB\n"}

# We're either e-mailing or outputting
if opt[:email_to]
  t = Time.now.strftime("%a %d %B %Y %H:%M:%S %Z")
  message = <<EOM
From: TxD Resource Meter <#{ENV['USER']}-noreply@#{`/bin/hostname`.chomp}>
To: #{opt[:email_to]}
Subject: [TxD Resource Meter] #{t} report for #{ENV['USER']}
X-Mailer: rsrcmeter | http://textsnippets.com/posts/show/842

Resource Report - #{t}
#{results}
---------------
Generated by rsrcmeter
EOM
  File.open("temp-emailresult", "w") { |file| file.print message }
  `cat temp-emailresult | /usr/sbin/sendmail -t`
  File.delete("temp-emailresult")
else
  puts results
end