#!/usr/bin/perl -w

use strict;

# ipquota.pl
#
# (c) 2003 Kim Holburn <kim@holburn.net>
#
# Download from :
# http://www.holburn.net/software/ipquota.perl.txt
#
# released under the GPL v2 http://www.gnu.org/copyleft/gpl.html
# 
# This script enforces a host-based quota on a linux firewall.
# It requires net-acct and ipstat.pl
#    net-acct is a network accounting program
#    ipstat.pl is a script to analyse net-acct logs.
#
# To limit the resources used on the firewall this script
# is intended to be run once per hour and keep logs.  
# It could be run at much smaller intervals.  It is intended to
# send an email when the quota for an IP is exceeded. It could also
# activate tc bandwidth throttling or iptables style rules.
#
#  Author =   $Author: kimh $
#  Date =     $Date: 2006/05/31 03:55:08 $
#  Id =       $Id: ipquota.pl,v 1.24 2006/05/31 03:55:08 kimh Exp $
#  Revision = $Revision: 1.24 $
#  State =    $State: Exp $
#  Log =      $Log: ipquota.pl,v $
#  Log =      Revision 1.24  2006/05/31 03:55:08  kimh
#  Log =      added upload and download counts
#  Log =
#  Log =      Revision 1.23  2006/05/17 08:03:04  kimh
#  Log =      log recent check faster
#  Log =
#  Log =      Revision 1.22  2006/05/17 07:55:07  kimh
#  Log =      log recent check faster
#  Log =
#  Log =      Revision 1.21  2006/05/17 06:27:28  kimh
#  Log =      typos added external checks
#  Log =
#  Log =      Revision 1.20  2006/05/17 05:53:57  kimh
#  Log =      typos added external checks
#  Log =
#  Log =      Revision 1.19  2006/05/17 05:53:16  kimh
#  Log =      typos added external checks
#  Log =
#  Log =      Revision 1.18  2006/05/17 03:10:13  kimh
#  Log =      added external checks
#  Log =
#  Log =      Revision 1.17  2006/05/17 01:22:12  kimh
#  Log =      only look at current log file most of the time
#  Log =
#  Log =      Revision 1.16  2006/05/12 04:29:00  kimh
#  Log =      added connection alarms
#  Log =
#  Log =      Revision 1.15  2006/05/11 23:37:28  kimh
#  Log =      added help debug option
#  Log =
#  Log =      Revision 1.14  2006/05/11 23:34:47  kimh
#  Log =      new options prelude to adding connection counts
#  Log =
#  Log =      Revision 1.11  2006/03/13 00:28:10  kimh
#  Log =      typos
#  Log =
#  Log =      Revision 1.10  2006/03/13 00:00:51  kimh
#  Log =      format more flexible
#  Log =
#  Log =      Revision 1.9  2005/08/18 22:56:01  kimh
#  Log =      changed internalnets.txt to free.txt
#  Log =
#  Log =      Revision 1.8  2005/08/18 05:01:28  kimh
#  Log =      fixing rel time bug
#  Log =
#  Log =      Revision 1.7  2004/11/24 06:19:03  kimh
#  Log =      added command line of ipstat command
#  Log =
#  Log =      Revision 1.6  2004/08/14 00:46:06  kimh
#  Log =      fixed stdout message
#  Log =
#  Log =      Revision 1.5  2004/06/20 23:25:26  kimh
#  Log =      normalised names of conf, var etc.
#  Log =
#  Log =      Revision 1.4  2004/03/29 00:11:29  kimh
#  Log =      added download URL, added port comment
#  Log =
#  Log =      Revision 1.3  2004/03/24 00:38:07  kimh
#  Log =      changed comment
#  Log =
#  Log =      Revision 1.2  2004/01/29 05:19:25  kimh
#  Log =      iq fixed stderr if no dump file
#  Log =
#  Log =      Revision 1.1  2003/12/16 10:28:44  kimh
#  Log =      enforce host-based ipquotas
#  Log =


my $quotamaster="root";
# config is a list of hosts and quotas
my $docheck = 1;
my $dohour = 1;
my $var = "/var/log";
my $etc = "/etc";
my $config = "$etc/ipquota.conf";
my $dir = "$var/ipquota";
my $quotalog="$dir/ipquota.log";
my $file1="$var/net-acct/net-acct.log";
my $file2="$file1.0";
my $files1="$file1 $file2";
my $files=$files1;
# dump is net-acct logs of current connections
my $dump="$var/net-acct/dump";
my $dumpout="$dir/dump";
my $mail = "/usr/bin/mail";
my $ipstat_pl = "$etc/ipstat.pl";
# these are internal (free networks so exclude them from consideration)
my $internals = "$etc/free.txt";
my $internale = "";
if (-e $internals) { $internale = "-e $internals "; }
my $debug=0;
my $verbose=0;
if ( ! -d "$dir" ) { mkdir $dir; }

