#!/usr/bin/perl

# Copyright (C) 2012 STRATO.  All rights reserved.
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public
# License v2 as published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# General Public License for more details.
#
# You should have received a copy of the GNU General Public
# License along with this program; if not, write to the
# Free Software Foundation, Inc., 59 Temple Place - Suite 330,
# Boston, MA 021110-1307, USA.

use strict;
use warnings;
use Getopt::Std;
use Time::HiRes qw(time);
use IPC::Open3;
use Symbol qw(gensym);

my @fs = qw(btrfs zfs);
my %opts = ();
my $ssh_chld = 0;

sub usage {
	print STDERR <<EO_USAGE;
$0 -s source -d destination [...]
  -d  destination type:device (required; always local)
  -k  keep temporary files
  -p  pedanctic mode (stop on error)
  -q  quiet
  -r  remote host (needs ssh root login)
  -s  source type:device (required; remote with -r, otherwise local)
  -t  test spec (e.g. "1,2,3", "1-3,6:1,7:2-4,8-12")
  -v  verbose
type must be one of: @fs
EO_USAGE
	exit(shift);
}

sub out {
	print @_ if !$opts{q};
}

sub verbose {
	print @_ if $opts{v};
}

sub do_local {
	my @out = capture_local(@_);
	print map {"$_\n"} @out if $opts{v};
}

my $remote_host;
my $ssh_socket;
sub args_for_remote {
	return @_ if !$remote_host;
	return ("ssh", "-l", "root", "-S", $ssh_socket, $remote_host,
		map {ref($_) ? $_ : quotemeta($_)} @_);
}

sub do_remote {
	return do_local(args_for_remote(@_));
}

sub capture_local {
	# we only open() to local handles and let perl do the close()s
	my $args = pop if (ref $_[-1]);
	verbose("running @_ ",
		$args->{input} ? "< $args->{input} " : "",
		$args->{output} ? "> $args->{output} " : "",
		"... ");
	my $t_start = time;
	local (*OUTFILE, *INFILE, *DATAFILE);
	my ($fd_in, $fd_out, $fd_out_err, $fd_infile) = ();
	my $data = "";
	if ($args->{input}) {
		open(INFILE, "<", $args->{input})
			or die "\nERROR from command: @_\ncannot redirect ".
			       "input from $args->{input}: $!\n";
		$fd_infile = \*INFILE;
	}
	if ($args->{data}) {
		open(DATAFILE, "<", $args->{data})
			or die "\nERROR from command: @_\ncannot get input ".
			       "data from $args->{data}: $!\n";
	}
	if ($args->{output}) {
		open(OUTFILE, ">", $args->{output})
			or die "\nERROR from command: @_\ncannot redirect ".
			       "output to $args->{output}: $!\n";
		$fd_out_err = gensym;
	}
	my $pid = open3($fd_in,
			$args->{output} ? ">&OUTFILE" : $fd_out,
			$fd_out_err, @_);
	if (!$args->{input}) {
		close($fd_in);
		undef $fd_in;
	}
	$fd_out = $fd_out_err if $args->{output};
	my ($rin, $win, $rout, $wout, @out) = ("", "");
	vec($rin, fileno($fd_out), 1) = 1 if $fd_out;
	vec($win, fileno($fd_in), 1) = 1 if $fd_in;
	while ($fd_in || $fd_out) {
		$rout = $rin;
		$wout = $win;
		select($fd_out ? $rout : undef, $fd_in ? $wout : undef,
		       undef, undef);
		if ($fd_in && $wout eq $win) {
			my $buf = <$fd_infile>;
			if (defined $buf) {
				syswrite($fd_in, $buf)
					or die "\nERROR from command: @_\n".
					       "cannot write to it: $!\n";
			} elsif ($args->{data} && $fd_infile != \*DATAFILE) {
				my $ret = syswrite($fd_in, "\n__END__\n");
				die if $ret != 9;
				$fd_infile = \*DATAFILE;
			} else {
				undef $fd_in;
			}
		}
		if ($fd_out && $rout eq $rin) {
			my $ret = sysread($fd_out, $data, 8192, length($data));
			if (!defined $ret) {
				die "\nERROR from command: @_\ncannot ".
				    "read from it: $!\n";
			}
			undef $fd_out if !$ret;
		}
	}
	waitpid($pid, 0);
	if ($? && !$args->{mayfail}) {
		die "\nERROR from command ($?): @_\noutput:\n$data\n";
	}
	my $t_elapsed = int((time - $t_start) * 100)/100;
	verbose("done ($t_elapsed sec)\n");
	return $data;
}

sub capture_remote {
	return capture_local(args_for_remote(@_));
}

sub ignore_local {
	capture_local(@_, {mayfail => 1});
}

