pogo-rexec

#!/usr/bin/env perl

# Copyright (c) 2010-2011 Yahoo! Inc. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

use 5.008;
use common::sense;

use Crypt::OpenSSL::RSA qw();
use File::Slurp qw(read_file);
use File::Temp qw();
use Getopt::Long qw(GetOptions);
use IPC::Open3 qw(open3);
use JSON qw(decode_json);
use Log::Log4perl qw(:easy);
use MIME::Base64 qw(decode_base64);
use POSIX qw(WEXITSTATUS mkfifo);
use Digest::SHA qw(sha1_hex);

#Register all signals if they exist, then clean up the private key file
$SIG{HUP} = $SIG{INT} = $SIG{TERM} = $SIG{__DIE__} = \&cleanup;

BEGIN
{
  File::Temp->safe_level(File::Temp::HIGH);
}

my $opts = {
  expect_wrapper => "pogo-pw",
  worker_key     => undef,
  scp_options    => [
    qw(-q -o RSAAuthentication=no -o StrictHostKeyChecking=no -o LogLevel=error -o GlobalKnownHostsFile=/dev/null -o UserKnownHostsFile=/dev/null -o ForwardAgent=yes)
  ],
  ssh_options => [
    qw(-o RSAAuthentication=no -o StrictHostKeyChecking=no -o LogLevel=error -o GlobalKnownHostsFile=/dev/null -o UserKnownHostsFile=/dev/null -o ForwardAgent=yes -t)
  ],
  command_prefix => [],
  tempdir =>
    "/tmp",    #The default directory where the named pipe will be created for storing private key.
};

my $remote_env = { PATH => '/usr/local/bin:/usr/local/sbin:/usr/bin:/usr/sbin:/bin:/sbin', };

my $pkfile        = undef;
my $ssh_agent_pid = undef;
my $ssh_sock      = undef;

sub cleanup
{
  my $msg = shift;
  if ( $msg =~ m/^[A-Z]+$/ ) { $msg = 'got SIG' . $msg; }
  if ($pkfile)
  {
    unlink $pkfile;
  }
  &kill_agent;
  LOGDIE "ERROR: pogo-rexec died: $msg\n\n";
}

sub kill_agent
{
  #Kill the ssh-agent if one was created
  if ( $ENV{SSH_AGENT_PID} )
  {
    my $agent_command = ("eval `ssh-agent -k`");

    #start the agent
    my ( $writer, $reader ) = ( IO::Handle->new, IO::Handle->new );
    my $pid;
    eval { $pid = open3( $writer, $reader, undef, $agent_command ); };
    if ($@)
    {
      LOGDIE "Error: $@";
    }

  }

}

sub remote_env
{
  my @env = qw(env);
  push @env, map { $_ . '=' . $remote_env->{$_} } sort keys %{$remote_env};
  return \@env;
}

sub execute
{
  my $args = shift;
  my @command = ( $opts->{expect_wrapper}, @_ );

  my ( $writer, $reader ) = ( IO::Handle->new, IO::Handle->new );
  my $pid;
  eval { $pid = open3( $writer, $reader, undef, @command ); };
  if ($@)
  {
    LOGDIE "Error: $@";
  }

  # Send timeout
  $writer->print( $args->{timeout}, "\n" );

  if ( !$args->{password} )
  {
    # Send blank password
    $writer->printf( "%s", pack( q{u}, "\n" ) );
  }
  else
  {
    # Send password
    $writer->printf( "%s",
      pack( q{u}, $opts->{private_key}->decrypt( decode_base64( $args->{password} ) ) ) );
  }

  if ( !$args->{pvt_key_passphrase} )
  {
    # Send blank passphrase
    $writer->printf( "%s", pack( q{u}, "\n" ) );
  }
  else
  {
    # Send passphrase
    $writer->printf( "%s",
      pack( q{u}, $opts->{private_key}->decrypt( decode_base64( $args->{pvt_key_passphrase} ) ) ) );
  }
  $writer->close;

  # Copy remote command's output to ours.
  while ( defined( my $line = $reader->getline() ) )
  {
    STDOUT->print($line);
  }

  # Reap child and return its exit status to our caller.

  my $p = waitpid( $pid, 0 );
  return WEXITSTATUS($?);
}