sub fail_usage {
  if (scalar @_) {
    print "$0: Error: \n";
    map {print "  $_\n";} @_;
  }
  print <<EOM;

  $0 Usage:
    $0 [option] 
      -s (hour stats)
      -c (check)
      -b (both (default))
      -d (debug)
    $0 -h|--help
       show this screen
      

EOM
  exit ;
} 
  
my $ARGV = $ARGV[0];
while (scalar @ARGV) {
#  if ($ARGV =~ /^-C/ ) {
#    if ($ARGV eq '-C') {
#      my $myarg = $ARGV;
#      shift @ARGV; $ARGV = $ARGV[0];
#      if (!$ARGV) { fail_usage ("$myarg needs an argument"); }
#      if ($ARGV=~/^-/) {fail_usage("$myarg must be followed by an argument");}
#    }
#    else {  #elsif ($ARGV =~ /^-C[^\s]/ ) 
#      $ARGV =~ s/^-C//;
#    }
#    $opt_m = $ARGV;
#    shift @ARGV; $ARGV = $ARGV[0];
#  }

  if ($ARGV eq '-h' || $ARGV eq '--help') {
    fail_usage;
  }
  elsif ($ARGV =~ /^-b|^--both$/ ) {
    $dohour = $docheck = 1; 
    if ($ARGV =~ /^-b$|^--both$/) { shift @ARGV; $ARGV = $ARGV[0]; }
    else { $ARGV =~ s/^-b/-/; }
  }
  elsif ($ARGV =~ /^-c|^--check$/ ) {
    $docheck = 1; 
    $dohour = 0; 
    if ($ARGV =~ /^-c$|^--check$/) { shift @ARGV; $ARGV = $ARGV[0]; }
    else { $ARGV =~ s/^-c/-/; }
  }
  elsif ($ARGV =~ /^-s|^--stats$/ ) {
    $dohour = 1; 
    $docheck = 0; 
    if ($ARGV =~ /^-s$|^--stats$/) { shift @ARGV; $ARGV = $ARGV[0]; }
    else { $ARGV =~ s/^-s/-/; }
  }
  elsif ($ARGV =~ /^-d|^--debug$/ ) {
    $debug = 1; 
    if ($ARGV =~ /^-d$|^--debug$/) { shift @ARGV; $ARGV = $ARGV[0]; }
    else { $ARGV =~ s/^-d/-/; }
  }
  elsif ($ARGV =~ /^-v/ ) {
    $verbose ++ ;
    if ($ARGV eq "-v") { shift @ARGV; $ARGV = $ARGV[0]; }
    else { $ARGV =~ s/^-v/-/; }
  }
  elsif ($ARGV =~ /^-/ ) {
    fail_usage " invalid option ($ARGV)";
  } else {
    last ;
  }
}

my $epoch = time;
use POSIX qw(strftime);
my $hour=strftime('%Y-%m-%dT%H', localtime($epoch)); 
my $time=strftime('%Y/%m/%d-%H:%M:%S', localtime($epoch)); 

# we only need the current log most of the time
# only need last log in the first 2 hours after log rotation
my $time1 = getfiletime ($file1);
if ($epoch - $time1 > 720) { $files = "$file1"; }

sub getfiletime {
  my $file = shift;
  my $file1 = $file;
  if ($file =~ /z$/i) { $file1 = "zcat $file|"; }
  open (FILE, $file1) or die "Couldn't open file ($file)";
  my $line = <FILE> ;
  close FILE;
  my ($time, $junk) = split (/\t/, $line, 2);
  $time;
}

my @logs;
sub do_logs {
  # get a list of all recent entries
  if (open (QUOTA, $quotalog)) {
    @logs = grep {
      my ($date, $htime, $host, $value) = split (" ", $_, 4);
      $epoch - $htime < 60*60*24;
    } <QUOTA>;
    close QUOTA;
  }
  if (! defined ($logs[0])) { $logs[0] = ""; }
}

