#!/usr/bin/perl -w
use strict;
use File::Basename;
use POSIX 'strftime';

my $VERSION = "1.04, 2010-03-04";

=head1
File /etc/sysconfig/vz-scripts/4.conf, FIREWALL directive
---------------------------------------------------------
FIREWALL="
	host.allowed.to.every.port
	yet.another.host
	* # means "any host"

	[25]
	host.allowed.to.access.smtp
	* # means "any"

	[80,443]
	hosts.allowed.to.access.two.ports
	
	[udp:53]
	*

	[CUSTOM]
	# You may use "$THIS" macro which is replaced by this machine IP
	# (and, if the machine has many IPs, it will be multiplicated).
	-A INPUT -i eth2 -d $THIS -j ACCEPT
	# Or you may use commands with no references to $THIS (only
	# such commands are allowed for 0.conf file).
	-A INPUT -i eth1 -j ACCEPT
"
----------------------------------------------------------------
We use FIREWALL directive in plain VE configs to allow to
vzmigrate it easily from one node to another.
=cut

my $IPTABLES = "/etc/sysconfig/iptables";
my ($DIR) = grep { -d $_ } ("/etc/sysconfig/vz-scripts", "/etc/vz/conf");

if (@ARGV && $ARGV[0] eq "-a") {
	do_apply(0, $ARGV[1] && $ARGV[1] eq "-f");
} elsif (@ARGV && $ARGV[0] eq "-t") {
	do_apply(1);
} else {
	die 
		"dkLab vzfirewall: simple rules for openvz firewall.\n" .
		"Version: $VERSION\n" .
		"Homepage: http://en.dklab.ru/lib/dklab_vzfirewall\n" .
		"Usage:\n" .
		"  - Apply rules in $DIR/*.conf (FIREWALL directives):\n" .
		"    $0 -a [-f]\n" .
		"  - Preview rules in $DIR/*.conf without activation:\n" .
		"    $0 -t\n";
}

sub do_apply {
	my ($test_mode, $force) = @_;
	
	my @cmds = ();
	push @cmds, "##\n## PLEASE DO NOT EDIT THIS FILE MANUALLY!!!\n##\n";
	push @cmds, "## It is generated by " . basename($0) . "\n";
	push @cmds, "## All changes will be lost on re-generation!\n##\n";
	push @cmds, "*filter\n";
	push @cmds, ":INPUT ACCEPT [0:0]\n"; # MUST be ACCEPT by default for safety (see footer)
	push @cmds, ":FORWARD ACCEPT [0:0]\n";
	push @cmds, ":OUTPUT ACCEPT [0:0]\n";
	push @cmds, "\n\n";
	push @cmds, "##\n## Default opened channels.\n##\n";
	push @cmds, "-A INPUT -i lo -j ACCEPT\n";
	push @cmds, "-A OUTPUT -o lo -j ACCEPT\n";
	push @cmds, "-A FORWARD -i lo -j ACCEPT\n";
	for ("INPUT", "OUTPUT", "FORWARD") {
		push @cmds, "-A $_ -p icmp -j ACCEPT\n";
		push @cmds, "-A $_ -m state --state ESTABLISHED,RELATED -j ACCEPT\n";
	}
	push @cmds, "# Open SSH port on hardware node - for safety.\n";
	push @cmds, "-A INPUT -p tcp --dport 22 -j ACCEPT\n";
	push @cmds, "\n\n";
	
	# Collect all data.
	my @parsed = ();
	foreach my $conf (glob($DIR . "/*.conf")) {
		my $basename = basename($conf);
		my $opts = read_conf($conf);
		my ($rules, $custom) = read_rules($opts->{FIREWALL});
		my @dst_ips;
		if ($basename ne "0.conf") {
			my $ips = $opts->{IP_ADDRESS} or die "Cannot find IP_ADDRESS in $conf\n";
			$ips =~ s/^\s+|\s+$//sg;
			@dst_ips = split /\s+/, $ips;
		} else {
			# "" meands "host node IPs" (everything)
			@dst_ips = ("");
		}
		push @parsed, [ $basename, \@dst_ips, $rules, $custom ];
	}
	
	# Generate OPEN commands.
	foreach (@parsed) {
		my ($basename, $dst_ips, $rules, $custom) = @$_;
		next if !@$rules && !@$custom;
		push @cmds, "##\n## Rules from $basename\n##\n";
		foreach my $dst_ip (@$dst_ips) {
			foreach my $pair (@$rules) {
				push @cmds, generate_open_rule(
					$pair->[1], # src ip
					$dst_ip,    # dst ip
					$pair->[0], # dst port(s)
					$pair->[2], # comment
				);
			}
			my @ccmds = ();
			foreach my $rule_c (@$custom) {
				my $rule = $rule_c;
				$rule =~ s/\$THIS/$dst_ip/sg or next;
				die "You cannot use \$THIS in $basename: it is for VEs only\n" if !$dst_ip;
				push @ccmds, $rule . "\n";
			}
			if (@ccmds) {
				push @cmds, "# CUSTOM for $dst_ip ($basename):\n";
				push @cmds, @ccmds;
				push @cmds, "\n";
			}
		}
		my @ccmds = ();
		foreach my $rule (@$custom) {
			next if $rule =~ /\$THIS/s;
			push @ccmds, $rule . "\n";
		}
		if (@ccmds) {
			push @cmds, "# CUSTOM for VE $basename:\n";
			push @cmds, @ccmds;
		}
		push @cmds, "\n\n";
	}
	
	# Generate CLOSE rules for all destination IPs.
	push @cmds, "##\n## All other access to these IPs is closed.\n##\n";
	foreach (@parsed) {
		my ($basename, $dst_ips, $rules, $custom) = @$_;
		push @cmds, map { generate_close_rule($_) } @$dst_ips;
	}

	# Generate OPEN rules for outgoing connections.
	push @cmds, "\n\n##\n## Outgoing connections are permitted.\n##\n";
	foreach (@parsed) {
		my ($basename, $dst_ips, $rules, $custom) = @$_;
		push @cmds, map { generate_outgoing_rule($_) } @$dst_ips;
	}

	push @cmds, "\n\n##\n## Default action for incoming packets - reject.\n##\n";
	push @cmds, "-A INPUT -j DROP\n";
	push @cmds, "-A FORWARD -j DROP\n";
	push @cmds, "COMMIT\n";
	my $cmds = join "", @cmds;
	
	my $prev = "";
	open(local *F, ">>", $IPTABLES) and close(F);
	if (open(local *F, $IPTABLES)) {
		local $/;
		$prev = <F>;
		close(F);
	}

	if ($prev ne $cmds) {
		print STDERR "--DIFF--\n";
		open(local *P, "| diff $IPTABLES - >&2");
		print P $cmds;
		close(P);
		print STDERR "--RULES--\n" if $test_mode;
	}
	
	if ($test_mode) {
		print $cmds;
	} else {
		if ($force || $prev ne $cmds) {
			print STDERR "Testing new rules...\n";
			open(local *P, "| iptables-restore -t");
			print P $cmds;
			if (!close(P)) {
				die "Test failed. Apply nothing.\n";
			}
			rename($IPTABLES, $IPTABLES . ".bak." . strftime("%Y-%m-%d_%H-%M-%S", localtime(time)));
			open(local *F, ">", $IPTABLES) or die "Cannot create $IPTABLES: $!\n";
			print F $cmds;
			close(F);
			print STDERR "Applying new rules...\n";
			system("iptables-restore < $IPTABLES");
		} else {
			die "Nothing is changed.\n";
		}
	}
}

