#!/usr/bin/perl
# ddos.pl 2012.04.03

$| = 1;

# Net::Pcap (debian: libnet-pcap-perl)
# IPTables::ChainMgr (debian: libiptables-chainmgr-perl)

#--- EDIT THESE VALUES AS NEEDED ---#
use constant THRESH	=> 6;
use constant MAXPERSEC	=> 40;
use constant MAXCOUNT	=> 64;
use constant MAXMULTI	=> 16;
use constant MAXBLOCKS	=> 4;
use constant BASETIME	=> 450/256;
use constant BLOCKFILE	=> "/root/ddos.test";
use constant NETDEV	=> "eth0";
use constant NETMASK	=> "10.0.0.0/16";
#--- DO NOT EDIT PAST THIS POINT ---#


use POSIX qw(strftime);
use Switch;
use IPTables::ChainMgr;
#use Net::Pcap;
#use NetPacket::Ethernet;
#use NetPacket::IP;
#use NetPacket::TCP;
#use strict;

my $blocked;
my @data;
my $dom;
my @tmp;
my $IP;
my $port;
my $key;
my $line;
my $hms;
my %pkt;
my $tmr = time();

my %opts = (
  'iptables' => '/sbin/iptables',
  'iptout'   => '/tmp/iptables.out',
  'ipterr'   => '/tmp/iptables.err',
  'debug'    => 0,
  'verbose'  => 0
);
my $ipt_obj = new IPTables::ChainMgr(%opts)
  or die "[*] Could not acquire IPTables::ChainMgr object";


#---------------------- Error Handlers ----------------------#
#$SIG{__DIE__} = 'INT_handler';
$SIG{'INT'} = 'INT_handler';
$SIG{'TERM'} = 'INT_handler';
sub INT_handler {
  print "\nEnding script...\n";
  foreach $IP (keys %pkt) {
	foreach $port (keys %{$pkt{$IP}}) {
	  if ($pkt{$IP}{$port}{block} == 1) {
		print "* Unloading $IP:$port\n";
		iptUpdate($IP, $port, "D");
	  }
	}
  }
  exit(0);
}


#------------ Initialize array with stored data ------------#
if (-e BLOCKFILE) {
  # Read blockfile into memory
  open FH, BLOCKFILE; my @LINES = <FH>; close(FH);

  foreach $line (@LINES) {
	@data = split(/[ \t]+/, $line);
	if (($data[0]) && ($data[0] ne "#IP:Port")) {
	  @tmp = split(/:/, $data[0]);
	  $IP = $tmp[0]; $IP =~ s/^\#//;
	  $port = $tmp[1];
	  $pkt{$IP}{$port}{multi} = $data[1];
		if ($pkt{$IP}{$port}{multi} eq "") { $pkt{$IP}{$port}{multi} = 0; }
	  $pkt{$IP}{$port}{count} = $data[2];
	  $pkt{$IP}{$port}{timer} = $pkt{$IP}{$port}{count};
	  $pkt{$IP}{$port}{start} = $data[3];
	  $pkt{$IP}{$port}{block} = 1;
	  if (substr($data[0], 0, 1) eq "#") {
		$pkt{$IP}{$port}{block} = 2;
	  } else {
		iptUpdate($IP, $port, "I");
	  }
	  print "Loading ($pkt{$IP}{$port}{block}) $IP:$port\n";
	}
  }
} else {
  # Create new blockfile
  open FH, ">", BLOCKFILE;
  print FH "#IP:Port		#Multi	#Count	#Start\n";
  close(FH);
}

