#!/usr/bin/perl -w

# $Id: clus 24455 2011-08-03 07:46:44Z zhur $

=head1 NAME

    clus - скрипт для выполнения ряда команд на нескольких хостах

=head1 SYNOPSIS

    clus ci01d.yandex.ru uptime
    clus FRONTEND exec 'ls -la /var/www/ppc.yandex.ru'

=head1 DESCRIPTION

    Общий формат вызова таков:
    clus [OPTIONS] HOSTS (COMMAND PARAMETER*)+

    - HOSTS - список хостов, групп хостов, масок разделённых запятой, возможно с отрицаниами
      Примеры:
        ci01d - один хост
        ci%,ppcsoap% - все хосты, начинающиеся с ci или ppcsoap
        PERL,-ci% - все хосты, входящие в группу PERL за исключением начинающихся с ci
      конечный список хостов вычисляется применением операций по порядку, с последующей уникализацией

    - COMMANDS - список команд для выполнения, возможно с параметрами
      Примеры:
        uptime - на каждом сервере выполнить определённую в конфиге команду uptime
        exec 'date' - на каждом сервере выполнить команду date
        lexec 'scp file %HOST:./' - локально выполнить команду scp, вместо %HOST подставятся по очереди все хосты

    Дополнительные опции:
      --help - справка
      --config=file - прочитать конфиг из указанного файла
      --clear-default-configs - не читать умолчательные конфиги /etc/clusrc, ~/.clusrc
      --check-config - распечатать конфиг и выйти
      --hosts - распечатать через запятую получившийся список хостов
      --parallel - запустить вылолнение на всех серверах в параллельных потоках
      --parallel-level=16 - во сколько потоков запускать
      --parallel-cols=0 - сколько стобцов делать при выводе результатов параллельного исполнения 
                 0 - выбирать автоматически
      --check / --no-check - нужно ли проверять успех выполнения команд
      --complete - для использования для автодополнения в bash
      --multitail - выводить логи через multitail

=head1 CONFIG

    TODO
    в конфиге, после команды, через тильду можно указывать флаги local и nocheck
    например:
    cmd:distribute-file~local = rsync -va %1 %HOST:%1

=cut

use strict;
use Getopt::Long;
use List::Util qw//;
use List::MoreUtils qw/uniq/;
use Data::Dumper;
use File::Basename;
use File::Spec;
use Parallel::ForkManager;
use POSIX qw/ceil/;
$Data::Dumper::Sortkeys = 1;

# каскадный список конфигов
my @default_configs = ("/etc/clusrc", "$ENV{HOME}/.clusrc");
my @configs;
my $CHECK_CONFIG;
my $COMPLETE;

# нужно ли проверять ошибки команд
our $CHECK_SYSTEM = 1;
# локальная ли команда
our $LOCAL_EXEC = 0;
# нужно ли запускать в параллель
our $PARALLEL = 0;
# во сколько потоков запускать
our $PARALLEL_LEVEL = 16;
# сколько стобцов делать при выводе результатов параллельного исполнения 
our $PARALLEL_COLS = 0;
# вывести список хостов через запятую
our $HOSTS = 0;
# выводить ли все сообщений от ssh
our $VERBOSE;
# запускать ли multitail
our $MULTITAIL;

# опции по-умолчанию
my %options = (
    # список хостов для подстановок и проверок
    hosts => [],        
    # нужно ли проверять хосты на вхождение в список
    hosts_check => 1,
    # домен по-умолчанию
    default_domain => '',
    # дирректория по-умолчанию
    default_dir => '',
    # команды
    #    %N  - параметр N
    #    %sN - параметр N, если он равен "-" - считать STDIN
    #    %qN - заквоченный параметр N
    cmd => {
        '' => { 
                'exec'      => '%s1',
                #'wc'        => 'wc %1',
                'cat'       => 'cat %1',
                #'cvs'       => 'CVS_RSH=ssh cvs %s1',
                #'df'        => 'df -h',
                #'ls'        => 'ls -la',
                'uptime'    => 'uptime',
                'grep'      => {cmd => 'zgrep %q1 %2', nocheck => 1},
                'egrep'     => {cmd => 'zgrep -E %q1 %2', nocheck => 1},
                #'ps'        => 'ps uxgww',
                'lexec'     => {cmd => '%1', local => 1},
                }
    }
);

