#!/opt/perl/bin/perl

use strict;
use warnings;

#use Smart::Comments;
use Term::ReadKey;
use Net::SSH::Expect;
use List::Util qw(max);
use Getopt::Std;
use Parallel::ForkManager;

my %opts;
getopts('ht:u:c:', \%opts) or help(1);

if ($opts{h}) { help(0) }
my $concurrency = $opts{c} || 10;
my $user = $opts{u} || $ENV{USER} || $ENV{LOGNAME} or
    die "No -u <user> option specified.\n";
my $timeout = $opts{t} || 2;

sub help {
    my $code = shift;
    my $out = <<'_EOC_';
Usage: batch-ssh.pl [options] <command> <hostname pattern>...
Options:
    -c <concurrency>       Specify the maximum number of concurrent
                           SSH sessions.
    -h                     Print this help.
    -t <timeout>           Timeout for command execution in seconds.
    -u <user>              Specify the user name for login (default
                           to the environment USER or LOGNAME).
Examples:
    batch-ssh.pl -c 10 -t 2 $'echo \'hello\'' 'tq9010[01-18,20,22-25].foo.com'
    batch-ssh.pl -u agentzh 'sudo rm /logs/*' a.com [aa-zz].com t[ab1-ab3].com
    batch-ssh.pl 'df -h' 'a[1-2].b[5-16].com'
_EOC_
    if ($code == 0) { print $out } else { warn $out }
    exit($code);
}

my $fork = new Parallel::ForkManager($concurrency);

my $cmd = shift or die "No command specified.\n";
my @hosts = map { parse_hostname_pat($_) } @ARGV;
print "Hosts: @hosts\n";

my $password;
if ($cmd =~ /\bsudo\b/) { read_password(); }

for my $host (@hosts) {
    $fork->start and next;
    exec_cmd_on_host($cmd, $host);
    $fork->finish;
}
$fork->wait_all_children;

sub exec_cmd_on_host {
    my ($cmd, $host) = @_;
    my $output_buf;
    $output_buf .= "\n================== $host ==========================\n";
    my $ssh = Net::SSH::Expect->new (
        host => $host,
        #password=> $host,
        user => $user,
        raw_pty => 1,
    );
    $ssh->run_ssh() or die "SSH process couldn't start: $!";
    #my $prompt = $ssh->read_all(1);
    #$prompt =~ /\$\s*\z/ or die "Where's the prompt? ** $prompt\n";
    #print "Prompt: $prompt\n";
    $ssh->exec("stty raw -echo", 2);

    if ($cmd =~ /\bsudo\b/) {
        if (!$password) {
            read_password();
        }
        $ssh->send($cmd);
        if ($ssh->waitfor('Password:\s*\z', 2)) {
            $ssh->send($password);
        } else {
            #warn "prompt 'sudo su -' not found after 1 second\n";
        }
        my $output = $ssh->read_all($timeout);
        $output =~ s/\n\S+\$\s*\z/\n/;
        $output_buf .= $output;
    } else {
        my $output = $ssh->exec($cmd, $timeout);
        $output =~ s/\n\S+\$\s*\z/\n/;
        $output_buf .= $output;
    }
    print $output_buf;
    $ssh->close;
}

sub parse_hostname_pat {
    my $pat = shift;
    return unless $pat;
    ### $pat
    my @seg;
    while (1) {
        if ($pat =~ /\G\[([^\]]+)\]/gc) {
            my $range = $1;
            my $interval_pat = qr/\w+(?:-\w+)?/;
            if ($range !~ m/^$interval_pat(?:,$interval_pat)*$/) {
                die "Bad number range: [$range]\n";
            }
            my @intervals = split /,/, $range;
            my @num;
            for (@intervals) {
                my ($a, $b) = split /-/, $_;
                #if (defined $b && ($a =~ /\D/ || $b =~ /\D/) && length $a ne length $b) {
                    #die "End points are not of equal lengths in the host range: $a-$b\n";
                #}
                push @num, defined $b ? $a..$b : $a;
                #print "@num";
            }
            push @seg, \@num;
        } elsif ($pat =~ /\G[^\[]+/gc) {
            push @seg, [$&];
            next;
        } else {
            last;
        }
    }
    my $res = expand_seg(\@seg);
    ### $res
    return @$res;
}

sub expand_seg {
    my ($list, $prefixes) = @_;
    my $cur = shift @$list;
    return $prefixes unless defined $cur;
    my @new_prefixes;
    if (!$prefixes) {
        for my $alt (@$cur) {
            push @new_prefixes, $alt;
        }
    } else {
        for my $prefix (@$prefixes) {
            for my $alt (@$cur) {
                push @new_prefixes, $prefix . $alt;
            }
        }
    }
    return expand_seg($list, \@new_prefixes);
}

sub read_password {
    print "Enter the password: ";

    ReadMode(2);
    while (not defined ($password = ReadLine(0))) {
    }
    ReadMode(0);

    $password =~ s/\n//s;
    print "\n";
    if (!$password) { die "No password specified.\n"; }
}