sub ignore_remote {
	return ignore_local(args_for_remote(@_));
}

sub check_blkdev {
	my $dev = shift;
	my $is_loop = 0;
	if (!-b $dev) {
		foreach (split /\n/, capture_local("losetup", "--all")) {
			if (/^([^\s:]+):.*$dev/) {
				ignore_local("umount", $1);
			}
		}
		$dev = capture_local("losetup", "--show", "-f", $dev);
		chomp $dev;
		$is_loop = 1;
	}
	return ($dev, $is_loop);
}

my $p_refgen = "./refgen.pl";
my $p_refgen_actions = "./actions";
my $p_fssum = "fssum";
my $p_fardir = "/tmp/far-$$";
my $p_sumfile = "$p_fardir/fssum.out";
my $p_btrfs = "btrfs";
my $p_zfs = "zfs";
my $p_diff_outfile = "$p_fardir/send.far";
my $p_perl = "perl";

my $allok = getopts("d:hkpqr:s:t:v", \%opts);

if (!$opts{d}) {
	print STDERR "missing -d destination\n";
	$allok = 0;
}

if (!$opts{s}) {
	print STDERR "missing -s source\n";
	$allok = 0;
}

my $allowed_fs = join("|", @fs);

if (!$allok || $opts{h}) {
	usage(!$opts{h});
}

if ($opts{d} !~ /^([^:]+):(.*)$/) {
	print STDERR "wrong format for -d destination ($opts{d})\n";
	$allok = 0;
}
my ($dst_type, $dst_dev, $dst_is_loop) = ($1, check_blkdev($2));
if (!grep $dst_type, @fs) {
	print STDERR "unsupported type for destination\n";
	$allok = 0;
}
if ($opts{s} !~ /^([^:]+):(.*)$/) {
	print STDERR "wrong format for -s source ($opts{s})\n";
	$allok = 0;
}
my ($src_type, $src_dev, $src_is_loop) = ($1, $opts{r} ? $2 : check_blkdev($2));
if (!grep $src_type, @fs) {
	print STDERR "unsupported type for source\n";
	$allok = 0;
}
if ($dst_type eq "zfs") {
	print STDERR "zfs destination type not implemented\n";
	$allok = 0;
}
my $dst_mnt = "$p_fardir/mnt-dst";
my $src_mnt = "$p_fardir/mnt-src";
my $verbose = $opts{v} ? "-v" : "";
my $keep_temp = $opts{k};
$remote_host = $opts{r};
$ssh_socket = "$p_fardir/ssh_sock";

my ($src_type_opt, @p_send, $src_send, $dst_subvol, $send_incr_opt);
if ($src_type eq "zfs") {
	$src_type_opt = "-z";
	@p_send = qw(zfs send -F);
	$send_incr_opt = "-i";
	$src_send = $src_dev;
	$dst_subvol = $src_dev;
	$dst_subvol =~ s{.*/}{};
} elsif ($src_type eq "btrfs") {
	$src_type_opt = "-b";
	@p_send = qw(btrfs send);
	$send_incr_opt = "-i";
	$src_send = "$src_mnt/";
	$dst_subvol = "";
} else {
	die;
}

sub cleanup_temp {
	if (!$keep_temp) {
		ignore_local("umount", $dst_mnt);
		ignore_remote("umount", $src_mnt);
		ignore_local("rm", "-rf", $p_fardir);
		if ($remote_host) {
			ignore_remote("rm", "-rf", $p_fardir);
		}
	}
}

sub cleanup {
	cleanup_temp();
	if ($ssh_chld) {
		kill "TERM", $ssh_chld;
		unlink $ssh_socket;
	}
	if ($dst_is_loop) {
		ignore_local("losetup", "-d", $dst_dev);
	}
	if ($src_is_loop) {
		ignore_local("losetup", "-d", $src_dev);
	}
}
local $SIG{INT} = sub {
	$SIG{INT} = "DEFAULT";
	cleanup();
	kill "INT", $$;
};

sub snapshot_path {
	return "$src_mnt/\@$_[0]" if ($src_type eq "btrfs");
	return "$src_mnt/.zfs/snapshot/$_[0]" if ($src_type eq "zfs");
	die;
}

sub parse_range {
	my $range = shift;
	if ($range !~ /^(\d+)(?:-(\d+))?$/) {
		die qq{failed to parse test spec at "$range"};
	}
	my $start = $1;
	my $end = defined $2 ? $2 : $start;
	if ($start > $end) {
		die qq{reverse range in test spec at "$range"};
	}
	return ($start, $end);
}

my $number_of_digits_in_expand = 3;
sub subtest_range {
	my ($s_start, $s_end) = @_;
	my $n = $number_of_digits_in_expand;
	my $test_spec = ":(?:";
	for ($s_start .. $s_end) {
		$test_spec .= sprintf("%0${n}d", $_);
		$test_spec .= "|";
	}
	chop $test_spec;
	$test_spec .= ")-";
	return $test_spec;
}