GetOptions(
    "help" => \&usage_full,
    # отсюда будем читать конфигурацию
    "config=s" => \@configs,
    # обнулить умолчательный список конфигов
    "clear-default-configs|cleardefaultconfigs" => sub {@default_configs = ();},
    # распечатать конфиг и выйти
    "check-config|checkconfig" => \$CHECK_CONFIG,
    # вывести список хостов через запятую
    "hosts" => \$HOSTS,
    # распечатать список подсказок по текущим параметрам
    "complete" => \$COMPLETE,
    # нужно ли запускать в параллель
    "parallel" => \$PARALLEL,
    # во сколько потоков запускать
    "parallel-level=i" => \$PARALLEL_LEVEL,
    # сколько столбцов у multitail
    "parallel-cols=i" => \$PARALLEL_COLS,
    # нужно ли проверять ошибки команд
    "check!" => \$CHECK_SYSTEM,
    # выводить ли все сообщений от ssh
    "verbose!" => \$VERBOSE,
    # запускать ли multitail
    "multitail" => \$MULTITAIL,
    ) || exit 0;

# каскадно читаем конфиги
for my $conf (@default_configs) {
    read_config(\%options, $conf, 0);
}
for my $conf (@configs) {
    read_config(\%options, $conf, 1);
}
$options{default_domain} =~ s/^\.+//;

if ($CHECK_CONFIG) {
    # если нужно всего дишь распечатать коннфиг - делаем это
    print Dumper(\%options);
    exit;
} elsif ($COMPLETE) {
    # или печатаем complete для bash
    complete();
    exit;
}

# определяем, с какими хостами будем работать
die "Usage: $0 hosts (command parameter*)+" unless @ARGV >= 2 || @ARGV >= 1 && $HOSTS;
my @hosts = get_hosts_by_mask(shift @ARGV);


if ($HOSTS) {
    print join(",", @hosts), "\n";
    exit;
}

# Для каждого указанного хоста выполняем команды
if (!$PARALLEL) {
    for my $host ( @hosts ) {
        exec_host_cmds($host, @ARGV);
    }
} else {
    my $tmpdir = sprintf('%s/clus', File::Spec->tmpdir());
    mkdir $tmpdir;
    chmod 0777, $tmpdir;

    my %HOST_DATA;
    for my $host ( @hosts ) {
        my $logfile = "$tmpdir/clus-".[getpwuid($<)]->[0]."-$host.log";
        open(my $lh, ">", $logfile) || die "ERROR: Can't open $logfile: $!\n";
        $HOST_DATA{$host} = {logfile => $logfile, lh => $lh};
    }

    my $forks = $PARALLEL_LEVEL;
    $forks++ if $MULTITAIL;

    my $pm = new Parallel::ForkManager($forks);
    my $multitail_pid;
    if ($MULTITAIL) {
        # запускаем multitail
        $multitail_pid = $pm->start;
        if (!$multitail_pid) {
            # расчитываем количество столбцов
            my $cols = $PARALLEL_COLS;
            if ($cols <= 0) {
                if (@hosts > 9) {
                    $cols = 3;
                } elsif (@hosts > 6) {
                    $cols = 2;
                } else {
                    $cols = 1;
                }
            }
            # расчитываем количество пустых файлов, чтобы выглядело красиво
            my $nulls = (ceil(scalar(@hosts)/$cols)*$cols-scalar(@hosts));
            exec( multitail =>
                    -o => 'check_mail:0',
                    ($cols > 1 ? (-s => $cols) : ()),
                    (map {$HOST_DATA{$_}->{logfile}} @hosts),
                    (map {"/dev/null"} 1..$nulls),
                );
            die "Can't exec multitail";
        }
    } else {
        print "Starting deploy, see logs in $tmpdir/ ...\n";
    }

    # обработчик ошибок
    my @error_hosts;
    my $hosts_finished = 0;
    $pm->run_on_finish(sub {
        my ($pid, $exit_code, $ident) = @_;
        if ($ident) {
            push @error_hosts, $ident if $exit_code;
            $hosts_finished++;
            if ($hosts_finished == @hosts && defined $multitail_pid) {
                sleep 2;
                kill 15 => $multitail_pid;
            }
        } else {
            # убили multitail
            $multitail_pid = undef;
        }
    });

    # запуск команд в параллель
    for my $host (@hosts) {
        $pm->start($host) and next;
        open(STDOUT, ">>&=", fileno($HOST_DATA{$host}->{lh}));
        open(STDERR, ">>&=", fileno($HOST_DATA{$host}->{lh}));
        exec_host_cmds($host, @ARGV);	    
        print "\n.\n.\nFINISH!!!\n";
        $pm->finish();
    }

    $pm->wait_all_children();
    print "\n\n";
    if (@error_hosts) {
        die "Error occured for hosts: ".join(", ", @error_hosts)."\n"
            ."Logs:\n"
            .join("", map {"$HOST_DATA{$_}->{logfile}\n"} @error_hosts)
            ;
    } else {
        print "Summary: ".scalar(@hosts)." successes\n";
    }
}

