#!/usr/bin/perl

use strict;

use Symbol qw();

use constant DEBUG => 1;
use constant NEUTER => 1;

use constant INTERVAL => 2;
use constant EMAIL_NOTIFY => 'dharris@drh.net';

use constant UPDATE_LOCAL_IP_INTERVAL => 3600;
use constant IPTABLES_TABLE_NAME => "input";

$|++;

use constant TCP_ESTABLISHED  => 1;
use constant TCP_SYN_SENT     => 2;
use constant TCP_SYN_RECV     => 3;
use constant TCP_FIN_WAIT1    => 4;
use constant TCP_FIN_WAIT2    => 5;
use constant TCP_TIME_WAIT    => 6;
use constant TCP_CLOSE        => 7;
use constant TCP_CLOSE_WAIT   => 8;
use constant TCP_LAST_ACK     => 9;
use constant TCP_LISTEN       => 10;
use constant TCP_CLOSING      => 11;

my $status_lookup = {
	scalar( TCP_ESTABLISHED ) => "ESTABLISHED",
	scalar( TCP_SYN_SENT    ) => "SYN_SENT",
	scalar( TCP_SYN_RECV    ) => "SYN_RECV",
	scalar( TCP_FIN_WAIT1   ) => "FIN_WAIT1",
	scalar( TCP_FIN_WAIT2   ) => "FIN_WAIT2",
	scalar( TCP_TIME_WAIT   ) => "TIME_WAIT",
	scalar( TCP_CLOSE       ) => "CLOSE",
	scalar( TCP_CLOSE_WAIT  ) => "CLOSE_WAIT",
	scalar( TCP_LAST_ACK    ) => "LAST_ACK",
	scalar( TCP_LISTEN      ) => "LISTEN",
	scalar( TCP_CLOSING     ) => "CLOSING",
};

use constant INDEX_NUM_CONN => 0;
use constant INDEX_NUM_SYN  => 1;
use constant INDEX_NETSTAT  => 2;


my @local_ip_list;
my $local_ip_hash;
my $local_ip_last_updated;

## Create configuration (TODO: we should read this from a file in future versions, and re-read with a HUP signal)

my @defaults = (
	syn_weight => 0.5,
	syn_contrib => 0.4,
	kill => 1,
	bantime => 3600*2,
	cutoff_block_factor => 0.25,
);

my $config = {
	"apache1" => {
		ports => [ 80, 443, 1001 ],
		cutoff => 40,
		@defaults,
	},
	"apache9" => {
		ports => [ 911, 912, 913, 914, 1009 ],
		cutoff => 40,
		@defaults,
	},
	"imapd" => {
		ports => [ 143 ],
		cutoff => 20,
		@defaults,
		syn_contrib => 0,
	},
	"ipop3d" => {
		ports => [ 110 ],
		cutoff => 20,
		@defaults,
		syn_contrib => 0,
	},
	"smtpd" => {
		ports => [ 25 ],
		cutoff => 20,
		@defaults,
		syn_contrib => 0,
	},
	"smtpd" => {
		ports => [ 225 ],
		cutoff => 20,
		@defaults,
		syn_contrib => 0,
	},
};

## Process configuration

my $class_by_port = {};

foreach my $class_name ( keys %$config ) {
	my $class = $config->{$class_name};
	foreach my $attribute ( qw( ports cutoff syn_weight syn_contrib ) ) {
		die "class ($class) missing required attribute ($attribute)"
			if ( not defined $class->{$attribute} );
	}
	$class->{class_name} = $class_name;
	foreach my $port ( @{$class->{ports}} ) {
		die "duplicate port ($port) in two classes ($class_name) ($class_by_port->{$port}{class_name})"
			if ( exists $class_by_port->{$port} );
		$class_by_port->{$port} = $class;
	}
}

&main;

sub convert_hex_ip
{
	my $hex_ip = shift;
	$hex_ip =~ /^([0-9a-fA-F]{2})([0-9a-fA-F]{2})([0-9a-fA-F]{2})([0-9a-fA-F]{2})$/
		or die "invalid hex ip ($hex_ip)";
	return join(".", map { hex($_) } ( $4, $3, $2, $1 ) );
}

sub get_class_c_block
{
	my $ip = shift;
	$ip =~ /^(\d+\.\d+\.\d+)\.\d+$/ or die "unable to parse ip ($ip)";
	return( $1 . ".0/24" );
}

sub exceedes_dos_cutoff
{
	my ($ip_counts, $class, $cutoff_scale_factor) = @_;

	my $aa = $ip_counts->[INDEX_NUM_SYN] * $class->{syn_weight};
	my $syn_max = $class->{syn_contrib} * $class->{cutoff};
	my $bb = $ip_counts->[INDEX_NUM_CONN] + ($aa > $syn_max ? $syn_max : $aa);

	return ( $bb >= ( defined($cutoff_scale_factor) ? $class->{cutoff} * $cutoff_scale_factor : $class->{cutoff} ) );
}

