MrShell.pm

   1 package App::MrShell;
   2
   3 use strict;
   4 use warnings;
   5
   6 use Carp;
   7 use POSIX;
   8 use Config::Tiny;
   9 use POE qw( Wheel::Run );
  10 use Term::ANSIColor qw(:constants);
  11 use Text::Balanced;
  12
  13 our $VERSION = '2.0210';
  14 our @DEFAULT_SHELL_COMMAND = (ssh => '-o', 'BatchMode yes', '-o', 'StrictHostKeyChecking no', '-o', 'ConnectTimeout 20', '[%u]-l', '[]%u', '%h');
  15
  16 # new {{{
  17 sub new {
  18     my $class = shift;
  19     my $this  = bless { hosts=>[], cmd=>[], _shell_cmd=>[@DEFAULT_SHELL_COMMAND] }, $class;
  20
  21     return $this;
  22 }
  23 # }}}
  24
  25 # _process_space_delimited {{{
  26 sub _process_space_delimited {
  27     my $this = shift;
  28     my $that = shift;
  29
  30     my @output;
  31     while( $that ) {
  32         if( $that =~ m/^\s*['"]/ ) {
  33             my ($tok, $rem) = Text::Balanced::extract_delimited($that, qr(["']));
  34
  35             ($tok =~ s/^(['"])// and $tok =~ s/$1$//) or die "internal error processing space delimited";
  36
  37             push @output, $tok;
  38             $that = $rem;
  39
  40         } else {
  41             my ($tok, $rem) = split ' ', $that, 2;
  42
  43             push @output, $tok;
  44             $that = $rem;
  45         }
  46     }
  47
  48     return @output
  49 }
  50 # }}}
  51 # _process_hosts {{{
  52 sub _process_hosts {
  53     my $this = shift;
  54
  55     my @h = do {
  56         my @tmp = map { my $k = $_; $k =~ s/^\@// ? @{$this->{groups}{$k} or die "couldn't find group: \@$k\n"} : $_ } @_;
  57         my %h; @h{@tmp} = ();
  58         for(keys %h) {
  59             if(my ($k) = m/^\-(.+)/) {
  60                 delete $h{$_};
  61                 delete $h{$k};
  62             }
  63         }
  64         sort keys %h;
  65     };
  66
  67     my $o = my $l = $this->{_host_width} || 0;
  68     for( map { length $this->_host_route_to_nick($_) } @h ) {
  69         $l = $_ if $_>$l
  70     }
  71
  72     $this->{_host_width} = $l if $l != $o;
  73
  74     return @h;
  75 }
  76 # }}}
  77 # _host_route_to_nick {{{
  78 sub _host_route_to_nick {
  79     my $this = shift;
  80
  81     return join "", shift =~ m/(?:!|[^!]+$)/g
  82 }
  83 # }}}
  84
  85 # set_shell_command_option {{{
  86 sub set_shell_command_option {
  87     my $this = shift;
  88     my $arg = shift;
  89
  90     $arg = [ $this->_process_space_delimited($arg||"") ] unless ref $arg;
  91
  92     if( ref($arg) eq "ARRAY" ) {
  93         if( not $arg->[0] ) {
  94             $this->{_shell_cmd} = sub {
  95                 shift @$arg;
  96                 my $perl_program = shift @$arg;
  97
  98                 local @ARGV = $this->subst_cmd_vars(@$arg, @_);
  99
 100                 unless( defined( do $perl_program ) ) {
 101                     die $@ if $@;
 102                     die $! if $!;
 103                     die "execution failure for $perl_program";
 104                 }
 105
 106                 return;
 107             };
 108
 109         } else {
 110             $this->{_shell_cmd} = [ @$arg ]; # make a real copy
 111         }
 112
 113     } elsif( ref($arg) eq "CODE" ) {
 114         $this->{_shell_cmd} = $arg;
 115     }
 116
 117     return $this;
 118 }
 119 # }}}
 120 # set_group_option {{{
 121 sub set_group_option {
 122     my $this   = shift;
 123     my $groups = ($this->{groups} ||= {});
 124
 125     my ($name, $value);
 126     while( ($name, $value) = splice @_, 0, 2 and $name and $value ) {
 127         if( ref($value) eq "ARRAY" ) {
 128             $groups->{$name} = [ @$value ]; # make a real copy
 129
 130         } else {
 131             $groups->{$name} = [ $this->_process_space_delimited( $value ) ];
 132         }
 133     }
 134
 135     my @groups = keys %{ $this->{groups} };
 136     my $replace_limit = 30;
 137     REPLACE_GROPUS: {
 138         my $replaced = 0;
 139
 140         for my $group (@groups) {
 141             my $hosts = $groups->{$group};
 142
 143             my $r = 0;
 144             for(@$hosts) {
 145                 if( m/^@(.+)/ ) {
 146                     if( my $g = $groups->{$1} ) {
 147                         $_ = $g;
 148
 149                         $r ++;
 150                     }
 151                 }
 152             }
 153
 154             if( $r ) {
 155                 my %h;
 156                 @h{ map {ref $_ ? @$_ : $_} @$hosts } = ();
 157                 $groups->{$group} = [ keys %h ];
 158                 $replaced ++;
 159             }
 160         }
 161
 162         $replace_limit --;
 163         last if $replace_limit < 1;
 164         redo if $replaced;
 165     }
 166
 167     return $this;
 168 }
 169 # }}}
 170 # set_logfile_option {{{
 171 sub set_logfile_option {
 172     my $this = shift;
 173     my $file = shift;
 174     my $trunc = shift;
 175
 176     unless( our $already_compiled++ ) {
 177         my $load_ansi_filter_package = q {
 178             package App::MrShell::ANSIFilter;
 179             use Symbol;
 180             use Tie::Handle;
 181             use base 'Tie::StdHandle';
 182
 183             my %orig;
 184
 185             sub PRINT {
 186                 my $this = shift;
 187                 my @them = @_;
 188                 s/\e\[[\d;]+m//g for @them;
 189                 print {$orig{$this}} @them;
 190             }
 191
 192             sub filtered_handle {
 193                 my $pfft = gensym();
 194                 my $it = tie *{$pfft}, __PACKAGE__ or die $!;
 195                 $orig{$it} = shift;
 196                 $pfft;
 197             }
 198
 199         1};
 200
 201         eval $load_ansi_filter_package or die $@; ## no critic -- sometimes this kind of eval is ok
 202         # (This probably isn't one of them.)
 203     }
 204
 205     open my $log, ($trunc ? ">" : ">>"), $file or croak "couldn't open $file for write: $!"; ## no critic -- I mean to pass this around, shut up
 206
 207     $this->{_log_fh} = App::MrShell::ANSIFilter::filtered_handle($log);
 208
 209     return $this;
 210 }
 211 # }}}
 212 # set_debug_option {{{
 213 sub set_debug_option {
 214     my $this = shift;
 215     my $val = shift;
 216
 217     # -d 0 and -d 1 are the same
 218     # -d 2 is a level up, -d 4 is even more
 219     # $val==undef clears the setting
 220
 221     if( not defined $val ) {
 222         delete $this->{debug};
 223         return $this;
 224     }
 225
 226     $this->{debug} = $val ? $val : 1;
 227
 228     return $this;
 229 }
 230 # }}}
 231 # set_no_command_escapes_option {{{
 232 sub set_no_command_escapes_option {
 233     my $this = shift;
 234
 235     $this->{no_command_escapes} = shift || 0;
 236
 237     return $this;
 238 }
 239 # }}}
 240
 241 # groups {{{
 242 sub groups {
 243     my $this = shift;
 244
 245     return unless $this->{groups};
 246     return wantarray ? %{$this->{groups}} : $this->{groups};
 247 }
 248 # }}}
 249
 250 # set_usage_error($&) {{{
 251 sub set_usage_error($&) { ## no critic -- prototypes are bad how again?
 252     my $this = shift;
 253     my $func = shift;
 254     my $pack = caller;
 255     my $name = $pack . "::$func";
 256     my @args = @_;
 257
 258     $this->{_usage_error} = sub {
 259         no strict 'refs'; ## no critic -- how would you call this by name without this?
 260         $name->(@args)
 261     };
 262
 263     return $this;
 264 }
 265 # }}}
 266 # read_config {{{
 267 sub read_config {
 268     my ($this, $that) = @_;
 269
 270     $this->{_conf} = Config::Tiny->read($that) if -f $that;
 271
 272     for my $group (keys %{ $this->{_conf}{groups} }) {
 273         $this->set_group_option( $group => $this->{_conf}{groups}{$group} );
 274     }
 275
 276     if( my $c = $this->{_conf}{options}{'shell-command'} ) {
 277         $this->set_shell_command_option( $c );
 278     }
 279
 280     if( my $c = $this->{_conf}{options}{'logfile'} ) {
 281         my $t = $this->{_conf}{options}{'truncate-logfile'};
 282         my $v = ($t ? 1:0);
 283            $v = 0 if $t =~ m/(?:no|false)/i;
 284
 285         $this->set_logfile_option($c, $v);
 286     }
 287
 288     if( my $c = $this->{_conf}{options}{'no-command-escapes'} ) {
 289         my $v = ($c ? 1:0);
 290            $v = 0 if $c =~ m/(?:no|false)/i;
 291
 292         $this->set_no_command_escapes_option( $v );
 293     }
 294
 295     return $this;
 296 }
 297 # }}}
 298 # set_hosts {{{
 299 sub set_hosts {
 300     my $this = shift;
 301
 302     $this->{hosts} = [ $this->_process_hosts(@_) ];
 303
 304     return $this;
 305 }
 306 # }}}
 307 # queue_command {{{
 308 sub queue_command {
 309     my $this = shift;
 310     my @hosts = @{$this->{hosts}};
 311
 312     unless( @hosts ) {
 313         if( my $h = $this->{_conf}{options}{'default-hosts'} ) {
 314             @hosts = $this->_process_hosts( $this->_process_space_delimited($h) );
 315
 316         } else {
 317             if( my $e = $this->{_usage_error} ) {
 318                 warn "Error: no hosts specified\n";
 319                 $e->();
 320
 321             } else {
 322                 croak "set_hosts before issuing queue_command";
 323             }
 324         }
 325     }
 326
 327     for my $h (@hosts) {
 328         push @{$this->{_cmd_queue}{$h}}, [@_]; # make a real copy
 329     }
 330
 331     return $this;
 332 }
 333 # }}}
 334 # run_queue {{{
 335 sub run_queue {
 336     my $this = shift;
 337
 338     $this->{_session} = POE::Session->create( inline_states => {
 339         _start       => sub { $this->poe_start(@_) },
 340         child_stdout => sub { $this->line(1, @_) },
 341         child_stderr => sub { $this->line(2, @_) },
 342         child_signal => sub { $this->sigchld(@_) },
 343         stall_close  => sub { $this->_close(@_) },
 344         ErrorEvent   => sub { $this->error_event },
 345     });
 346
 347     POE::Kernel->run();
 348
 349     return $this;
 350 }
 351 # }}}
 352
 353 # std_msg {{{
 354 sub std_msg {
 355     my $this  = shift;
 356     my $host  = shift;
 357     my $cmdno = shift;
 358     my $fh    = shift;
 359     my $msg   = shift;
 360
 361     my $host_msg = $host ? $this->_host_route_to_nick($host) . ": " : "";
 362     my $time_str = strftime('%H:%M:%S', localtime);
 363
 364     print $time_str,
 365         sprintf(' %-*s', $this->{_host_width}+2, $host_msg),
 366             ( $fh==2 ? ('[',BOLD,YELLOW,'stderr',RESET,'] ') : () ), $msg, RESET, "\n";
 367
 368     if( $this->{_log_fh} ) {
 369         $time_str = strftime('%Y-%m-%d %H:%M:%S', localtime);
 370
 371         # No point in printing colors, stripped anyway.  Formatting columns is
 372         # equally silly -- in append mode anyway.
 373         $host_msg = $host ? "$host: " : "";
 374         print {$this->{_log_fh}} "$time_str $host_msg", ($fh==2 ? "[stderr] " : ""), $msg, "\n";
 375     }
 376
 377     return $this;
 378 }
 379 # }}}
 380
 381 # line {{{
 382 sub line {
 383     my $this = shift;
 384     my $fh   = shift;
 385     my ($line, $wid) = @_[ ARG0, ARG1 ];
 386     my ($kid, $host, $cmdno, $lineno) = @{$this->{_wid}{$wid}};
 387
 388     $$lineno ++;
 389     $this->std_msg($host, $cmdno, $fh, $line);
 390
 391     return;
 392 }
 393 # }}}
 394
 395 # sigchld {{{
 396 sub _sigchld_exit_error {
 397     my $this = shift;
 398     my ($pid, $exit) = @_[ ARG1, ARG2 ];
 399     $exit >>= 8;
 400
 401     $this->std_msg("?", -1, 0, BOLD.RED."-- sigchld received for untracked pid($pid, $exit), probably a bug in Mr. Shell --");
 402
 403     return;
 404 }
 405
 406 sub sigchld {
 407     my $this = shift; # ARG0 is the signal name string
 408     my ($kid, $host, $cmdno, @c) = @{ $this->{_pid}{ $_[ARG1] } || return $this->_sigchld_exit_error(@_) };
 409
 410     # NOTE: this usually isn't an error, sometimes the sigchild will arrive
 411     # before the handles are "closed" in the traditional sense.  We get error
 412     # eveents for errors.
 413     #### # $this->std_msg($host, $cmdno, 0, RED.'-- error: unexpected child exit --');
 414
 415     # NOTE: though, the exit value may indicate an actual error.
 416     if( (my $exit = $_[ARG2]) != 0 ) {
 417         # XXX: I'd like to do more here but I'm waiting to see what Paul
 418         # Fenwick has to say about it.
 419         $exit >>= 8;
 420
 421         my $reset = RESET;
 422         my $black = BOLD.BLACK;
 423         my $red   = RESET.RED;
 424
 425         $this->std_msg($host, $cmdno, 0, "$black-- shell exited with nonzero status: $red$exit$black --");
 426     }
 427
 428     $_[KERNEL]->yield( stall_close => $kid->ID, 0 );
 429
 430     return;
 431 }
 432 # }}}
 433 # _close {{{
 434 sub _close {
 435     my $this = shift;
 436     my ($wid, $count) = @_[ ARG0, ARG1 ];
 437
 438     return unless $this->{_wid}{$wid}; # sometimes we'll get a sigchild *and* a close event
 439
 440     # NOTE: I was getting erratic results with some fast running commands and
 441     # guessed that I was sometimes getting the close event before the stdout
 442     # event. Waiting through the kernel loop once is probably enough, but I
 443     # used 3 because it does't hurt either.
 444
 445     if( $count > 3 ) {
 446         my ($kid, $host, $cmdno, $lineno, @c) = @{ delete $this->{_wid}{$wid} };
 447
 448         $this->std_msg($host, $cmdno, 0, BOLD.BLACK.'-- eof --') if $$lineno == 0;
 449
 450         if( @c ) {
 451
 452             $cmdno ++;
 453
 454             $this->start_queue_on_host($_[KERNEL] => $host, $cmdno, @c);
 455             $this->std_msg($host, $cmdno, 0, BOLD.BLACK."-- starting: @{$c[0]} --");
 456         }
 457
 458         delete $this->{_pid}{ $kid->PID };
 459
 460     } else {
 461         $_[KERNEL]->yield( stall_close => $wid, $count+1 );
 462     }
 463
 464     return;
 465 }
 466 # }}}
 467 # error_event {{{
 468 sub error_event {
 469     my $this = shift;
 470     my ($operation, $errnum, $errstr, $wid) = @_[ARG0 .. ARG3];
 471     my ($kid, $host, $cmdno, @c) = @{ delete $this->{_wid}{$wid} || return };
 472     delete $this->{_pid}{ $kid->PID };
 473
 474     $errstr = "remote end closed" if $operation eq "read" and not $errnum;
 475     $this->std_msg($host, $cmdno, 0, RED."-- $operation error $errnum: $errstr --");
 476
 477     return;
 478 }
 479 # }}}
 480
 481 # set_subst_vars {{{
 482 sub set_subst_vars {
 483     my $this = shift;
 484
 485     while( my ($k,$v) = splice @_, 0, 2 ) {
 486         $this->{_subst}{$k} = $v unless exists $this->{_subst}{$k};
 487     }
 488
 489     return $this;
 490 }
 491 # }}}
 492 # subst_cmd_vars {{{
 493 sub subst_cmd_vars {
 494     my $this = shift;
 495     my %h = %{ delete($this->{_subst}) || {} };
 496     my $host = $h{'%h'};
 497
 498     my @c = @_; # copy this so it doesn't get altered upstream
 499                 # (I'd swear I shoulnd't need to do this at all, but it's
 500                 #  proovably true that I do.)
 501
 502     if( $host =~ m/\b(?!<\\)!/ ) {
 503         my @hosts = split '!', $host;
 504
 505         my @indexes_of_replacements;
 506         for(my $i=0; $i<@c; $i++) {
 507             if( $c[$i] eq '%h' ) {
 508                 splice @c, $i, 1, $hosts[0];
 509
 510                 push @indexes_of_replacements, $i;
 511
 512                 for my $h (reverse @hosts[1 .. $#hosts]) {
 513                     splice @c, $i+1, 0, @c[0 .. $i-1] => $h;
 514                     push @indexes_of_replacements, $i+1 + $indexes_of_replacements[-1];
 515
 516                     unless( $this->{no_command_escapes} ) {
 517                         for my $arg (@c[$i+1 .. $#c]) {
 518
 519                             # NOTE: This escaping is going to be an utter pain to maintain...
 520
 521                             $arg =~ s/([`\$])/\\$1/g;
 522
 523                             if( $arg =~ m/[\s()]/ ) {
 524                                 $arg =~ s/([\\"])/\\$1/g;
 525                                 $arg = "\"$arg\"";
 526                             }
 527                         }
 528                     }
 529                 }
 530             }
 531         }
 532
 533         my $beg = 0;
 534         for my $i (@indexes_of_replacements) {
 535             if( $c[$i] =~ s/^([\w.\-_]+)@// ) {
 536                 my $u = $1;
 537                 for(@c[$beg .. $i-1]) {
 538                     s/^(\[\%u\]|\[\](?=\%u))//;
 539                     $_ = $u if $_ eq '%u';
 540                 }
 541
 542             } else {
 543                 # NOTE: there's really no need to go through and remove [%u]
 544                 # conditional options, they'll automatically get nuked below
 545                 $c[$i] =~ s/\\@/@/g;
 546             }
 547             $beg = $i+1;
 548         }
 549
 550         delete $h{'%h'};
 551
 552     } else {
 553         $h{'%h'} =~ s/\\!/!/g;
 554     }
 555
 556     if( $h{'%h'} ) {
 557         $h{'%u'} = $1 if $h{'%h'} =~ s/^([\w.\-_]+)@//;
 558         $h{'%h'} =~ s/\\@/@/g;
 559     }
 560
 561     @c = map {exists $h{$_} ? $h{$_} : $_}
 562          map { m/^\[([^\[\]]+)\]/ ? ($h{$1} ? do{s/^\[\Q$1\E\]//; $_} : ()) : ($_) } ## no critic: why on earth not?
 563          map { s/\[\]\%(\w+)/[\%$1]\%$1/; $_ }                                       ## no critic: why on earth not?
 564          @c;
 565
 566     if( $this->{debug} ) {
 567         local $" = ")(";
 568         $this->std_msg($host, $h{'%n'}, 0, BOLD.BLACK."DEBUG: exec(@c)");
 569     }
 570
 571     return @c;
 572 }
 573 # }}}
 574 # start_queue_on_host {{{
 575 sub start_queue_on_host {
 576     my ($this, $kernel => $host, $cmdno, $cmd, @next) = @_;
 577
 578     my $program;
 579     if( ref($this->{_shell_cmd}) eq "CODE" ) {
 580         $program = sub {
 581             $this->set_subst_vars( '%h' => $host, '%n' => $cmdno );
 582             $this->{_shell_cmd}->(@$cmd)
 583         };
 584
 585     } else {
 586         $this->set_subst_vars( '%h' => $host, '%n' => $cmdno );
 587         $program = [ $this->subst_cmd_vars(@{$this->{_shell_cmd}} => @$cmd) ];
 588     }
 589
 590     my $kid = POE::Wheel::Run->new(
 591         Program     => $program,
 592         StdoutEvent => "child_stdout",
 593         StderrEvent => "child_stderr",
 594         CloseEvent  => "child_close",
 595     );
 596
 597     $kernel->sig_child( $kid->PID, "child_signal" );
 598
 599     my $lineno = 0;
 600     my $info = [ $kid, $host, $cmdno, \$lineno, @next ];
 601     $this->{_wid}{ $kid->ID } = $this->{_pid}{ $kid->PID } = $info;
 602
 603     return;
 604 }
 605 # }}}
 606
 607 # poe_start {{{
 608 sub poe_start {
 609     my $this = shift;
 610
 611     my %starting;
 612     my @hosts = keys %{ $this->{_cmd_queue} };
 613     for my $host (@hosts) {
 614         my @c = @{ $this->{_cmd_queue}{$host} };
 615
 616         $this->start_queue_on_host($_[KERNEL] => $host, 1, @c);
 617         push @{$starting{"@{$c[0]}"}}, $host;
 618     }
 619
 620     for my $message (keys %starting) {
 621         my @hosts = @{ $starting{$message} };
 622
 623         if( @hosts == 1 ) {
 624             $this->std_msg($this->_host_route_to_nick($hosts[0]), 1, 0, BOLD.BLACK."-- starting: $message --");
 625
 626         } else {
 627             $this->std_msg("", 1, 0, BOLD.BLACK."-- starting: $message on @hosts --");
 628         }
 629     }
 630
 631     delete $this->{_cmd_queue};
 632
 633     return;
 634 }
 635 # }}}
 636
 637 1;