sub read_rules {
	my ($text) = @_;
	return ([], []) if !$text;
	#return ([[ "*", "*", "ANYTHING" ]], []) if !$text;
	my @rules = ();
	my @custom = ();
	my $port = "*";
	foreach (split /\n/, $text) {
		s/^\s+|[#;].*|\s+$//sg;
		next if !$_;
		if (/^\s*\[(.*)\]\s*$/s) {
			$port = $1;
			$port =~ s/\s+//sg;
			next;
		}
		if (lc $port eq "custom") {
			push @custom, $_;
			next;
		}
		my $ips = resolve($_);
		foreach my $ip (@$ips) {
			push @rules, [ $port, $ip, $_ ];
		}
	}
	return (\@rules, \@custom);
}

sub generate_open_rule {
	my ($src_ip, $dst_ip, $dst_port, $comment) = @_;
	my @lines = ("# $comment -> $dst_ip:$dst_port\n");
	my $proto = "tcp";
	if ($dst_port =~ /^([a-z]+):(.*)$/is) {
		$proto = $1;
		$dst_port = $2;
	}
	my $rule = 
		($src_ip && $src_ip ne "*"? " -s $src_ip" : "") . 
		($dst_ip && $dst_ip ne "*"? " -d $dst_ip" : "") . 
		($dst_port ne "*"? " -m multiport -p $proto --dports $dst_port" : "") .
		" -j ACCEPT";
	$rule =~ s/\s+/ /sg;
	my $chain = $dst_ip? "FORWARD" : "INPUT";
	push @lines, "-A $chain" . $rule . "\n";
	return join "", @lines;
}

sub generate_close_rule {
	my ($dst_ip) = @_;
	return ($dst_ip && $dst_ip ne "*"? "-A FORWARD -d $dst_ip" : "-A INPUT") . " -j DROP\n";
}

sub generate_outgoing_rule {
	my ($src_ip) = @_;
	return ($src_ip && $src_ip ne "*"? "-A FORWARD -s $src_ip" : "-A OUTPUT") . " -j ACCEPT\n";
}

sub read_conf {
	my ($conf) = @_;
	open(local *F, $conf) or die "Cannot open $conf: $!\n";
	local $/;
	$_ = <F>;
	close(F);
	my %opts = ();
	my @matches = m/^\s* (\w+) \s* = \s* " ([^\"]*) "/mxg;
	for (my $i = 0; $i < @matches; $i += 2) {
		$opts{$matches[$i]} = $matches[$i + 1];
	}
	return \%opts;
}

my %resolved = ();
sub resolve {
	my ($host) = @_;
	return [ $host ] if $host =~ /^\d+\.\d+\.\d+\.\d+|^\*$/s;
	return $resolved{$host} if $resolved{$host};
	my @ips = (`host $host.`||"") =~ /(\d+\.\d+\.\d+\.\d+$)/mg or die "Cannot resolve $host\n";
	foreach (@ips) {
		die "Invalid address $_ resolved for $host\n" if /^(128|255)/s;
	}
	return $resolved{$host} = \@ips;
}