sub main
{
  my ( $args, $tmpfile, $exit_status );

  STDERR->autoflush(1);
  STDOUT->autoflush(1);

  # Override default option values with those in configuration file.
  if ( -r $ENV{"POGO_WORKER_CONFIG_FILE"} )
  {
    $opts = merge( $opts, YAML::LoadFile( $ENV{"POGO_WORKER_CONFIG_FILE"} ) );
  }

  if ( defined $opts->{rexec_log4perl} && -r $opts->{rexec_log4perl} )
  {
    Log::Log4perl::init( $opts->{log4perl} );
  }
  if ( defined $opts->{loglevel} )
  {
    Log::Log4perl::get_logger()->level( $opts->{loglevel} );
  }

  GetOptions(
    'e|expect-wrapper=s' => sub { $opts->{expect_wrapper} = $_[1] },
    'k|ssl-key=s'        => sub { $opts->{worker_key}     = $_[1] },
  );

  # Sanity checks
  unless ( -x $opts->{expect_wrapper} )
  {
    LOGDIE "Can't execute " . $opts->{expect_wrapper};
  }

  unless ( -r $opts->{worker_key} )
  {
    LOGDIE "Can't read private key " . $opts->{worker_key};
  }

  $opts->{private_key} =
       Crypt::OpenSSL::RSA->new_private_key( scalar read_file $opts->{worker_key} )
    || LOGDIE "unable to load worker private key!";

  # Read JSON-encoded args from STDIN.
  $args = decode_json( STDIN->getline() );

  # Check for required arguments
  for (qw(job_id command user run_as host timeout))
  {
    unless ( defined $args->{$_} )
    {
      LOGDIE "Missing required argument $_\n";
    }
  }

  #Check to make sure atleast one form of authentication is provided
  if ( !$args->{client_private_key} && !$args->{password} )
  {
    LOGDIE "Missing authentication parameters\n";
  }

  #If it is using private key authenticaion, set up ssh-agent
  if ( $args->{client_private_key} )
  {
    #prepare the command to start ssh-agent
    my $agent_command =
      'eval `ssh-agent` ; echo SSH_AUTH_SOCK=$SSH_AUTH_SOCK ; echo SSH_AGENT_PID=$SSH_AGENT_PID ; ';

    #start the agent
    my ( $writer, $reader ) = ( IO::Handle->new, IO::Handle->new );
    my $pid;
    eval { $pid = open3( $writer, $reader, undef, $agent_command ); };
    if ($@)
    {
      LOGDIE "Error: $@";
    }

    #fetch the socket and pid
    while ( my $line = $reader->getline )
    {
      $ssh_agent_pid = $1 if ( $line =~ m/SSH_AGENT_PID=(\d+)$/ );
      $ssh_sock      = $1 if ( $line =~ m/SSH_AUTH_SOCK=(.+)$/ );
    }

    #set the socket for this program and all its children
    $ENV{SSH_AUTH_SOCK} = $ssh_sock;
    $ENV{SSH_AGENT_PID} = $ssh_agent_pid;

    #add the key to the agent using ssh-add
    #decrypt the private key and send it to the named pipe
    my $pk_data;
    for my $data ( @{ $args->{client_private_key} } )
    {
      my $dkdata = $opts->{private_key}->decrypt( decode_base64($data) );
      $pk_data .= $dkdata;

    }

    # create a named pipe and make sure it does not already exists
    $pkfile = "$opts->{tempdir}/" . sha1_hex( sprintf "%0.5f-%d-%d", rand(), $$, time );
    while ( -p $pkfile )
    {
      $pkfile = "$opts->{tempdir}/" . sha1_hex( printf "%0.5f-%d-%d", rand(), $$, time );
    }

    # the named pipe blocks till some one reads and ssh-add also blocks till its completed,
    # so both these will deadlock. So we fork the process here, let the pipe run in the child,
    # start the ssh-add in the parent and once they are completed, we merge them back
    my $pipepid = fork;
    if ( not defined $pipepid )    # not enough resources
    {
      LOGDIE "Resources not available for forking.\n";
    }
    elsif ( $pipepid == 0 )        # Child
    {
      # ssh-add first adds without passphrase and then tries next time with passphrase
      # since fifo is read once only, if the given private key does not have a passphrase,
      # this writing has to be only once, if it has a correct passphrase, it has to be
      # done twice. So based on passphrase, we write it once or twice. If the passphrase is
      # incorrect, this process will exit with writing twice and the parent will keep
      # reading multiple times. EVentually the system timeout will take over and exit.

      my $write_count = $args->{pvt_key_passphrase} ? 2 : 1;
      for ( my $i = 0; $i < $write_count; $i++ )
      {
        unless ( -p $pkfile )
        {
          unlink($pkfile);
          mkfifo( $pkfile, 0600 ) or LOGDIE "can't mkfifo $pkfile: $!";
        }

        # next line blocks until there's a reader
        open( FIFO, "> $pkfile" ) or LOGDIE "can't write $pkfile: $!";
        print FIFO $pk_data;
        close(FIFO);
        sleep 1;    # to avoid dup signals and let ssh-add catch up

      }
      exit(0);
    }
    else            #parent
    {
      #run ssh-add with the private key
      execute( $args, "ssh-add", "$pkfile" );
      waitpid( $pipepid, 0 );    #wait for child

    }

    #Unlink the pk file
    if ($pkfile)
    {
      unlink $pkfile;
    }

    push @{ $opts->{scp_options} }, ( "-o", "PubkeyAuthentication=yes" );
    push @{ $opts->{ssh_options} }, ( "-o", "PubkeyAuthentication=yes" );

  }
  else
  {
    push @{ $opts->{scp_options} }, ( "-o", "PubkeyAuthentication=no" );
    push @{ $opts->{ssh_options} }, ( "-o", "PubkeyAuthentication=no" );
  }
  # split if hostname:rootname combination is present
  my @target;
  if ( $args->{host} =~ /:/ )
  {
    LOGDIE "Rootname present but missing command_root_transform property"
      unless ( $args->{command_root_transform} );
    @target = split( ":", $args->{host} );
    $args->{host} = $target[0];
  }

  if ( defined $args->{command_root_transform} )
  {
    LOGDIE "\${command} substring not found in $args->{command_root_transform} \n"
      unless $args->{command_root_transform} =~ /\${command}/;

    LOGDIE "\${rootname} substring not found in $args->{command_root_transform} \n"
      unless $args->{command_root_transform} =~ /\${rootname}/;
  }

  # insert POGO_JOB_ID into the remote environment
  $remote_env->{POGO_JOB_ID} = $args->{job_id};

  if ( ${ $args->{command} }[0] =~ /^POGOATTACHMENT!(.+)/s )
  {

    # Initial step: copy the attachment to our destination host.
    my $attachment = decode_base64($1);
    $tmpfile = new File::Temp();
    chmod( 0700, $tmpfile );
    $tmpfile->print($attachment);
    $tmpfile->close;
    execute( $args, "scp", @{ $opts->{scp_options} },
      $tmpfile, $args->{user} . "@" . $args->{host} . ":" . $tmpfile )
      and LOGDIE "scp failed";

    # $args->{command} = [$tmpfile->filename];
    # This would be more convenient:
    $args->{command}->[0] = $tmpfile->filename;

    # that way we can pass args to our script
  }

  # substitute the command_root_transform with the rootname
  # and the name of the pogo worker stub
  if ( defined $args->{command_root_transform} && defined $target[1] )
  {
    $args->{command_root_transform} =~ s/\${rootname}/$target[1]/g;
    $args->{command_root_transform} =~ s/\${command}/$args->{command}->[0]/g;
    $args->{command}->[0] = $args->{command_root_transform};
  }

  # Prepend "sudo -u <user>" if we are running as someone else
  if ( $args->{run_as} ne $args->{user} )
  {
    unshift( @{ $args->{command} }, "sudo", "-u", $args->{run_as} );
  }

  # Execute the provided command
  $exit_status = execute(
    $args, "ssh", @{ $opts->{ssh_options} },
    "--",
    $args->{user} . "@" . $args->{host},
    @{ remote_env() },
    @{ $opts->{command_prefix} },
    @{ $args->{command} }
  );

  # Remove the copied attachment on the target host, if necessary.
  if ($tmpfile)
  {
    execute(
      $args, "ssh", @{ $opts->{ssh_options} },
      "--", $args->{user} . "@" . $args->{host},
      "rm -f", $tmpfile
    );
  }

  &kill_agent;

  return $exit_status;
}

exit main();

1;

=pod

=head1 NAME

  CLASSNAME - SHORT DESCRIPTION

=head1 SYNOPSIS

CODE GOES HERE

=head1 DESCRIPTION

LONG_DESCRIPTION

=head1 METHODS

B<methodexample>

=over 2

methoddescription

=back

=head1 SEE ALSO

L<Pogo::Dispatcher>

=head1 COPYRIGHT

Apache 2.0

=head1 AUTHORS

  Andrew Sloane <andy@a1k0n.net>
  Michael Fischer <michael+pogo@dynamine.net>
  Mike Schilli <m@perlmeister.com>
  Nicholas Harteau <nrh@hep.cat>
  Nick Purvis <nep@noisetu.be>
  Robert Phan <robert.phan@gmail.com>

=cut

# vim:syn=perl:sw=2:ts=2:sts=2:et:fdm=marker