sub find_dos_attacks
{
	my $fh = Symbol::gensym();

	update_local_ip_list(UPDATE_LOCAL_IP_INTERVAL);

#	open($fh, "< /proc/net/tcp") or die "unable to open file (/proc/net/tcp): $!";
open($fh, "< proc-net-tcp") or die "unable to open file (/proc/net/tcp): $!";

	my $netstat_header = 
	"Proto Recv-Q Send-Q Local Address           Foreign Address         State       Inode      D";

	my $candidate = {};

	while( <$fh> ) {
		if ( /^\s+sl\s+local_address/ ) {
			; # ignore header line
		}
		elsif ( m/^
			\s*\d+:\s+
			([0-9A-Fa-f]{8}):([0-9A-Fa-f]{4})\s+
			([0-9A-Fa-f]{8}):([0-9A-Fa-f]{4})\s+
			([0-9A-Fa-f]{2})\s+
			([0-9A-Fa-f]{8}):([0-9A-Fa-f]{8})\s+
			[0-9A-Fa-f]{2}:[0-9A-Fa-f]{8}\s+
			[0-9A-Fa-f]{8}\s+
			\d+\s+
			\d+\s+
			(\d+)\s+
			.*$
		/x )
		{
			my ($local_ip, $local_port, $remote_ip, $remote_port, $status) =
				(convert_hex_ip($1), hex($2), convert_hex_ip($3), hex($4), hex($5));
			my ($tx_queue, $rx_queue, $inode) = (hex($6), hex($7), $8);

			print sprintf("%-25s -> %-25s %-15s %-5s\n",
				"$local_ip:$local_port", "$remote_ip:$remote_port", $status_lookup->{$status}, $inode
			) if ( DEBUG );

			## Add this connection to the pool of potential bad things.

			if (
				exists($class_by_port->{$local_port}) &&
				! exists($local_ip_hash->{$remote_ip}) &&
				$remote_ip ne "0.0.0.0" &&
				$local_ip ne "0.0.0.0"
			)
			{
				my $class_name = $class_by_port->{$local_port}{class_name};

				my $tally_slot;

				if (
					$status == TCP_ESTABLISHED
					|| ( $inode != 0 && (
						   $status == TCP_LAST_ACK
						|| $status == TCP_CLOSE_WAIT
						|| $status == TCP_TIME_WAIT
						|| $status == TCP_CLOSE
						|| $status == TCP_CLOSING
						|| $status == TCP_FIN_WAIT1
						|| $status == TCP_FIN_WAIT2
					) )
				) {
					$tally_slot = 0;
				}

				if ( $status == TCP_SYN_RECV ) {
					$tally_slot = 1;
				}

				my $block = get_class_c_block($remote_ip);

				my $netstat_line =
				sprintf("%-5s %6d %6d %-23s %-23s %-12s\n %-10u %-1s\n",
					"tcp",
					$rx_queue,
					$tx_queue,
					$local_ip  . ":" . ( $local_port  == 0 ? "*" : $local_port  ),
					$remote_ip . ":" . ( $remote_port == 0 ? "*" : $remote_port ),
					exists($status_lookup->{$status}) ? $status_lookup->{$status} : $status,
					$inode,
					( defined($tally_slot) ? $tally_slot : "-" )
				);

				$candidate->{$class_name}{$block}{$remote_ip}[$tally_slot]++ if ( defined $tally_slot );
				$candidate->{$class_name}{$block}{$remote_ip}[2] .= $netstat_line;
			};
		}
		else {
			warn "malformed line from /proc/net/tcp: $_";
		}
	}

	## Close the file

	close($fh) or die "error closing /proc/net/tcp file: ($!)";

	## For each class, find the dos attacks

	my @dos_attacks;

	foreach my $class_name ( keys %$candidate )
	{
		my $ip_block_hash = $candidate->{$class_name};
		my $class = $config->{$class_name};

		## Sort through the blocks

		while ( my ($ip_block_name, $ip_hash) = each(%$ip_block_hash) ) {

			## Sort through each IP, counting and detecting DOS IPs

			my @total_counts = (0, 0);
			my @total_counts_non_dos_ips = (0, 0);

			my @candidate_dos_ips;

			while ( my ($ip, $ip_counts) = each %$ip_hash ) {

				$total_counts[INDEX_NUM_CONN] += $ip_counts->[INDEX_NUM_CONN];
				$total_counts[INDEX_NUM_SYN]  += $ip_counts->[INDEX_NUM_SYN];

				if ( exceedes_dos_cutoff($ip_counts, $class) ) {
					push @candidate_dos_ips, {
						type    => "ip",
						class   => $class_name,
						ip      => $ip,
						ip_list => [ $ip ],
						netstat => $netstat_header . $ip_counts->[INDEX_NETSTAT],
					};
print "candidate DOS ip: $ip\n";
				} else {
					$total_counts_non_dos_ips[INDEX_NUM_CONN] += $ip_counts->[INDEX_NUM_CONN];
					$total_counts_non_dos_ips[INDEX_NUM_SYN]  += $ip_counts->[INDEX_NUM_SYN];
				}
			}

			## Choose between a DOS IP and DOS block, if detected

			# NOTE: We accept a DOS assessment on a block if:
			#
			#   (a) the block has enough connections to qualify as DOS
			#   &&
			#     (b) there are no individual IP DOS attacks detected inside the block
			#     ||
			#     (c) the connectinos not included in the individual IP DOS attacks
			#         ($total_count_non_dos_ips), execeedes $class->{cutoff_block_factor}
			#         of the way to a DOS, which allows us to assume the DOS is coming
			#         from the whole block

			if (
				exceedes_dos_cutoff(\@total_counts, $class)
				&& (
					( ! @candidate_dos_ips ) ||
					exceedes_dos_cutoff(\@total_counts_non_dos_ips, $class, $class->{cutoff_block_factor})
				)
			)
			{
				## Get a list of the netcat lines for all ips in thes block,
				## sorted by rightmost IP quad

				my $netstat_for_block = join("",
					map { $ip_hash->{$_->[0]}[INDEX_NETSTAT] }
					sort { $a->[1] <=> $b->[1] }
					map { [ $_, (split(".",$_))[3] ] }
					values %$ip_hash
				);

				## Add the attack

				push @dos_attacks, {
					type     => "block",
					class    => $class_name,
					ip       => $ip_block_name,
					ip_list  => [ $ip_block_name ],
					netstat  => $netstat_header . $netstat_for_block
				};
			}
			else {
				# this list of @candidate_dos_ips may be blank, but we at least know
				# there is no block DOS that should be added instead, so we add it it 
				# onto the list without reservation
				push @dos_attacks, @candidate_dos_ips;
			}
		}
	}

	## Do a second check to prevent any blocking of local addresses, this time with guranteed
	## fresh local_ip_list information

	if ( @dos_attacks )
	{
		update_local_ip_list(0);

		my @dos_attacks_filtered;

		foreach my $attack ( @dos_attacks ) {
			push @dos_attacks_filtered, $attack
				if ( ! grep { is_local_ip($_) } @{$attack->{ip_list}} );
		}

		@dos_attacks = @dos_attacks_filtered;
	}

	## Return the data we found

	print "dos_attacks = " . join(" ", map { "$_->{class},$_->{type},$_->{ip}" } @dos_attacks) . "\n" if ( DEBUG );

	return \@dos_attacks;
}