sub myalarm {
  my ($host, $size, $asize, $mess, $extra) = @_;
 
  if (! defined ($logs[0])) { do_logs; }
  my $found; 
#    print "host=($host)($size)\n";
  for (@logs) { if (/ $host / && /$mess mail$/) {
    if ($debug) { print "debug old skipping h($host)m($mess)\n"; }
    return;
  } }
#  if (open (QUOTA, $quotalog)) {
#    while (<QUOTA>) { if (/ $host / && /$mess mail$/) { $found=$_; } }
#    close QUOTA;
#    if ($found) {
#      chomp $found;
#      if ($debug) { print "debug found=($found) h($host)m($mess)\n"; }
#      my ($date, $htime, $host, $value) = split (" ", $found, 4);
#      if ($epoch - $htime < 60*60*24) { return; } # only once each day
#    }
#  }
  my $mess1 = "$time $epoch $host $size $asize $mess mail\n" ; 
  logit ($host, $size, $asize, $mess, $mess1);
  mailit ($host, $size, $asize, $mess, $mess1, $extra);
}

sub logit {
  my ($host, $size, $asize, $mess, $mess1) = @_;
  
  if ($debug) { print "host=($host)($size)\n"; }
  if (!$mess1) { $mess1 = "$time $epoch $host $size $asize $mess mail\n" ; }
#  print $mess1 ;
  if (!$debug) {
    if (open (QUOTA, ">>$quotalog")) {
      print QUOTA $mess1 ;
      close QUOTA;
    } else { print STDERR "Could not write to log ($quotalog)\n"; }
  }
  else { warn "debug ", $mess1 ; }
}
sub mailit {
  my ($host, $size, $asize, $mess, $mess1, $extra) = @_;

  my $xopt = ($extra ne "extra") ? "-w -U":"-D";
  if ($extra eq "extrax") { $internale = "-E"; $xopt = "-w"; }
  my $mailc = "nice -n 19 $ipstat_pl $internale $xopt -r -S \\
    -H $host --from \"-0/0/1 00:00:00\" --to \"now\" $files1 |";
  my $abbrsize = printbytes ($size); 
  my $abbrasize = printbytes ($asize); 
  my $hosti = $host;
  $host = ip_to_name ($host);
  my ($bytes, $quota, $down, $conn) =
       ("Bytes", "quota", "downloaded", "");
  if ($extra ne "extra") {
    ($bytes, $quota, $down, $conn) =
       ("Connections", "connections quota", "made","connections");
  }
  $mess1 = "host '$host' has gone over it's $mess $quota\n\n";
  $mess1 .= " $host has $down $size($abbrsize) ${bytes} \n\n";
  $mess1 .= " ${host}'s $mess $quota is $asize ($abbrasize)${bytes} \n\n";
  if ($extra ne 'extrax') {
    $mess1 .= " See: <https://network.example.com/assets/search_ips.pl" . 
      "?ipordns=$hosti> \n\n";
  }
 
  my $cmd="| $mail -s \"$host ($abbrsize) $mess $quota\" $quotamaster\n"; 
  if ($debug) {
    print "debug $cmd\n";
    print "debug \"$mess1\" \n";
  }
  else {
    if ($extra && open CMD, $cmd) {
      print CMD $mess1;
      if (open INFO, $mailc) {
        print CMD "\n", "-"x40, "\n";
        print CMD $mailc, "\n";
        print CMD "\n", "-"x40, "\n";
        print CMD <INFO>; 
        close INFO;
        print CMD "\n", "-"x40, "\n";
      } else { print CMD $mess1; }
      close CMD;
    } else { system "echo \"$mess1\" $cmd"; } 
  }
}

# ip_to_name: return the name of of the given ip address
sub ip_to_name {
  my ($ip) = @_;

  my $pattern_ip =
    "[[:digit:]]{1,3}\.[[:digit:]]{1,3}\.[[:digit:]]{1,3}\.[[:digit:]]{1,3}";

  if (!$ip) { return ("[no ip]"); }
  if (!($ip =~ /^$pattern_ip$/)) { return ("$ip\[bad ip]"); }

  my ($host_name, $aliases, $addrtype, $length) ;
  my @addrs;
  my $ipaddr = pack("C4", split(/\./, $ip));
  if (($host_name, $aliases, $addrtype, $length, @addrs)
             = gethostbyaddr($ipaddr, 2)) { return ("$host_name\[$ip]"); }
  else { return ("$ip\[lookup failed]"); }
  return ($ip);
}

