#!/usr/bin/perl -w
# Copyright (C) 2004, by Tomas Janousek <tomi@nomi.cz>
# Under GNU GPL
use strict;
use Data::Dumper;

# static conf
my $ip = "/sbin/ip";
my $iptables = "/sbin/iptables";
my $tc = "/sbin/tc";
my $rmmod = "/sbin/rmmod";
my $modprobe = "/sbin/modprobe";

# loaded
my $me;
my $failover;
my @czf_ifaces = ();
my $net_iface = "";
my @classes;
my %ranges;
my %protogroups;
my %protorules;
my $debug = 0;
my $verbose = 0;

# generated
my %classids;
my %gclassids;
my %gfiltprefs;
my %gprios;
my %devmarks;
my ($inroot, $outroot);
my $lastmark = 0xffffffff;
my %classdescs;


# utility funcs

our $dev;
sub fordev(@) {
    my $sub = pop @_;
    for my $cdev (@_) {
	local $dev = $cdev;
	$sub->();
    }
}

sub ex($) {
    my $cmd = shift;
    print STDERR "+ $cmd\n";
    unless ($debug) { return system($cmd); }
}

sub linkup() { ex("$ip link set $dev up"); }

# ipt(cmd)
sub ipt($) { ex("$iptables ".shift); }

# iptmangle(cmd)
sub iptmangle($) { ipt("-t mangle ".shift); }

# tc(cmd)
sub tc($) { ex("$tc ".shift); }

sub qddelroot() { tc("qdisc del dev $dev root"); }

# qdhtbadd(handle, [default])
sub qdhtbadd($@) {
    my ($qdisc, $def) = @_;
    my $cmd = "qdisc add dev $dev root handle $qdisc:0 htb";
    if (defined($def)) { $cmd .= " default $def"; }
    tc($cmd);
}

# clhtbadd(qdisc, parent, classid, rate, [ceil], [prio])
sub clhtbadd($$$$@) {
    my ($qdisc, $parent, $classid, $rate, $ceil, $prio) = @_;
    my $cmd = "class add dev $dev parent $qdisc:$parent".
    " classid $qdisc:$classid htb rate ${rate}kbit";
    if (defined($ceil)) { $cmd .= " ceil ${ceil}kbit"; }
    if (defined($prio)) { $cmd .= " prio $prio"; }
    tc($cmd);
}

# fltfwadd(qdisc, parent, classid, handle)
sub fltfwadd($$$$) {
    my ($qdisc, $parent, $classid, $handle) = @_;
    my $pref = fltalloc();
    tc("filter add dev $dev pref $pref parent $qdisc:$parent protocol ip ".
	"handle $handle fw classid $qdisc:$classid");
}

# fltu32ad(qdisc, parent, classid, matches)
sub fltu32ad($$$$) {
    my ($qdisc, $parent, $classid, $matches) = @_;
    my $pref = fltalloc();
    tc("filter add dev $dev pref $pref parent $qdisc:$parent protocol ip ".
	"u32 $matches classid $qdisc:$classid");
}

# clalloc([dev], [qdisc])
sub clalloc(@) {
    my ($cdev, $qdisc) = @_;
    unless (defined($cdev)) { $cdev = $dev; }
    unless (defined($qdisc)) { $qdisc = "1"; }
    return sprintf("%x", ++$gclassids{$cdev}->{$qdisc});
}

sub markalloc() { return $lastmark--; }

# fltalloc([dev])
sub fltalloc(@) {
    my ($cdev) = @_;
    unless (defined($cdev)) { $cdev = $dev; }
    return ++$gfiltprefs{$cdev};
}

# prioalloc(dev, qdisc, classid)
sub prioalloc($$$) {
    my ($cdev, $qdisc, $classid) = @_;
    return sprintf("%x", ++$gprios{$cdev}->{$qdisc}->{$classid} - 1);
}

sub cmt(@) {
    my $cmd = shift;
    if ($cmd) {
	print STDERR " -- $cmd\n";
    } else {
	print STDERR "\n";
    }
}

sub perr(@) {
    my $cmd = shift;
    if ($cmd) {
	print STDERR " !! $cmd\n";
    }
}


# shaper funcs

sub protochain($$$$$);
sub class($);