sub check_for_ip_in_table
{
	my $ip = shift;
	my $table = shift;

	my $fh = Symbol::gensym();

	open($fh, "iptables -n -L $table |") or die "unable to open pipe from iptables: ($!)";

	my $found = 0;

	while ( <$fh> ) {
		$found = 1 if ( grep { $_ eq $ip } split(/\s+/, $_) );
	}

	close($fh) or die "error closing pipe from iptables: ($!)";

	return $found;
}

sub match_ip
{
	my $ip_or_block = shift;
	my $ip = shift;

	die "match_ip error: second argument should not be a block" if ( $ip =~ m|/| );

	if ( $ip_or_block =~ m|/(\d+)$| ) {
		my $block_size = $1;
		if ( $block_size == 24 ) {
			return $ip_or_block eq get_class_c_block($ip);
		} else {
			die "unable to handle block size ($block_size) in block ($ip_or_block)";
		}
	} else {
		return $ip_or_block eq $ip;
	}
}

sub kill_connections
{
	my $ip_list = shift;
	my $port_list = shift;

	my $lsof_ip_config = join(" ", map { "-i TCP\@$_" } @$ip_list );

	my $fh = Symbol::gensym();
	open($fh, "lsof -n -P $lsof_ip_config |") or die "unable to open pipe from lsof: $!";

	my $header_line = "";
	my $connection_lines = "";
	my @pids;

	while (<$fh>)
	{
		if ( /^\s*COMMAND\s+PID/ ) {
			$header_line = $_;
		}
		elsif ( /TCP\s+\d+\.\d+\.\d+\.\d+\:(\d+)->(\d+\.\d+\.\d+\.\d+):\d+/ )
		{
			my ($local_port, $remote_ip) = shift;
			if ( 
				scalar( grep { $_ == $local_port } @$port_list ) &&
				scalar( grep { match_ip($_, $remote_ip) } @$ip_list )
			) {
				$connection_lines .= $_;
				/^\s*\S+\s+(\d+)\s/ or die "unable to parse out pid";
				push @pids, $1;
			}
		}
		else {
			warn "bad line from lsof: $_";
		}
	}

	close($fh) or die "error closing lsof pipe";

	if ( ! NEUTER ) {
		kill 15, @pids if ( @pids );
	}

	return $header_line . ( $connection_lines eq "" ? "no processes found\n" : $connection_lines );
}