sub parseabbrev {
  my $numb=shift;
  $numb =~ s/\s//g;
  if ($numb =~ /^\d*$/) { return $numb; }
  if ($numb =~ /^(\d+)([bkmgtx])$/i) {
    $numb = $1;
    my $factor = lc $2;
    if    ($factor eq 'x') { $numb *= 1000*1000*1000*1000*1000; }
    elsif ($factor eq 't') { $numb *= 1000*1000*1000*1000; }
    elsif ($factor eq 'g') { $numb *= 1000*1000*1000; }
    elsif ($factor eq 'm') { $numb *= 1000*1000; }
    elsif ($factor eq 'k') { $numb *= 1000; }
  } else { $numb = 0; }
  $numb;
}

sub printbytes {
  my ($val)= @_;
  my $suffix = "" ;
  if ($val>1024) { $val/=1024; $suffix="K"; }
  if ($val>1024) { $val/=1024; $suffix="M"; }
  if ($val>1024) { $val/=1024; $suffix="G"; }
  if ($val>1024) { $val/=1024; $suffix="T"; }
  my $val2 = sprintf("%.1f", $val);
  $val2 =~ s/\.0$//;
  sprintf("%s", "$val2$suffix");
}

if ($dohour) {
  my $dirs = $hour;
  $dirs =~ s/T.*$//g;
  $dirs =~ s#-#/#g;
  my @dirs = split /\//, $dirs;
  my $sdir=$dir;
  for (@dirs) {
    $sdir .= "/".$_;
    if ( ! -d $sdir) { mkdir $sdir; }
  }

  my $out="$dir/$dirs/$hour";
  my $lout="$dir/current.$hour";
# don't redo data for an hour
  if ( ! -e "$out" ) {
    if ($debug) { print "creating hourly data for $hour...\n"; }
    system ("nice -n 19 $ipstat_pl $internale -c -1 -f \\
             --from -01:00:00 --to \"$hour\" $files > \"$out\"");
#    system ("ln", "-s", $out, $lout);
    system ("(cd \"$dir\"; ln -s \"$dirs/$hour\" \"current.$hour\")");
    if ( -e $dump) {
      system ("nice -n 19 $ipstat_pl $internale -c -1 -f \\
             --from -01:00:00 --to \"$hour\" $dump > \"$dumpout\"");
    }
  } else { print "hourly data for $hour already present\n"; }
}

