#!/usr/bin/perl

use strict;
use POSIX ":sys_wait_h";
use Data::Dumper;
use warnings;


my %conf = (
	avail => 2,
	spare => 2,
	ipt_chain => "PREROUTING",
	ipt_rule => "-p tcp -m tcp --dport 777 -m statistic --mode random --probability %s -j REDIRECT --to-ports %s",
	ipt_table => 'nat',
	socks_offset => 9000,
	path => '/opt/tor',
	dns_check => [
		'www.microsoft.com',
		'www.google.com',
	],
	proto => {
		ipv4 => 1,
		ipv6 => 1,
	},
	loopinterval => 60,
);

# states: startup, ready, problem, dying, killed
# modes: spare, live

my @var;

ipt_clean();
while(1) { run(); }


sub run {

	check_processes();
	show_state();

	sleep($conf{loopinterval});
}

sub show_state {

	my $s = ipt_conn();

	printf("======================================\n");
	printf("%3s %-7s %-5s %4s %4s\n", "id", "state", "mode", "ipv4", "ipv6");
	for(my $id=0; defined $var[$id]; $id++) {
		printf("%3s %-7s %-5s %4s %4s\n",
			$id,
			$var[$id]{state},
			$var[$id]{mode},
			exists $s->{$id}{ipv4} ? $s->{$id}{ipv4} : 0,
			exists $s->{$id}{ipv6} ? $s->{$id}{ipv6} : 0,
		);
	}
	printf("======================================\n");
}

sub check_processes {
	my $inuse=0;
	my @spares;

	for(my $id=0; defined $var[$id]; $id++) {
		if ( $var[$id]{pid} eq waitpid($var[$id]{pid}, WNOHANG) ) {
			print "child $var[$id]{pid} died. resetting instance $id\n";
			ipt_del($id);
			$var[$id]{pid}=undef;
			$var[$id]{state}='killed';
			$var[$id]{mode}='none';
		}
	}

	for(my $id=0; defined $var[$id]; $id++) {
		if ( $var[$id]{state} =~ /^(startup|ready|problem)$/ ) {
			check_dns($id);
		}
	}

	for(my $id=0; defined $var[$id]; $id++) {
		if ( $var[$id]{state} eq "ready" && $var[$id]{mode} eq 'live' ) {
			$inuse++;
		}

		if ( $var[$id]{state} =~ /^(startup|ready|problem)$/ && $var[$id]{mode} eq 'spare' ) {
			push @spares, $id;
		}

	}

	print "found " . scalar(@spares) . " spare clients and $inuse live clients\n";

	if ( $inuse < $conf{avail} ) {
		print "missing live processes. looking for spares.\n";

		for(my $x=0; $x<($conf{avail}-$inuse); $x++) {
			my $new = pop @spares;
			if ( defined $new ) {
				if ( $var[$new]{state} eq 'ready' ) {
					print "moving spare to prod.\n";
					tor_prod($new);
				}
				else {
					push @spares, $new;
				}
			}
		}
	}

	if ( scalar(@spares) < $conf{spare} ) {
		print "missing spare processes, starting new ones.\n";

		for(my $x=0; $x < ($conf{spare}-scalar(@spares)); $x++) {
			tor_new();
		}
	}
}

sub check_dns {

	my($id) = @_;

	my $ok = 0;
	foreach my $hostname ( @{$conf{dns_check}} ) {
		print "Checking DNS ($hostname) for $id ";
		system("timeout 5 tor-resolve $hostname 127.0.0.1:$var[$id]{port} >/dev/null 2>&1");
		if ( $? ) {
			print "ERROR\n";
		}
		else {
			print "OK\n";
			$ok++;
		}
	}

	if ( $var[$id]{state} eq 'startup' ) {
		if ( $ok ) {
			$var[$id]{state} = 'ready';
		}
		else {
			$var[$id]{state} = 'problem';
		}
	}
	elsif ( $var[$id]{state} eq 'ready' && !$ok ) {
		$var[$id]{state} = 'problem';
	}
	elsif ( $var[$id]{state} eq 'problem' ) {
		if ( $ok ) {
			$var[$id]{state} = 'ready';
		}
		else {
			$var[$id]{state} = 'dying';
			ipt_del($id);
		}
	}
}

sub tor_new {
	my $next;

	print "initiating a new tor client\n";

	# try to recycle a position
	for(my $id=0; defined $var[$id]; $id++) {
		if ( $var[$id]{state} eq 'killed' ) {
			print "recovering id $id\n";
			$next = $id;
			last;
		}
	}

	# alternatively, create a new position
	if ( !defined $next ) {
		$next = @var;
		print "new id $next\n";
	}

	$var[$next] = {
		state => "startup",
		mode  => "spare",
		port  => $conf{socks_offset} + $next,
		pid   => undef,
	};

	print "creating configuration file and dependencies\n";

	# create a config file and the required files and directories
	my $datadir = "$conf{path}/var/lib/$next";
	my %files = (
		config  => "$conf{path}/etc/$next.conf",
		pidfile => "$conf{path}/run/$next.pid",
		socket  => "$conf{path}/run/$next.socks",
		ctrl    => "$conf{path}/run/$next.ctrl",
		cookie  => "$conf{path}/run/$next.cookie",
	);

	system("mkdir -p $datadir");
	system("chmod 700 $datadir");
	system("chown debian-tor:debian-tor $datadir");

	foreach my $f ( keys %files ) {
		system("touch $files{$f}");
		system("chown debian-tor:debian-tor $files{$f}");
	}

	open(C,">$files{config}");
	print C <<EOF;
DataDirectory $datadir
PidFile $files{pidfile}
RunAsDaemon 0
User debian-tor

ControlSocket $files{ctrl} GroupWritable RelaxDirModeCheck
ControlSocketsGroupWritable 1
#SocksPort unix:$files{socket} WorldWritable
SocksPort 0.0.0.0:$var[$next]{port}
SocksPort [::]:$var[$next]{port}

CookieAuthentication 1
CookieAuthFileGroupReadable 1
CookieAuthFile $files{cookie}

BandwidthRate 10MB

Log notice syslog

EOF
	close(C);

	print "forking...\n";
	my $pid = fork();

	if ( !defined $pid ) {
		die "Failed to fork, something has gone badly wrong!\n";
	}
	elsif ( $pid ) {
		print "tor client pid $pid\n";
		$var[$next]{pid} = $pid;
	}
	else {
		chdir $conf{path};
		exec("tor -f $files{config} >/dev/null");
		die "something weng really wrong, failed to exec tor";
		exit;
	}
}

