#!/usr/bin/env ruby

require 'ldap'
require 'yaml'
require 'date'
require 'fileutils'

CONFIGNAME = "./userblocker.conf"
MAXINT = 1 << 64
SECONDS_IN_DAY = 24 * 60 * 60

$config = {
  :debug => false,

  :config_file => '/etc/userblocker/userblocker.conf',

  :squid_command => '/usr/sbin/squid -k reconfigure',
  :squid_reload => false,

  :path_lightsquid => '/var/lib/lightsquid',
  :path_output => './',
  :file_users     => "users.all",
  :file_blocked   => "users.blocked",
  :file_banned    => "users.banned",
  :file_hold      => "users.hold",
  :file_expired   => "users.expired",

  :format_file_users    => "%u",
  :format_file_blocked  => "%u %1 %2",
  :format_file_banned   => "%u",
  :format_file_hold     => "%u %1",
  :format_file_expired  => "%u %1",

  :ldap_host => 'localhost',
  :ldap_port => LDAP::LDAP_PORT,
  :ldap_bind_dn => '',
  :ldap_bind_pw => '',
  :ldap_base_dn => '',
  :ldap_scope => LDAP::LDAP_SCOPE_SUBTREE,

  :uid_filter => 'objectClass=*',
  :uid_attr => 'uid',
  :traffic_attr => 'traffic',
  :hold_attr => 'holdUntil',
  :banned_attr => 'banned',
  :expire_attr => 'expires',

  :hold => true,
  :period => :month,
  :mb_size => 1024*1024,
}

def getConfig( file )

    configfile = FileTest.exists?( file ) ? file : CONFIGNAME

    if ! File.readable?( configfile ) then
        puts "Error opening config file #{ configfile }."
        exit 1
    end

    $config.merge!(YAML.load(File.open(configfile)))
end

def msgDbg( error )
    error = getTime() + ': ' + error

    $config[:debug] and puts error
end

def getTime()
    return Time.now.to_i.to_s
end

def check_paths()

  if ( ! File.directory?( $config[:path_lightsquid] )) ||
     ( ! File.readable?( $config[:path_lightsquid] ))
    puts "Error opening lightsquid directory #{ $config[:path_lightsquid] }."
    exit 1
  end
end

def connect_to_ldap()
  ldap = LDAP::Conn.open( $config[:ldap_host], $config[:ldap_port] )
  ldap.simple_bind( $config[:ldap_bind_dn], $config[:ldap_bind_pw] )
  return ldap
end

def get_userdata()
  users={}

 $ldap.search(
   $config[:ldap_base_dn],
   $config[:ldap_scope],
   $config[:uid_filter],
   [
     $config[:uid_attr],
     $config[:traffic_attr],
     $config[:hold_attr],
     $config[:expire_attr],
     $config[:banned_attr],
   ]
 ) { |res|
      username      = res[ $config[:uid_attr]     ]
      traffic_limit = res[ $config[:traffic_attr] ]
      hold_until    = res[ $config[:hold_attr]    ]
      expire_after  = res[ $config[:expire_attr]  ]
      banned        = res[ $config[:banned_attr]  ]

      username      = username[0]                         if username
      traffic_limit = traffic_limit[0].to_f               if traffic_limit
      banned        = (banned[0]=='TRUE')                 if banned
      hold_until    = hold_until   ? validate_date( hold_until[0] )   : Date.today-1
      expire_after  = expire_after ? validate_date( expire_after[0] ) : Date.today+1

      if username
        users[ username ] = {
          :traffic => 0,
          :traffic_limit => traffic_limit,
          :hold_until => hold_until,
          :expire_date => expire_after,
          :banned => banned,
        }
      end
   }
  return users
end

def get_traffic( users )

  today=Date.today

  case $config[:period]
  when :daily
    date_from=today
  when :monthly
    date_from=Date.new(today.year,today.month)
  when :last_month
    date_from=today-30
  when :weekly
    date_from=today-today.wday+1
  when :last_week
    today=today-7
  else
    date_from=Date.new(today.year,today.month)
  end

  ( date_from .. today ).each do |date|
    begin
      filename = "#{$config[:path_lightsquid]}/#{date.to_s.delete("-")}/.total"
      File.readlines( filename ).each do |line|
        line = line.chomp.split[0..1]
        user = line[0];       traffic = line[1].to_f

        next if ["user:","size:"].include?(user)
        next unless users.include?(user)

        users[user][:traffic] += traffic / $config[:mb_size]

      end
    rescue => m
      raise m unless m.class == Errno::ENOENT   # ignore the days with no traffic
    end
  end
end

def validate_date ( datestring )
  datestring.strip!
  if datestring.match(/^\d{4}-\d{2}-\d{2}/)       # => YYYY-MM-DD
    return Date.strptime( datestring )
  else
    return nil
  end
end

def write_data( filetype, users, format=nil )
  File.open( "#{$config[:path_output]}/#{$config[filetype]}", "w" ) { |file|
    users.each { |user|
      if format && user.is_a?(Array)
        result=$config[format].gsub(/%u/, user[0].to_s)
        result=result.gsub(/%1/, user[1].to_s) if user[1]
        result=result.gsub(/%2/, user[2].to_s) if user[2]
        result=result.gsub(/%3/, user[3].to_s) if user[3]
      else
        result=user
      end
      file.puts( result )
    }
  }
end

def main
  getConfig( $config[:config_file] )
  msgDbg( "Config loaded: #{$config.inspect}" )

  check_paths

  $ldap = connect_to_ldap

  blocked_users   = []
  banned_users    = []
  hold_users      = []
  expired_users   = []

  users=get_userdata
  get_traffic(users)

  users.each do |user,stats|
    traffic       = stats[:traffic]
    traffic_limit = stats[:traffic_limit]
    hold_until    = stats[:hold_until]
    banned        = stats[:banned]
    expire_date   = stats[:expire_date]

    if traffic >= traffic_limit
      if hold_until >= Date.today
        hold_users.push( [user, hold_until] )
      else
        blocked_users.push( [user, traffic.to_i, traffic_limit.to_i] )
      end
    end

    banned_users.push( user )                     if stats[:banned]

    expired_users.push( [user, expire_date] )     if expire_date < Date.today
  end

  $ldap.unbind   if $ldap.bound?

  write_data( :file_users,    users.map{|x| x[0]}, :format_file_users   )
  write_data( :file_blocked,  blocked_users,       :format_file_blocked )
  write_data( :file_banned,   banned_users,        :format_file_banned  )
  write_data( :file_hold,     hold_users,          :format_file_hold    )
  write_data( :file_expired,  expired_users,       :format_file_expired )

  system( $config[:squid_command] )               if $config[:squid_reload]

end

if ARGV[0] == "configdump" then
  case ARGV.size
  when 1
    YAML.dump( $config, STDOUT )
  when 2
    File.open( ARGV[1], 'w' ) do |out|
        YAML.dump( $config, out )
    end
  end
elsif ARGV.size == 0 then
    main()
else
    puts " Unknown command: #{ARGV[0]}"
    puts " Usage: #{__FILE__} configdump /path/to/dump.conf"
end
