#!/usr/bin/perl
# This script launch a oarstat command to retreive all data about the given job.
# Then all nodes are contacted and their /etc/security/access_cluster.conf
# are populated to enable the ssh connection to this OAR job owner if the whole
# nodes are allocated.
#
# Typically you have to write in oar.conf:
#
# SERVER_PROLOGUE_EXEC_FILE="/etc/oar/oar_server_proepilogue.pl START"
# SERVER_EPILOGUE_EXEC_FILE="/etc/oar/oar_server_proepilogue.pl STOP"
#
use Fcntl;

my $Taktuk_cmd    = "/usr/bin/taktuk -s -c 'ssh'";
my $Oarstat_cmd   = "oarstat";
my $Security_file = "/etc/security/access_cluster.conf";

my $Job_id = $ARGV[1];
if ($Job_id !~ m/^\d+$/m) {
    warn("ERROR: cannot get the job id\n");
    exit(1);
}

my $struct;
if (open(OARSTAT, "$Oarstat_cmd -fj $Job_id -D |")) {
    $struct = eval(<OARSTAT>);
    close(OARSTAT);
    if (!defined($struct->{state})) {
        warn("ERROR: Cannot read correctly the result of the command $Oarstat_cmd\n");
        exit(4);
    }
} else {
    warn("ERROR: Cannot execute command $Oarstat_cmd\n");
    exit(3);
}

if ($ARGV[0] eq "START") {
    my $str = '
my $file_str = "# File generated by OAR.\n";
if (! system("diff /dev/cpuset/oar/' .
      $struct->{cpuset_name} . '/cpus /dev/cpuset/cpus > /dev/null 2>&1")){
    $file_str .= "+:' . $struct->{owner} . ':ALL\n";
}
$file_str .= "-:ALL:ALL\n";
if (open(ACCESS, "> ' . $Security_file . '")){
    print(ACCESS "$file_str");
    close(ACCESS);
}else{
    exit(1);
}
';
    my @tmp = diffuse_cmd($struct->{assigned_network_address}, $str, $Job_id);
    if ($#tmp >= 0) {
        warn("ERROR: Some nodes are inaccessible @tmp\n");
        exit(5);
    }
} elsif ($ARGV[0] eq "STOP") {
    my $str = '
if (open(ACCESS, "> ' . $Security_file . '")){
    print(ACCESS "# File generated by OAR.\n");
    print(ACCESS "-:ALL:ALL\n");
    close(ACCESS);
}else{
    exit(1);
}
if (! system("diff /dev/cpuset/oar/' .
      $struct->{cpuset_name} . '/cpus /dev/cpuset/cpus > /dev/null 2>&1")){
#    system("su - ' . $struct->{owner} . ' -c \'kill -9 -1\'");
}
';
    my @tmp = diffuse_cmd($struct->{assigned_network_address}, $str, $Job_id);
    if ($#tmp >= 0) {
        warn("ERROR: Some nodes are inaccessible @tmp\n");
        exit(5);
    }
} else {
    warn("ERROR: unknown first argument $ARGV[0]\n");
    exit(2);
}

sub diffuse_cmd($$$) {
    my $hosts       = shift;
    my $perl_script = shift;
    my $job_id      = shift;

    pipe(tak_node_read,   tak_node_write);
    pipe(tak_stdin_read,  tak_stdin_write);
    pipe(tak_stdout_read, tak_stdout_write);
    my $pid = fork;
    if ($pid == 0) {

        #CHILD
        $SIG{CHLD} = 'DEFAULT';
        $SIG{TERM} = 'DEFAULT';
        $SIG{INT}  = 'DEFAULT';
        $SIG{QUIT} = 'DEFAULT';
        $SIG{USR1} = 'DEFAULT';
        $SIG{USR2} = 'DEFAULT';
        my $cmd =
          "$Taktuk_cmd " . '-o status=\'"STATUS $host $line\n"\'' . " -f '<&=" .
          fileno(tak_node_read) .
          "' broadcast exec [ sudo perl - SERVER_SCRIPT ], broadcast input file [ - ], broadcast input close";
        fcntl(tak_node_read, F_SETFD, 0);
        close(tak_node_write);
        close(tak_stdout_read);
        close(STDOUT);

        # Redirect taktuk output into the pipe
        open(STDOUT, ">& tak_stdout_write");

        # Use the child STDIN to send the user command
        close(tak_stdin_write);
        close(STDIN);
        open(STDIN, "<& tak_stdin_read");

        exec($cmd);
        warn("[ERROR] Cannot execute $cmd\n");
        exit(-1);
    }
    close(tak_node_read);
    close(tak_stdin_read);
    close(tak_stdout_write);

    my $tmp_node_hash;

    # Send node list
    foreach my $n (@{$hosts}) {
        $tmp_node_hash{$n} = 1;
        print(tak_node_write "$n\n");
    }
    close(tak_node_write);

    eval {
        $SIG{ALRM} = sub { die "alarm\n" };
        alarm(50);

        # Send data structure to all nodes
        print(tak_stdin_write $perl_script);
        close(tak_stdin_write);

        # Check good nodes from the stdout taktuk
        while (<tak_stdout_read>) {
            if ($_ =~ /^STATUS ([\w\.\-\d]+) (\d+)$/) {
                if ($2 == 0) {
                    delete($tmp_node_hash{$1}) if (defined($tmp_node_hash{$1}));
                }
            } else {
                print("[TAKTUK OUTPUT] $_");
            }
        }
        close(tak_stdout_read);
        alarm(0);
    };
    @bad = keys(%tmp_node_hash);
    return (@bad);
}