# запускаем на хосте все команды
sub exec_host_cmds {
    my ($host, @lcmds) = @_;
    while(@lcmds) {
        host_cmd($host, shift(@lcmds), \@lcmds);
    }
}

#############################################3

# чтение конфига
sub read_config {
    my ($opt, $file, $die_flag) = @_;
    if (!-f $file) {
        if ($die_flag) {
            die "No config file: '$file'";
        } else {
            return;
        }
    }
    open(my $fh, "<", $file) || die "Can't open config file '$file': $!";
    while(<$fh>) {
        s/\s+/ /g; s/^\s+|\s+$//g;
        # пропускаем пустые строки и комментарии
        next if !$_ || $_ =~ /^#/;
        if (/^include\s+(.*?)\s*$/) {
            die "Recursive include for $1" if $opt->{_seen}->{$1}++;
            read_config($opt, $1, 1);
        } elsif (/^([a-z0-9\.:_-]+)(~\w+(?:,\w+)*)?\s*=\s*(.*)$/i) {
            my ($key, $flags, $val) = ($1, $2, $3);
            my @flags = defined $flags ? ($flags =~ /(\w+)/g) : ();
            if (my @err = grep {!/^(local|nocheck)$/} @flags) {
                die "Incorrect flags ".join(',', @err)." in line $., file $file";
            }
            $val =~ s/\$\$\{([a-z0-9_\-]+)\}/defined $opt->{$1} ? $opt->{$1} : ''/ge;
            add_config_key($opt, $key, $val, \@flags);
        } else {
            die "Incorrect line $. in file $file";
        }
    }
}

# Добавить ключ
sub add_config_key {
    my ($opt, $key, $val, $flags) = @_;
    if ($key =~ /^(default_domain|default_dir)$/) {
        # простые значения
        $opt->{$key} = $val;
    } elsif ($key =~ /^(hosts_check)$/) {
        # флаги
        if ($val =~ /^(yes|1)$/i) {
            $opt->{$key} = 1;
        } elsif ($val =~ /^(no|0)$/i) {
            $opt->{$key} = 0;
        } else {
            die "Incorrect boolean value '$val' for key '$key'";
        }
    } elsif ($key =~ /^(hosts)$/) {
        # массивы
        push @{$opt->{$key}}, map {$opt->{host_group}->{$_} ? @{$opt->{host_group}->{$_}} : $_} split /\s*,\s*/, $val;
    } elsif ($key =~ /^(host_group):(\w+)$/i) {
        # группа хостов
        my @hosts;
        for my $v (split /\s*,\s*/, $val) {
            if ($v =~ /[_%]/) {
                my $re = mask2re($v);
                push @hosts, grep {/^$re$/} @{$opt->{hosts}};
                push @hosts, map {@{$opt->{host_group}->{$_}}} grep {/^$re$/} keys %{$opt->{host_group}};
            } elsif ($opt->{host_group}->{$v}) {
                push @hosts, @{$opt->{host_group}->{$v}};
            } else {
                push @hosts, $v;
            }
        }
        $opt->{$1}->{$2} = [uniq @hosts];
        @{$opt->{hosts}} = uniq @{$opt->{hosts}}, @hosts;
    } elsif ($key =~ /^(?:cmd|command):([^:]+)$/) {
        # глобальные команды
        save_cmd($opt->{cmd}->{''}, $1, $val, $flags);
    } elsif ($key =~ /^(?:cmd|command):([^:]+):([^:]+)$/) {
        # команды для хоста
        my ($mask, $cmd) = ($1, $2);
        for my $host (get_hosts_by_mask($mask)) {
            save_cmd($opt->{cmd}->{$host} ||= {}, $cmd, $val, $flags);
        }
    } else {
        die "Unknown key: '$key'";
    }
}