my $test_spec = "";
if ($opts{t}) {
	my $n = $number_of_digits_in_expand;
	foreach (split /,/, $opts{t}) {
		my ($tnum, $subtests) = split /:/;
		my ($t_start, $t_end) = $tnum ? parse_range($tnum) : (1, 1);
		my ($s_start, $s_end) = parse_range($subtests) if $subtests;
		for ($t_start .. $t_end) {
			$test_spec .= "^";
			$test_spec .= $tnum ? sprintf("%0${n}d", $_) : "\\d+";
			if ($subtests) {
				$test_spec .= subtest_range($s_start, $s_end);
			} else {
				$test_spec .= ":";
			}
			$test_spec .= "|";
		}
	}
	chop $test_spec;
	eval {
		$test_spec = qr{$test_spec};
	};
	if ($@) {
		die "failed to compile regexp for test spec: $@"
	}
}

mkdir($p_fardir) or die "mkdir for temp dir $p_fardir failed: $!\n";
mkdir($dst_mnt) or die "mkdir $dst_mnt failed: $!\n";
if ($remote_host) {
	# setup shared socket
	$ssh_chld = fork();
	if (!defined $ssh_chld) {
		die "fork failed: $!\n";
	}
	if (!$ssh_chld) {
		close(STDIN);
		close(STDOUT);
		close(STDERR);
		exec("ssh", "-M", "-N", "-l", "root", "-S",
		     $ssh_socket, $remote_host);
		die "exec failed: $!\n";
	}
	verbose("ssh master has pid $ssh_chld, socket $ssh_socket\n");
	do_remote("mkdir", $p_fardir);
	do_remote("mkdir", $src_mnt);
} else {
	mkdir($src_mnt) or die "mkdir $src_mnt failed: $!\n";
}

ignore_local("umount", $dst_dev);
ignore_remote("umount", $src_dev);

my @files = split /\n/, capture_local("ls", $p_refgen_actions);
my @failed_tests = ();
my $cnt = 0;
my $ex_ret = 0;
foreach my $action_file (@files) {
	if ($test_spec && $action_file !~ $test_spec) {
		next;
	}
	my $i = 0;
	my $n_snaps = 0;
	my $t_start = time;
	eval {
		out("running check $action_file... ");
		if ($cnt++) {
			do_local("umount", $dst_mnt);
			do_remote("umount", $src_mnt);
		}
		do_local("mkfs.btrfs", "-L", "fits", $dst_dev);
		do_local("mount", "-o", "noatime", $dst_dev, $dst_mnt);

		do_remote($p_perl, "-", "-p", $src_mnt, $src_type_opt, $src_dev,
			  ($verbose ? $verbose : ()), {
				input => $p_refgen,
				data => "$p_refgen_actions/$action_file",
		});

		$n_snaps = capture_local("grep", "-c", "^snapshot\$",
					 "$p_refgen_actions/$action_file");
		chomp $n_snaps;

		my @exclude = ();
		for ($i = 1; $i <= $n_snaps; $i++) {
			my ($data_file);
			if ($i == 1) {
				$data_file = $p_diff_outfile.".full";
				do_remote(@p_send, "$src_send\@1",
					  {output => $data_file});
			} else {
				my $prev = $i - 1;
				$data_file = $p_diff_outfile.".incr.$prev-$i";
				do_remote(@p_send, $send_incr_opt,
					  "$src_send\@$prev", "$src_send\@$i",
					  {output => $data_file});
			}

			do_local($p_btrfs, "receive",
				 ($verbose ? $verbose : ()),
				 $dst_mnt, {input => $data_file});

			my @x = map { ("-x", "$src_mnt/\@$i/\@$_") } @exclude
				if $src_type eq "btrfs";
			do_remote($p_fssum, @x, "-f", snapshot_path($i),
				  {output => $p_sumfile});
			do_local($p_fssum, "-r", $p_sumfile,
				 "$dst_mnt/$dst_subvol\@$i");
			verbose("done\n");
			push @exclude, $i;
		}
	};

	if ($@) {
		die($@) if $opts{p};
		out($i ? "step $i/$n_snaps" : "preparation", " failed\n")
			if !$verbose;
		verbose($@);
		$ex_ret++;
		push @failed_tests, $action_file;
	} else {
		my $t_elapsed = int((time - $t_start) * 100)/100;
		out("ok ($t_elapsed sec)\n");
	}
}

if (@failed_tests) {
	print "\nFailed tests: ", join(", ", @failed_tests), "\n";
} else {
	print "\nAll tests ok\n";
}

exit($ex_ret);

END {
	cleanup();
};