sub prepare() {
    cmt("imq init");
    ex("$modprobe imq numdevs=3");
    fordev qw/imq0 imq1 imq2/, sub {
	linkup();
	qddelroot();
    };
    cmt();

    cmt("ipt init");
    fordev $net_iface->{dev}, sub {
	my $mark = markalloc();
	$devmarks{$dev} = $mark;
	iptmangle("-A INPUT -i $dev -j CONNMARK --set-mark $mark");
	iptmangle("-A OUTPUT -o $dev -j MARK --set-mark $mark");
    };
    for my $proto (keys %protorules) {
	for my $rule (@{$protorules{$proto}}) {
	    if ($rule->{used} && $rule->{type} eq "ipt") {
		my $mark = markalloc();
		iptmangle("-A PREROUTING $rule->{param} -j CONNMARK --set-mark $mark");
		iptmangle("-A POSTROUTING $rule->{param} -j CONNMARK --set-mark $mark");
		iptmangle("-A OUTPUT $rule->{param} -j CONNMARK --set-mark $mark");
		$rule->{param} = $mark;
		$rule->{type} = "mark";
	    }
	}
    }
    iptmangle("-A PREROUTING -j CONNMARK --restore-mark");
    iptmangle("-A POSTROUTING -j CONNMARK --restore-mark");
    fordev eval{ map { $_->{dev} } @czf_ifaces }, sub {
	my $mark = markalloc();
	$devmarks{$dev} = $mark;
	iptmangle("-A POSTROUTING -s ! 10.0.0.0/8 -o $dev -j IMQ --todev 0");
	iptmangle("-A PREROUTING -d ! 10.0.0.0/8 -i $dev -j IMQ --todev 1");
	iptmangle("-A PREROUTING -d 10.0.0.0/8 -i $dev -j MARK --set-mark $mark");
    };
    fordev $net_iface->{dev}, eval{ map { $_->{dev} } @czf_ifaces }, sub {
    	iptmangle("-A PREROUTING -m mark --mark $devmarks{$dev} -j IMQ --todev 2");
    };
    cmt();

    cmt("net qdisc init");
    fordev "imq0", sub {
	$inroot = clalloc();
	my $infail = clalloc();
	qdhtbadd(1, $infail);
	clhtbadd(1, 0, $inroot, $net_iface->{in});
	clhtbadd(1, $inroot, $infail, $failover->{in}, undef, 255);
    };
    fordev "imq1", sub {
	$outroot = clalloc();
	my $outfail = clalloc();
	qdhtbadd(1, $outfail);
	clhtbadd(1, 0, $outroot, $net_iface->{out});
	clhtbadd(1, $outroot, $outfail, $failover->{out}, undef, 255);
    };
    cmt();

    cmt("local limit");
    fordev "imq2", sub {
	qdhtbadd(1, 0);
	my $localin = clalloc();
	clhtbadd(1, 0, $localin, $me->{in});
	fltfwadd(1, 0, $localin, $devmarks{$net_iface->{dev}});

	protochain(1, $localin, $me->{protogroup}, $me->{in}, undef);
    };
    fordev $net_iface->{dev}, sub {
	qddelroot();
	qdhtbadd(1, 0);
	my $localout = clalloc();
	clhtbadd(1, 0, $localout, $me->{out});
	fltfwadd(1, 0, $localout, $devmarks{$net_iface->{dev}});

	protochain(1, $localout, $me->{protogroup}, $me->{out}, undef);
    };
    cmt();

    for (@czf_ifaces) {
	cmt("czf iface $_->{dev}");
	my $root;
	fordev $_->{dev}, sub {
	    qddelroot();
	    qdhtbadd(1, 0);
	    $root = clalloc();
	    clhtbadd(1, 0, $root, $_->{out});
	    fltu32ad(1, 0, $root, "match ip src 10.0.0.0/8");

	    protochain(1, $root, $_->{protogroup}, $_->{out}, undef);
	};

	# ingress
	fordev "imq2", sub {
	    my $in = clalloc();
	    clhtbadd(1, 0, $in, $_->{in});
	    fltfwadd(1, 0, $in, $devmarks{$_->{dev}});

	    protochain(1, $in, $_->{protogroup}, $_->{in}, undef);
	};
	cmt();
    }
}