# записать команду
sub save_cmd {
    my ($hash, $cmd, $val, $flags) = @_;
    my $prop = 'cmd';
    if ($cmd =~ /^(.*)\.(.*)$/) {
        $cmd = $1;
        $prop = $2;
    }
    if ($hash->{$cmd} && !ref $hash->{$cmd}) {
        $hash->{$cmd} = {cmd => $hash->{$cmd}};
    }
    $hash->{$cmd}->{$prop} = $val;
    $hash->{$cmd}->{$_} = 1 for @$flags;
}


# По маске хостов определяем список
sub get_hosts_by_mask {
	my $hosts_mask = shift || '';
	my (@res, %res_hash);
	for my $mask (split /\s*,\s*/, $hosts_mask) {
        my $minus = $mask =~ s/^-//;
	    die "Incorrect host definition: '$mask'" if $mask !~ /^[a-z0-9%_\.-]+$/i;

        my $_proc_host = sub {
            my $host = shift;
            if (!$minus) {
                push(@res, $host) if !$res_hash{$host}++;
            } else {
                @res = grep {$_ ne $host} @res;
                delete $res_hash{$host};
            }
        };

	    if ($mask =~ /[%_]/) {
	        my $re = mask2re($mask);
	        my $found = 0;
	        for my $host (@{$options{hosts}}) {
	            if ($host =~ /^$re$/ 
	                || $options{default_domain} && $host =~ /^$re\.$options{default_domain}$/i
                ) {
                    $_proc_host->($host);
                    $found++;
	            }
	        }
	        for my $group (grep {/^$re$/} sort keys %{ $options{host_group} || {}}) {
                for my $host (@{$options{host_group}->{$group}}) {
                    $_proc_host->($host);
                    $found++;
                }
	        }
	        die "Incorrect mask '$mask' - no one domain matched" if !$found;
        } elsif ($options{host_group} && $options{host_group}->{$mask}) {
            for my $host (@{$options{host_group}->{$mask}}) {
                $_proc_host->($host);
            }
	    } else {
            $_proc_host->($mask);
	    }
	}
	return uniq @res;
}

# Вылолнить команду на указаном сервере
sub host_cmd {
	my ($host, $cmd, $cmds) = @_;
	# Выводим заголовок
	print STDERR "---$host: $cmd\n";
	my $execstr = defined $options{cmd}->{$host}->{$cmd}
	                ? $options{cmd}->{$host}->{$cmd}
	                : $options{cmd}->{''}->{$cmd};
    # Если вместо команды - хэш
    local $CHECK_SYSTEM = $CHECK_SYSTEM;
    local $LOCAL_EXEC = $LOCAL_EXEC;
    if (ref $execstr eq 'HASH') {
        $CHECK_SYSTEM = !$execstr->{nocheck} if defined $execstr->{nocheck};
        $LOCAL_EXEC = $execstr->{local} if defined $execstr->{local};
        $execstr = $execstr->{cmd};
    }
    if (!defined $execstr) {
        die "Command '$cmd' is not defined";
    } elsif (!$execstr) {
        print STDERR "skip...\n";
        return;
    }
    # если есть переменные - заменяем их
    my $vnum = List::Util::max( $execstr =~ /%[qs]*(\d+)/gi );
    if ($vnum) {
        if (@$cmds < $vnum) {
            die "Not enough params for command '$cmd'";
        }
        my @params = splice @$cmds, 0, $vnum;
        $execstr =~ s/%([qs]*)(\d+)/prepare_param($params[$2-1], $1)/gei;
    }
    $execstr =~ s/%HOST/$host/g;
    # Обрабатываем default_dir
    if ($options{default_dir}) {
        $execstr = "cd $options{default_dir}; $execstr";
    }
    if ($host eq 'localhost' || $LOCAL_EXEC) {
        my_system($execstr);
    } else {
        my_system(ssh => (!$VERBOSE ? '-q' : ()), '-p', '10022', $host, $execstr);
    }
}