#--------------------- Begin main loop ---------------------#
open DUMP, "tcpdump -pnt -i ".NETDEV." not src net ".NETMASK." and not udp src port 53 and udp dst port 53 and '(ip[2:2] != 0)' 2>/dev/null |" or die;
while (<DUMP>) {
  chomp ($line = <DUMP>);

  # Values for {block}
  # 0 - Port traffic detected, but no action taken
  # 1 - IP:port is currently blocked
  # 2 - IP:port block expired, now monitoring for repeat offense

  # Record each packet in array
  @data = split(/ /, $line);
  @tmp = split(/\./, $data[1]);
  $IP = "$tmp[0].$tmp[1].$tmp[2].$tmp[3]";
  $port = $tmp[4];
  $url = ($data[5] eq "[1au]") ? $data[7] : $data[6];
  #@bits = reverse (split(/\./, $url));
  if ($IP) {
	# Add new IP to array
	if (! $pkt{$IP}) {
	  $pkt{$IP}{0}{count} = 0;
	  $pkt{$IP}{0}{multi} = 0;
	}
	$pkt{$IP}{0}{count}++;

	# New port number seen for this IP
	if (! $pkt{$IP}{$port}) {
	  if ($pkt{$IP}{$port}{block} > 0) { print "WARNING! $IP:$port adding existing block\n"; }
	  $pkt{$IP}{$port}{block} = 0;
	  $pkt{$IP}{$port}{count} = 0;
	  $pkt{$IP}{$port}{multi} = 0;
	  $pkt{$IP}{$port}{start} = $tmr;
	}

	# Do unless port is currently blocked
	if ($pkt{$IP}{$port}{block} != 1) {
	  $pkt{$IP}{$port}{count}++;
	  #$pkt{$IP}{$port}{"$bits[1].$bits[0]"}++;
	  if ($pkt{$IP}{$port}{count} > MAXCOUNT) { $pkt{$IP}{$port}{count}=MAXCOUNT; }
	}

	# If IP:port exceeds limits while being monitor, prepare to block again
	if (($pkt{$IP}{$port}{block} == 2) && ($pkt{$IP}{$port}{count} >= THRESH)) { $pkt{$IP}{$port}{block} = 0; }

  }

  # If one second has passed since the last check, examine collected data
  if ($tmr != time()) {
	$tmr = time();
	$h = strftime("%H", localtime);
	$m = strftime("%M", localtime);
	$s = strftime("%S", localtime);
	$hms = "$h:$m:$s";
	print "$hms Count: ", scalar(keys %pkt), "  \r";

	foreach $IP (keys %pkt) {
	  $pkt{$IP}{0}{portblock} = 0;
	  foreach $port (reverse sort keys %{$pkt{$IP}}) {
		if ($port == 0) {
		  if (($pkt{$IP}{0}{block} == 2) && ($pkt{$IP}{0}{portblock} >= MAXBLOCKS)) { $pkt{$IP}{0}{block} = 0; }
		} else {
		  # Count the number of ports blocked for this IP
		  if ($pkt{$IP}{$port}{block} == 1) { $pkt{$IP}{0}{portblock}++; }
		}

		# Add new block
		if ($pkt{$IP}{$port}{block} == 0) {
		  $max = "";
		  if (($port > 0) && ($pkt{$IP}{$port}{count} >= THRESH)) { $max = "port"; }
		  if (($port == 0) && ($pkt{$IP}{0}{count} >= MAXPERSEC)) { $max = "maxpersec"; }
		  if (($port == 0) && ($pkt{$IP}{0}{portblock} >= MAXBLOCKS)) { $max = "maxblock"; }

		  if ($max) {
			$pkt{$IP}{$port}{block} = 1;
			$pkt{$IP}{$port}{start} = $tmr;
			$pkt{$IP}{$port}{timer} = $pkt{$IP}{$port}{count};
			print "$hms [$pkt{$IP}{$port}{multi}] ",$max,"Blocking $IP:$port - $pkt{$IP}{$port}{count} hits\n";
			iptUpdate($IP, $port, "I");
			updateBlockfile($IP, $port, %pkt);
		  }

		# Remove expired block
		} else {
		  $expire = $pkt{$IP}{$port}{start} + ((($pkt{$IP}{$port}{block} * 2) ** $pkt{$IP}{$port}{multi}) * $pkt{$IP}{$port}{timer} * BASETIME);
		  if ($tmr > $expire) {
			$blocked = $tmr - $pkt{$IP}{$port}{start};
			if ($pkt{$IP}{$port}{block} == 2) {
			  print "$hms Deleting $IP:$port after $blocked seconds\n";
			  $pkt{$IP}{$port}{block} = 0;
			  updateBlockfile($IP, $port, %pkt);
			  delete ($pkt{$IP}{$port});

			} else {
			  print "$hms Unblocking $IP:$port after $blocked seconds\n";
			  iptUpdate($IP, $port, "D");
			  $pkt{$IP}{$port}{block} = 2;
			  $pkt{$IP}{$port}{count} = 0;
			  $pkt{$IP}{$port}{start} = $tmr;
			  $pkt{$IP}{$port}{multi}++;
			  if ($pkt{$IP}{$port}{multi}>MAXMULTI) {$pkt{$IP}{$port}{multi}=MAXMULTI;}
			  updateBlockfile($IP, $port, %pkt);
			}
		  }
		}

		# If record still exists, make updates
		if ($pkt{$IP}{$port}) {

		  # Unblocked IPs should decrease their count at a normal rate
		  if ($pkt{$IP}{$port}{block} == 2) {
			$pkt{$IP}{$port}{count}--;
			if ($pkt{$IP}{$port}{count} < 1) { $pkt{$IP}{$port}{count} = 1; }
		  }

		  # Remove array element when no longer needed
		  if (($port > 0) && ($pkt{$IP}{$port}{block} == 0)) {
			$pkt{$IP}{$port}{count}--;
			if ($pkt{$IP}{$port}{count} <= 0) {
			  if ($pkt{$IP}{$port}{timer} > 0) { print "WARNING! Removed $IP:$port from array [$pkt{$IP}{$port}{block}]\n"; }
			  delete ($pkt{$IP}{$port});
			}
		  }

		} #if exists

	  } #foreach PORT

	  # Clean up unused IPs
	  $portcount = scalar(keys %{$pkt{$IP}});
	  if ($portcount <= 1) {
		if ($pkt{$IP}{0}{portblock} > 0) { print "$hms [$pkt{$IP}{0}{portblock}] Cleanup $IP:port\n"; }
		delete ($pkt{$IP});
	  } else {
		$pkt{$IP}{0}{count} = 0;
	  }

	} #foreach IP
  }
}


