#!/usr/bin/perl

use strict;
use POSIX ":sys_wait_h";
use List::Util qw(shuffle);
use Data::Dumper;
use warnings;

$|=1;

my %conf = (
	avail => 10,
	spare => 2,
	ipt_chain => "PREROUTING",
	ipt_rule => "-s %s -p tcp -m tcp --dport 777 -m statistic --mode random --probability %s -j REDIRECT --to-ports %s",
	ipt_table => 'nat',
	ipt_source4 => '193.110.95.0/24',
	ipt_source6 => '2001:1a8f:ffff::/48',
	socks_offset => 9000,
	path => '/opt/tor',
	dns_check => [
		'www.google.com',
		'www.youtube.com',
		'www.facebook.com',
		'www.baidu.com',
		'www.wikipedia.org',
		'www.qq.com',
		'www.taobao.com',
		'www.tmall.com',
		'www.yahoo.com',
		'www.amazon.com',
		'www.twitter.com',
		'www.suho.com',
	],
	getip => 'http://www.spale.com/cgi-bin/ip.cgi',
	proto => {
		ipv4 => 1,
		ipv6 => 1,
	},
	loopinterval => 10,
	statefile => 'var/status.txt',
);

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

my @var;

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


sub run {

	# check state / start new spares / etc..
	check_processes();

	# update conntrack session usage per client
	ipt_conn();

	# display state of all clients
	show_state();

	# commit log
	log_commit();

	# wait for next loop
	sleep($conf{loopinterval});
}

sub show_state {

	log_add("=======================================================================================\n");
	log_add("%3s %-7s %-5s %8s %8s %15s %35s\n", "id", "state", "mode", "ipt_ipv4", "ipt_ipv6", "ipv4", "ipv6");
	log_add("---------------------------------------------------------------------------------------\n");
	for(my $id=0; defined $var[$id]; $id++) {
		log_add("%3s %-7s %-5s %8s %8s %15s %35s\n",
			$id,
			$var[$id]{state},
			$var[$id]{mode},
			$var[$id]{ipt}{ipv4},
			$var[$id]{ipt}{ipv6},
			$var[$id]{ipv4},
			$var[$id]{ipv6},
		);
	}
	log_add("=======================================================================================\n");
}

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

	# check for dead tor clients (killed or crashed)
	for(my $id=0; defined $var[$id]; $id++) {
		if ( defined $var[$id]{pid} && $var[$id]{pid} eq waitpid($var[$id]{pid}, WNOHANG) ) {
			log_add("child $var[$id]{pid} died. resetting instance $id\n");
			tor_dead($id);
		}
	}

	# check for conntrack on dying processes
	for(my $id=0; defined $var[$id]; $id++) {
		if ( $var[$id]{state} eq 'dying' &&
			$var[$id]{ipt}{ipv4} eq 0 &&
			$var[$id]{ipt}{ipv6} eq 0 ) {

			log_add("client $id was dying and has no conntrack anymore, killing\n");
			tor_kill($id);
		}
	}
	

	# check dns query for startup, read and problem clients
	check_dns();

	# count/list available live and available spare clients
	for(my $id=0; defined $var[$id]; $id++) {
		if ( $var[$id]{state} =~ /^(ready|problem)$/ && $var[$id]{mode} eq 'live' ) {
			$inuse++;
		}

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

	}

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

	# check if available live processes are sufficient
	# and get from spare if missing
	if ( $inuse < $conf{avail} ) {
		log_add("missing live processes. looking for spares.\n");

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

	# if spares arent enough, start some more
	# especially useful for quick availability after initial startup
	for(my $x=0; $x<($conf{avail}-$inuse)/2; $x++) { tor_new(); }

	# check if there are enough spares
	if ( scalar(@spares) < $conf{spare} ) {
		log_add("missing spare processes, starting new ones.\n");

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

	# update public ipv4/ipv6 for tunnels
	for(my $id=0; defined $var[$id]; $id++) {
		if ( $var[$id]{state} =~ /^(ready|problem)$/ ) {
			getip($id,4) if $var[$id]{ipv4} eq '-';
			getip($id,6) if $var[$id]{ipv6} eq '-';
		}
	}

}

sub getip {
	my($id,$proto) = @_;


	open(CURL,"timeout 10 curl -f -s -$proto --socks5 localhost:$var[$id]{port} $conf{getip} |");
	while(<CURL>) {
		if ( /^([0-9a-f\:\.]+)$/ ) {
			$var[$id]{"ipv".$proto} = $1;
		}
	}
	close(CURL);
}

sub check_dns {

	my @pids;

	log_add("DNS Check Start ");
	for(my $id=0; defined $var[$id]; $id++) {

		# ignore clients not in bad state
		if ( $var[$id]{state} !~ /^(startup|ready|problem)$/ ) {
			next;
		}

		my $pid = fork();


		if ( $pid ) {
			push @pids, { id => $id, pid => $pid };
			log_add(".");
		}
		else {
			my $ok = 0;
			foreach my $hostname ( shuffle @{$conf{dns_check}} ) {
				system("timeout 5 tor-resolve $hostname 127.0.0.1:$var[$id]{port} >/dev/null 2>&1");
				if ( ! $? ) { $ok=1; last; }
			}

			if ( $ok ) { exit 0; }
			else { exit 1; }
		}
	}

	log_add("\n");

	log_add("DNS Check results: ");

	foreach my $x ( @pids ) {
		my $id  = $x->{id};
		my $pid = $x->{pid};

		waitpid($pid, 0);

		my $ok = $? ? 0 : 1;

		log_add($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 {
				tor_dying($id);
			}
		}
	}
	log_add("\n");
}

sub tor_new {
	my $next;

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

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

	$var[$next] = {
		state => "startup",
		mode  => "spare",
		port  => $conf{socks_offset} + $next,
		pid   => undef,
		ipt   => { ipv4 => 0, ipv6 => 0 },
		ipv4  => '-',
		ipv6  => '-',
	};

	# 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);

	my $pid = fork();

	if ( !defined $pid ) {
		die "Failed to fork, something has gone badly wrong!\n";
	}
	elsif ( $pid ) {
		log_add("new 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 tor_dying {
	my($id) = @_;

	$var[$id]{state} = 'dying';
	$var[$id]{mode}  = 'none';
	ipt_del($id);
}

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

	ipt_del($var[$id]{port});

	$var[$id]{pid}=undef;
	$var[$id]{state}='killed';
	$var[$id]{mode}='none';
}

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

	kill 'KILL', $var[$id]{pid};
}

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

	system("iptables  -t $conf{ipt_table} -A $conf{ipt_chain} " . sprintf($conf{ipt_rule}, $conf{ipt_source4}, "1.0", $port)) if $conf{proto}{ipv4};
	system("ip6tables -t $conf{ipt_table} -A $conf{ipt_chain} " . sprintf($conf{ipt_rule}, $conf{ipt_source6}, "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}, $conf{ipt_source4}, $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}, $conf{ipt_source6}, $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);
	}

	for(my $id=0; defined $var[$id]; $id++) {
		$var[$id]{ipt}{ipv4} = exists $s{$id}{ipv4} ? $s{$id}{ipv4} : 0;
		$var[$id]{ipt}{ipv6} = exists $s{$id}{ipv6} ? $s{$id}{ipv6} : 0;
	}

}

sub log_add {
	open(F,'>>'.$conf{path}.'/'.$conf{statefile}.'.tmp');
	printf F @_;
	close(F);
}

sub log_commit {
	rename $conf{path}.'/'.$conf{statefile}.'.tmp', $conf{path}.'/'.$conf{statefile};
}