sub class($) {
    my $c = shift;

    return if (defined($classids{$c->{name}}));

    my $inparent = $inroot;
    my $outparent = $outroot;
    if ($c->{parent}) {
	unless (defined($classids{$c->{parent}})) {
	    for (@classes) {
		if ($_->{name} eq $c->{parent}) {
		    class($_);
		}
	    }
	}
	$inparent = $classids{$c->{parent}}->{"imq0"};
	$outparent = $classids{$c->{parent}}->{"imq1"};
    }
    cmt("class $c->{name}");

    my $inid;
    my $outid;

    fordev "imq0", sub {
	$inid = clalloc();
	clhtbadd(1, $inparent, $inid, $c->{irate}, $c->{iceil}, $c->{prio});
	push @{$classdescs{"$dev: 1:$inid"}}, "IN: class $c->{name};  (parent 1:$inparent)";
    };
    fordev "imq1", sub {
	$outid = clalloc();
	clhtbadd(1, $outparent, $outid, $c->{orate}, $c->{oceil}, $c->{prio});
	push @{$classdescs{"$dev: 1:$outid"}}, "OUT: class $c->{name};  (parent 1:$outparent)";
    };

    $classids{$c->{name}} = { "imq0" => $inid, "imq1" => $outid };

    if ($ranges{$c->{name}}) {
	fordev "imq0", sub {
	    protochain(1, $inid, $c->{protogroup}, $c->{irate}, $c->{iceil});
	};
	fordev "imq1", sub {
	    protochain(1, $outid, $c->{protogroup}, $c->{orate}, $c->{oceil});
	};
    }

    for (@{$ranges{$c->{name}}}) {
	fordev "imq0", sub {
	    fltu32ad(1, 0, $inid, "match ip dst $_");
	    push @{$classdescs{"$dev: 1:$inid"}}, $_;
	};
	fordev "imq1", sub {
	    fltu32ad(1, 0, $outid, "match ip src $_");
	    push @{$classdescs{"$dev: 1:$outid"}}, $_;
	};
    }

    cmt();
}

# protochain(qdisc, parent, protogroup, rate, ceil)
sub protochain($$$$$) {
    my ($qdisc, $parent, $protogroup, $rate, $ceil) = @_;

    my $var = {
	qdisc => $qdisc,
	parent => $parent,
	protogroup => $protogroup,
	rate => $rate,
	ceil => $ceil,
	other => undef,
    };

    for my $proto (@{$protogroups{$protogroup}}) {
	my $clsid = clalloc($dev, $qdisc); $var->{clsid} = $clsid;
	my $prate = int($rate * $proto->{prate} / 100.0);
	my $pceil = defined($ceil) ? int($ceil * $proto->{pceil} / 100.0) : undef;

	clhtbadd($qdisc, $parent, $clsid, $prate, $pceil, prioalloc($dev, $qdisc, $parent));
	push @{$classdescs{$dev . ": $qdisc:" . $clsid}}, "proto $proto->{name};  (parent 1:$parent)";

	sub procrules($$);
	sub procrules($$) {
	    my $name = shift;
	    my $var = shift;
	    for my $rule (@{$protorules{$name}}) {
		if ($rule->{type} eq 'u32') {
		    fltu32ad($var->{qdisc}, $var->{parent}, $var->{clsid}, $rule->{param});
		    push @{$classdescs{"$dev: $var->{qdisc}:$var->{clsid}"}}, "$rule->{type} $rule->{param}";
		} elsif ($rule->{type} eq 'mark') {
		    fltfwadd($var->{qdisc}, $var->{parent}, $var->{clsid}, $rule->{param});
		    push @{$classdescs{"$dev: $var->{qdisc}:$var->{clsid}"}}, "$rule->{type} $rule->{param}";
		} elsif ($rule->{type} eq 'include') {
		    push @{$classdescs{"$dev: $var->{qdisc}:$var->{clsid}"}}, "$rule->{type} $rule->{param}";
		    procrules($rule->{param}, $var);
		} elsif ($rule->{type} eq 'other') {
		    if (defined($var->{other})) {
			perr("Duplicate definition of 'other' in proto '$name' (original in '$var->{other}->{name}')");
		    } else {
			$var->{other} = { clsid => $var->{clsid}, name => $name };
		    }
		} else {
		    perr("Unsupported rule type \"$rule->{type}\"");
		}
	    }
	}

	procrules($proto->{name}, $var);
    }

    if (!defined($var->{other})) {
	$var->{other}->{clsid} = clalloc($dev, $qdisc);
	clhtbadd($qdisc, $parent, $var->{other}->{clsid}, $rate, $ceil, prioalloc($dev, $qdisc, $parent));
    }
    fltu32ad($qdisc, $parent, $var->{other}->{clsid}, "match ip src 0.0.0.0/0");
    push @{$classdescs{"$dev: $qdisc:$var->{other}->{clsid}"}}, "other";
}