# подготовить параметр с учётом флагов
{
my $stdin_text = undef;
sub prepare_param {
    my ($text, $flags) = @_;
    if ($flags =~ s/s//i && $text eq '-') {
        # считываем из STDIN
        if (!defined $stdin_text) {
            print STDERR "Enter param text:\n";
            $stdin_text = join '', <STDIN>;
            $stdin_text =~ s/\r?\n?$//;
        }
        $text = $stdin_text;
    }
    if ($flags =~ s/q//i) {
        # если параметр нужно заквотить - делаем это
        $text = quote_param($text);
    }
    die "excess flags '$flags' founded" if $flags;
    return $text;
}
}

# Заквотить шеловский параметр
sub quote_param {
    my $param = shift;
    $param =~ s/(["\$\\!])/\\$1/g;
    return "\"$param\"";;
}

# Обёрнутый фатальный system
sub my_system {
    #print "system: ".Dumper(\@_);
    my $res = system(@_);
    if ($res != 0) {
        print STDERR "system call failed\n";
        exit(1) if $CHECK_SYSTEM;
    } 
}

# более-менее универсальная функция, претендует на Yandex::...
sub complete {
    my $line = $ENV{COMP_LINE} || '';
    my @all_params = split /\s+/, $line;
    my $is_new_param = $line =~ /\s$/;
    # удаляем команду
    shift @all_params;
    # находим начало подсказки
    my $startwith = !$is_new_param ? pop(@all_params) : '';
    $startwith = '' if !defined $startwith;

    my @words = _get_complete_words($startwith, \@all_params);

    print join("", map {"$_\n"} @words), "\n";
}

# скриптозависимая функция - выдаёт список слов
sub _get_complete_words {
    my ($startwith, $all_params) = @_;
    my @params = grep {!/^-/} @$all_params;
    if (defined $startwith && $startwith =~ /^-/) {
        # опции 
        return grep {/^\Q$startwith\E/} qw/--help --config= --clear-default-configs --check-config --parallel --check --no-check/;
    } elsif (!@params) {
        # список хостов/групп
        my ($start, $pattern) = $startwith =~ /^(.*,-?)?(.*)/;
        $start = '' if !defined $start;

        (my $start_mask = $start) =~ s/[,-]+$//;
        my @already_hosts = get_hosts_by_mask($start_mask);
        my %already_hosts = map {$_ => 1} @already_hosts;
        
        my @suggest_hosts = (keys(%{$options{host_group}}), @{$options{hosts}});

        if ($start =~ /-$/) {
            # выбираем тех, кто может что-то эффективно вычесть
            @suggest_hosts = grep {my $mask = $_; scalar(grep {$already_hosts{$_}} get_hosts_by_mask($mask))} @suggest_hosts;
        } else {
            # выбираем тех, кто может что-то эффективно добавить
            @suggest_hosts = grep {my $mask = $_; scalar(grep {!$already_hosts{$_}} get_hosts_by_mask($mask))} @suggest_hosts;            
        }
        return map {"$start$_"} grep {/^\Q$pattern\E/} @suggest_hosts;
    } else {
        # выдаём те команды, которые определены для всех хостов
        my @already_hosts = get_hosts_by_mask($params[0]);
        my %cmds;
        for my $host (@already_hosts) {
            for my $cmd (uniq keys(%{$options{cmd}{$host}}), keys(%{$options{cmd}{''}})) {
                $cmds{$cmd}++;
            }
        }
        return grep {/^\Q$startwith\E/} grep {$cmds{$_} == scalar(@already_hosts)} keys %cmds;
    }
}

sub usage_full {
    system("pod2text-utf8 <$0");
    exit(0);
}

sub mask2re {
    my $mask = shift;

    my $re = $mask;
    $re =~ s/\./\\./g;
    $re =~ s/%/.*/g;
    $re =~ s/_/./g;

    return $re;
}