sub tor_prod {
	my($id) = @_;

	# move process from spare to prod

	$var[$id]{mode} = 'live';
	ipt_add($var[$id]{port});
}

sub ipt_add {
	my($port) = @_;

	system("iptables  -t $conf{ipt_table} -A $conf{ipt_chain} " . sprintf($conf{ipt_rule}, "1.0", $port)) if $conf{proto}{ipv4};
	system("ip6tables -t $conf{ipt_table} -A $conf{ipt_chain} " . sprintf($conf{ipt_rule}, "1.0", $port)) if $conf{proto}{ipv6};

	ipt_recalc();
}

sub ipt_del {
	my($port) = @_;

	if ( $conf{proto}{ipv4} ) {
		open(IPT,"iptables -t $conf{ipt_table} -L $conf{ipt_chain} -n --line-numbers |");
		while(<IPT>) {
			chomp;
			if ( /^(\d+) .*mode random probability .* redir ports $port$/ ) {
				system("iptables -t $conf{ipt_table} -D $conf{ipt_chain} $1");
			}
		}
		close(IPT);
	}

	if ( $conf{proto}{ipv6} ) {
		open(IPT,"ip6tables -t $conf{ipt_table} -L $conf{ipt_chain} -n --line-numbers |");
		while(<IPT>) {
			chomp;
			if ( /^(\d+) .*mode random probability .* redir ports $port$/ ) {
				system("ip6tables -t $conf{ipt_table} -D $conf{ipt_chain} $1");
			}
		}
		close(IPT);
	}
	

	ipt_recalc();
}

sub ipt_clean {
	system("iptables  -t $conf{ipt_table} -F $conf{ipt_chain}") if $conf{proto}{ipv4};
	system("ip6tables -t $conf{ipt_table} -F $conf{ipt_chain}") if $conf{proto}{ipv6};
}

sub ipt_recalc {
	my @rules;

	if ( $conf{proto}{ipv4} ) {
		open(IPT,"iptables -t $conf{ipt_table} -L $conf{ipt_chain} -n --line-numbers |");
		while(<IPT>) {
			chomp;
			if ( /^(\d+) .*mode random probability .* redir ports (\d+)$/ ) {
				push @rules, [ $1, $2 ];
			}
		}
		close(IPT);

		my $nums = @rules;

		for(my $num = 0; defined $rules[$num]; $num++) {
			my($rulenum,$port)=@{$rules[$num]};
	
			my $prob = 1 / ($nums - $num);
	
			system("iptables -t $conf{ipt_table} -R $conf{ipt_chain} $rulenum " . sprintf($conf{ipt_rule}, $prob, $port));
		}
	}

	if ( $conf{proto}{ipv6} ) {
		open(IPT,"ip6tables -t $conf{ipt_table} -L $conf{ipt_chain} -n --line-numbers |");
		while(<IPT>) {
			chomp;
			if ( /^(\d+) .*mode random probability .* redir ports (\d+)$/ ) {
				push @rules, [ $1, $2 ];
			}
		}
		close(IPT);

		my $nums = @rules;

		for(my $num = 0; defined $rules[$num]; $num++) {
			my($rulenum,$port)=@{$rules[$num]};
	
			my $prob = 1 / ($nums - $num);
	
			system("ip6tables -t $conf{ipt_table} -R $conf{ipt_chain} $rulenum " . sprintf($conf{ipt_rule}, $prob, $port));
		}
	}

}

sub ipt_conn {

	my %s;

	if ( $conf{proto}{ipv4} ) {
		open(IPT,"conntrack -L -f ipv4 2>/dev/null |");
		while(<IPT>) {
			chomp;
			if ( / dport=777 .* sport=(\d+) / ) {
				my $port = $1;
				if ( $port >= $conf{socks_offset} ) {
					my $id = $port - $conf{socks_offset};
					$s{$id}{ipv4}++;
				}
			}
		}
		close(IPT);
	}

	if ( $conf{proto}{ipv6} ) {
		open(IPT,"conntrack -L -f ipv6 2>/dev/null |");
		while(<IPT>) {
			chomp;
			if ( / dport=777 .* sport=(\d+) / ) {
				my $port = $1;
				if ( $port >= $conf{socks_offset} ) {
					my $id = $port - $conf{socks_offset};
					$s{$id}{ipv6}++;
				}
			}
		}
		close(IPT);
	}

	return \%s;
}