sub checkconf() {
    my $err = 0;

    my %classes;
    for (@classes) {
	if (defined($classes{$_->{name}})) {
	    perr("Class name collision for \"$_->{name}\"");
	    $err = 1;
	}
	$classes{$_->{name}} = 1;
    }
    
    for (@classes) {
	if (defined($_->{parent}) && !defined($classes{$_->{parent}})) {
	    perr("Undefined parent \"$_->{parent}\" for \"$_->{name}\"");
	    $err = 1;
	}
	if (!defined($protogroups{$_->{protogroup}})) {
	    perr("Undefined protogroup \"$_->{protogroup}\" for \"$_->{name}\"");
	    $err = 1;
	} else {
	    for my $proto (@{$protogroups{$_->{protogroup}}}) {
		sub markused($);
		sub markused($) {
		    my $name = shift;
		    for my $rule (@{$protorules{$name}}) {
			$rule->{used} = 1;
			if ($rule->{type} eq 'include') {
			    markused($rule->{param});
			}
		    }
		}

		markused($proto->{name});
		$proto->{used} = 1;
	    }
	}
    }

    for (keys %ranges) {
	if (!defined($classes{$_})) {
	    perr("Undefined class \"$_\" for ranges");
	    $err = 1;
	}
    }

    exit if ($err);
}


# main

{
    # ugly getopt. sorry.

    my @newARGV;

    for (@ARGV) {
	if (/^-d$/) { $debug = 1; }
	elsif (/^-v$/) { $verbose = 1; }
	else { push @newARGV, $_; }
    }

    @ARGV = @newARGV;
}

my @conf;
my $fh;
my $confbase = "";
if (defined $ARGV[0]) {
    $confbase = $ARGV[0]; $confbase =~ s/[^\/]*$//;
    open($fh,$ARGV[0]) or die("Could not open conf!");
} else {
    $fh = *STDIN;
}
unshift @conf, $fh;

sub confre($$@) {
    my ($cmd, $pars, $opars) = @_;
    if (!defined($opars)) { $opars = 0; }
    my $re = '^\s*' . quotemeta($cmd) . '\s+(\S+)'x$pars . '(?:\s+(\S+))?'x$opars . '\s*$';
    return qr/$re/;
}

while (@conf) {
    while (readline($conf[0])) {
	s/#.*$//;
	next if (/^\s*$/);
	my ($mask,$master,$irate,$iceil,$orate,$oceil);

	if ($_ =~ confre("me", 2, 1)) {
	    $me = { in => $1, out => $2, protogroup => "local" };
	    if (defined($3)) { $me->{protogroup} = $3; }
	} elsif ($_ =~ confre("failover", 2)) {
	    $failover = { in => $1, out => $2 };
	} elsif ($_ =~ confre("czf_iface", 3, 1)) {
	    my $czf_iface = { dev => $1, in => $2, out => $3, protogroup => "czf" };
	    if (defined($4)) { $czf_iface->{protogroup} = $4; }
	    push @czf_ifaces, $czf_iface;
	} elsif ($_ =~ confre("net_iface", 3)) {
	    $net_iface = { dev => $1, in => $2, out => $3 };
	} elsif ($_ =~ /^\s*~\s*(?:(\S+)>)?(\S+)\s+(\S+)\s+(\S+)\s+(\S+)\s+(\S+)(?:\s+(\S+))?(?:\s+(\S+))?\s*$/) {
	    my $class = {
		parent => $1, name => $2, irate => $3, iceil => $4,
		orate => $5, oceil => $6, prio => 3, protogroup => "default"
	    };
	    my ($prio, $protogroup) = ($7, $8);

	    if (defined($prio) && $prio =~ /^\d+$/) {
		$class->{prio} = $prio;
	    }

	    if (defined($protogroup)) {
		$class->{protogroup} = $protogroup;
	    }

	    push @classes, $class;
	} elsif ($_ =~ /^\s*:\s*(\S+?)\s+(\S+)\s*$/) {
	    push @{$ranges{$1}}, $2;
	} elsif (/^\s*\?(\S+)\s+(\S+)\s+(\d+)\s+(\d+)\s*$/) {
	    push @{$protogroups{$1}}, { name => $2, prate => $3, pceil => $4 };
	} elsif (/^\s*!(\S+)\s+(\S+)(?:\s+(.+?))?\s*$/) {
	    push @{$protorules{$1}}, { type => $2, param => $3 };
	} elsif ($_ =~ confre("@", 1)) {
	    my $fh;
	    open($fh,$confbase . $1) or die("Could not open conf!");
	    unshift @conf, $fh;
	} else {
	    if (defined $.) { print STDERR "Parse error on line $.\n"; }
	    else { print STDERR "Parse error on $_"; }
	}
    }

    close(shift(@conf));
}

if (!defined($protogroups{default})) { $protogroups{default} = []; }
if (!defined($protogroups{czf})) { $protogroups{czf} = []; }
if (!defined($protogroups{local})) { $protogroups{local} = []; }

checkconf();
prepare();
for (@classes) {
    class($_);
}

if ($verbose) {
    for my $cls (sort keys %classdescs) {
	print $cls . "\n";
	for (@{$classdescs{$cls}}) {
	    print "\t" . $_ . "\n";
	}
	print "\n";
    }
}