if ($docheck) {
  my %warn;
  my %stop;
  my %kill;
  my %warnc;
  my %stopc;
  my %killc;
#default:500M:1G:2G:2k:4k:8k
  $warn{'default'} =  parseabbrev '500M';
  $stop{'default'} =  parseabbrev '1G';
  $kill{'default'} =  parseabbrev '2G';
  $warnc{'default'} = parseabbrev '2k';
  $stopc{'default'} = parseabbrev '4k';
  $killc{'default'} = parseabbrev '8k';
# read config file and store config data
  if (!open CONF, $config) { die "couldn't read file \"$config\""; }
  map {
    chomp;
    if (/:/) {
      my ($ip, $warn, $stop, $kill, $warnc, $stopc, $killc) = split /:/;
      $warn = parseabbrev ($warn);
      $stop = parseabbrev ($stop);
      $kill = parseabbrev ($kill);
      if ($kill && $kill < $stop) { $stop = $kill; }
      if ($stop && $stop < $warn) { $warn = $stop; }
      $warn{$ip} = $warn;
      $stop{$ip} = $stop;
      $kill{$ip} = $kill;
      $warnc = parseabbrev ($warnc);
      $stopc = parseabbrev ($stopc);
      $killc = parseabbrev ($killc);
      if ($killc && $killc < $stopc) { $stopc = $killc; }
      if ($stopc && $stopc < $warnc) { $warnc = $stopc; }
      $warnc{$ip} = $warnc;
      $stopc{$ip} = $stopc;
      $killc{$ip} = $killc;
    } elsif (/\@/) {
      my ($ip, $name) = split /\@/;
      if (defined($warn{$name})) {
        $warn{$ip} = $warn{$name};
        $stop{$ip} = $stop{$name};
        $kill{$ip} = $kill{$name};
        $warnc{$ip} = $warnc{$name};
        $stopc{$ip} = $stopc{$name};
        $killc{$ip} = $killc{$name};
      }
    }
  } grep !/^#/, <CONF>;
  close CONF;

  opendir DIR, $dir or die "couldn't open directory ($dir)";
  my @list = grep /^current/, readdir(DIR);
  closedir DIR;
  @list = sort map { "$dir/$_" } @list;
  while (scalar @list > 24) {
    my $old = shift @list;
  if ($debug) {  print "removing old links ($old)\n"; }
    if (-l $old) { system ("rm", $old); }
  }
  push @list, $dumpout;

  my %ip;
  my %ipc;
  if ($debug) { print "reading data for last 24 hours...\n"; }
  for my $list (@list) {
    if (!open FILE, $list) { print STDERR "Couldn't read file ($list)"; next; }
    if ($debug) { print "reading file ($list)\n"; }
    my $line=0;
    map {
#      if (!/^#/) { #}
      chomp;
      if (/^\d+\.\d+\.\d+\.\d+(?::\d+)+$/) {
        my ($ip, $size, $inc, $outc, $rest) = split /:/, $_, 5; 
        $ip{$ip} += $size; 
        if (!defined($ipc{$ip}) || $ipc{$ip} < $outc) { $ipc{$ip} = $outc; }
        if ($debug && $ip =~ /214\.47/)
          { print "ip($ip)=s($size)c($ip{$ip})\n"; }
      }
      elsif (!/^#/ && !/^\s*$/) {
        warn "ipquota format error: file($list) line($line)=($_)\n";
      }
      $line++;
    } <FILE>;
    close FILE;
  }
  for my $ipn (keys %ip) {
    my $warn = (defined($warn{$ipn}))?$warn{$ipn}:$warn{'default'};
#    if ($debug && $ipn =~ /\.47/)
#      { print "ip($ipn)  warn=($warn) size=($ip{$ipn})\n"; }
    if ($warn && ($ip{$ipn} > $warn))
      { myalarm ($ipn, $ip{$ipn}, $warn, "warn", "extra" ); } 
    my $stop = (defined($stop{$ipn}))?$stop{$ipn}:$stop{'default'};
#    if ($debug && $ipn =~ /\.47/)
#      { print "ip($ipn)  stop=($stop) size=($ip{$ipn})\n"; }
    if ($stop && ($ip{$ipn} > $stop)) {
      myalarm ($ipn, $ip{$ipn}, $stop, "stop", "extra" );
#      system ( "iptables -I fw_quick -d $ipn -j DROP" );
#      system ( "iptables -I fw_quick -s $ipn -j DROP" );
    } 
  }
  for my $ipn (keys %ipc) {
    my $warn = (defined($warnc{$ipn}))?$warnc{$ipn}:$warnc{'default'};
    if ($warn && ($ipc{$ipn} > $warn))
      { myalarm ($ipn, $ipc{$ipn}, $warn, "warnc", "extrac" ); } 
    my $stop = (defined($stopc{$ipn}))?$stopc{$ipn}:$stopc{'default'};
    if ($stop && ($ipc{$ipn} > $stop)) {
      myalarm ($ipn, $ipc{$ipn}, $stop, "stopc", "extrac" );
#      system ( "iptables -I fw_quick -d $ipn -j DROP" );
#      system ( "iptables -I fw_quick -s $ipn -j DROP" );
    }
  }
  my $conntest="$ipstat_pl -E -w -U -r -f --from -01:00:00 --to $hour $files";

  my $warnc=(defined($warnc{'external'}))?$warnc{'external'}:$warnc{'default'};
  my $stopc=(defined($stopc{'external'}))?$stopc{'external'}:$stopc{'default'};
  my $killc=(defined($killc{'external'}))?$killc{'external'}:$killc{'default'};
  if (!open FILE, "-|", $conntest)
    { print STDERR "Couldn't read pipe ($conntest)"; exit 1; }
  if ($debug) { print "reading pipe ($conntest)\n"; }
  my $line=0;
  map {
    chomp;
    if (/^#/ || /^\s*$/) { ; }
    elsif (/^\d+\.\d+\.\d+\.\d+(?::\d+)+$/) {
      my ($ip, $size, $inc, $outc, $rest) = split /:/, $_, 5;
      if ( $warnc && $inc > $warnc) 
        { myalarm ($ip, $inc, $warnc, "warnx", "extrax" ); }
      if ( $stopc && $inc > $stopc) 
        { myalarm ($ip, $inc, $stopc, "stopx", "extrax" ); }
      if ( $killc && $inc > $killc) 
        { myalarm ($ip, $inc, $killc, "killx", "extrax" ); }
  #    if ($debug && $ip =~ /214\.47/)
  #      { print "ip($ip)=s($size)c($ip{$ip})\n"; }
    }
    else { warn "ipquota format error: file($conntest) line($line)=($_)\n"; }
    $line++;
  } <FILE>;
  close FILE;
}