#----------------------- Subroutines ------------------------#

sub updateBlockfile { # IP, port, %pkt
  my $IP = shift;
  my $port = shift;
  my $pkt = shift;
  my $key = "$IP:$port";
  my $found = 0;

  open FH, BLOCKFILE;
  my @LINES = <FH>;
  close(FH);

  open FH, ">", BLOCKFILE;
  foreach $line (@LINES) {
	switch ($pkt{$IP}{$port}{block}) {
	  case 1 {
		if (($line =~ /^$key/) || ($line =~ /^\#$key/)) {
		  print FH setLen("$IP:$port"), "	$pkt{$IP}{$port}{multi}	$pkt{$IP}{$port}{count}	$pkt{$IP}{$port}{start}\n";
		  $found++;
		} else { print FH $line; }
	  }

	  case 2 {
		if (($line =~ /^$key/) || ($line =~ /^\#$key/)) {
		  print FH setLen("#$IP:$port"), "	$pkt{$IP}{$port}{multi}	$pkt{$IP}{$port}{count}	$pkt{$IP}{$port}{start}\n";
		  $found++;
		} else { print FH $line; }
	  }

	  else {
		print FH $line unless (($line =~ /^$key/) || ($line =~ /^\#$key/));
	  }
	}
  }

  if (($pkt{$IP}{$port}{block} == 1) && (!$found)) {
	print FH setLen("$IP:$port"), "	$pkt{$IP}{$port}{multi}	$pkt{$IP}{$port}{count}	$pkt{$IP}{$port}{start}\n";
  }
  close(FH);
}

sub setLen {
  my $str = $_[0];
  if (length $str < 16) { $str .= "	"; }
  return $str;
}

sub iptUpdate { # IP, port, I/D
  my $ip = $_[0];
  if (substr($ip, 0, 1) eq "#") { $ip = substr($ip, 1); }
  my $port = $_[1];
  my $act = $_[2];

  if ($ip && $act) {
	if ($act eq "D") {
	  if ($port > 0) {
		$ipt_obj->delete_ip_rule($ip, '0.0.0.0/0', 'filter', 'INPUT', 'DROP', {'protocol' => 'udp', 's_port' => $port});
	  } else {
		$ipt_obj->delete_ip_rule($ip, '0.0.0.0/0', 'filter', 'INPUT', 'DROP', {'protocol' => 'udp', 'd_port' => 53});
	  }
	}
	if ($act eq "I") {
	  if ($port > 0) {
		$ipt_obj->add_ip_rule($ip, '0.0.0.0/0', 1, 'filter', 'INPUT', 'DROP', {'protocol' => 'udp', 's_port' => $port});
	  } else {
		$ipt_obj->add_ip_rule($ip, '0.0.0.0/0', 1, 'filter', 'INPUT', 'DROP', {'protocol' => 'udp', 'd_port' => 53});
	  }
	}
  }
}