sub update_local_ip_list
{
	my $max_old = shift;

	if (
		   ! defined($local_ip_last_updated)
		|| ! defined($max_old)
		|| $max_old == 0
		|| time() - $local_ip_last_updated > $max_old
	)
	{
		my $fh = Symbol::gensym();
		open($fh, "ifconfig |") or die "unable to open a pipe from ifconfig: $!";

		$local_ip_hash = {};
		@local_ip_list = ();
		$local_ip_last_updated = time();

		foreach (<$fh>) {
			if ( /inet\s*addr\s*:\s*(\d+\.\d+\.\d+\.\d+)(?:\s|$)/ ) {
				my $local_ip = $1;
				$local_ip_hash->{$1} = 1;
				push(@local_ip_list, $1);
			}
		}

		close($fh) or die "error closing pipe from ifconfig";

		die "no local ip addresses found: perhaps ifconfig changed the output format?"
			if ( ! @local_ip_list )
	}

}

sub is_local_ip
{
	my $ip_or_block = shift;

	return scalar ( grep { match_ip($ip_or_block, $_) } @local_ip_list );
}

sub send_email
{
	my $to = shift;
	my $subject = shift;
	my $body = shift;

	my $fh = Symbol::gensym();

	open($fh, "|qmail-inject") or die "error opening pipe to qmail-inject: ($!)";

	print $fh "To: $to\n";
	print $fh "Subject: $subject\n";
	print $fh "\n";
	print $fh $body;

	close($fh) or die "error closing pipe to qmail-inject: ($!)";
}

sub print_and_capture
{
	my $capture_ref = shift;
	my $print_str = shift;

	$$capture_ref .= $print_str;
	print $print_str;
}

sub perform_check
{
	my $dos_attacks = find_dos_attacks();

	foreach my $attack ( @$dos_attacks )
	{
		my $type    = $attack->{type};
		my $class   = $attack->{class};
		my $ip      = $attack->{ip};
		my $ip_list = $attack->{ip_list};
		my $netstat = $attack->{netstat};

		my $port_list = $config->{$class}{ports};

		my $port_list_str = join(", ", sort @$port_list);
		my $ip_list_str   = join(", ", sort @$ip_list);

		my $hostname = `hostname`;
		chomp($hostname);

		my $capture = "";

		print "============================================================\n";

		print_and_capture(\$capture, <<EOT);

Detected and blocked DOS attack!

         from ip: $ip_list_str
  type of attack: $type
service attacked: $class (ports $port_list_str)
            time: @{[scalar(localtime())]}
   host attacked: $hostname

EOT

		print_and_capture(\$capture, "netstat reports:\n$netstat\n");

		if ( check_for_ip_in_table($ip, IPTABLES_TABLE_NAME) )
		{
			print_and_capture(\$capture, "already blocked in iptables\n");
		}
		else
		{
			if ( ! NEUTER ) {
				print_and_capture(\$capture, "blocking in iptables...\n");
				system("iptables", "-A", "INPUT", "-s", $ip, "-j", "DROP") == 0 or die "error running iptables";
				print_and_capture(\$capture, "blocked\n");
			}
			else {
				print_and_capture(\$capture, "this instance is neutered, so refusing to block in iptables\n");
			}
		}

		print_and_capture(\$capture, "killing connections...\n");
		my $connections_killed = kill_connections($ip_list, $port_list);
		print_and_capture(\$capture, "done\n\n");

		print_and_capture(\$capture, "killed these connections:\n$connections_killed\n");

		if ( EMAIL_NOTIFY ne "" ) {
			print_and_capture(\$capture, "\nsending e-mail notification to @{[EMAIL_NOTIFY]}\n");
			send_email(EMAIL_NOTIFY, "Blocked DOS attack from $ip", $capture . "\n");
			print_and_capture(\$capture, "sent\n");
		}
		print_and_capture(\$capture, "\n\n");
	}
}

sub main
{
	update_local_ip_list();

#	while (1) {
#		print "checking....\n\n";
#		perform_check();
#		sleep(INTERVAL);
#	}
	find_dos_attacks();
}