#!/usr/bin/perl
#
# Danga's Statistics Wrapper
# This little program simply runs a command and pipes the output to a
# statistics server so that you can see the output of the command.  Data
# will be stored in $KEY.output and $KEY.status, so you can determine
# exactly how long a run took.
#
# Command line options:
#
#    -s STRING    server to connect to
#    -p INT       port statistics server is running on
#    -k STRING    key to put data in ($KEY.status, $KEY.output)
#    -e STRING    command to execute
#    -c STRING    file to load config data from
#
# Configuration file format:
#
#    server = STRING     location of server
#    port = INT          port of server
#
# Copyright 2004, Danga Interactive
#
# Authors:
#   Mark Smith <marksmith@danga.com>
#
# License:
#   undecided.
#

# uses
use strict;
use IO::Socket;
use IPC::Open3;
use Getopt::Long;
use Carp;
use POSIX;

# Command-line options will override
my ($daemonize, $conf_file, $key, $host, $port, $cmd);
Getopt::Long::Configure( "bundling" );
Getopt::Long::GetOptions(
    'D|daemon'    => \$daemonize,
    'c|config=s'  => \$conf_file,
    'k|key=s'     => \$key,
    'p|port=i'    => \$port,
    's|server=s'  => \$host,
    'e|exec=s' => \$cmd,
);

# if we don't have all the data we need, parse a config file
die "No key (-k KEY) specified.\n" unless $key;
unless ($port && $host) {
    die "You must at least specify --config=FILENAME for me to work.\n" unless $conf_file;
    die "File '$conf_file' doesn't exist for configuration.\n" unless -e $conf_file;

    # parse the config file
    open FILE, "<$conf_file"
        or die "Unable to open config file: $!\n";
    while (my $line = <FILE>) {
        if ($line =~ /^\s*([^#].*)\s*=\s*(.*)\s*/) {
            my ($l, $r) = (trim($1), trim($2));
            if ($l eq 'server') { $host = $r; }
            elsif ($l eq 'port') { $port = $r+0; }
        }
    }
    close FILE;
}

# if cmd was specified, use that
@ARGV = split /\s+/, $cmd if $cmd;

# error checking
die "No server to connect to.\n" unless $port && $host;
die "No arguments to run.\n" unless @ARGV;
die "File to run ($ARGV[0]) doesn't seem to exist.\n" unless -e $ARGV[0];

# daemonize?
daemonize() if $daemonize;

# connect to stats server we're replicating TO
my $socket = IO::Socket::INET->new(Proto => 'tcp',
                                   PeerAddr => $host,
                                   PeerPort => $port,)
    or die "Unable to connect to local stats server ($host:$port): $!\n";
$socket->autoflush(1);

# now setup our open3 socket
my ($reader, $writer);
my $procpid = open3($writer, $reader, $reader, @ARGV);

# parent pid
my $ppid = $$;

# now fork
my $kidpid;
die "Unable to fork: $!\n" unless defined($kidpid = fork());

# basic parts that handle reading/writing from the two sides
if ($kidpid) {
    # parent
    while (defined (my $line = <$socket>)) {
        # we do this to flush the read buffer, but we don't need to
        # echo this information to anybody, as all of this is going
        # to be in the form of "OK" and such saying that data got set
    }
    kill("TERM" => $kidpid); # death to the children
    waitpid $kidpid, 0;
} else {
    # child
    print $socket "QSET $key.status started\n";
    while (defined (my $line = <$reader>)) {
        # take the output and set it
        $line = trim($line);
        print $socket "QSET $key.output $line\n";
    }
    print $socket "QSET $key.status finished\n";
    $socket->close;

    # now kill off our parent
    kill("TERM" => $ppid);
}

# daemonizer routine
sub daemonize {
    my($pid, $sess_id, $i);

    ## Fork and exit parent
    if ($pid = fork) { exit 0; }

    ## Detach ourselves from the terminal
    croak "Cannot detach from controlling terminal"
        unless $sess_id = POSIX::setsid();

    ## Prevent possibility of acquiring a controling terminal
    $SIG{'HUP'} = 'IGNORE';
    if ($pid = fork) { exit 0; }

    ## Change working directory
    chdir "/";

    ## Clear file creation mask
    umask 0;

    ## Close open file descriptors
    close(STDIN);
    close(STDOUT);
    close(STDERR);

    ## Reopen stderr, stdout, stdin to /dev/null
    open(STDIN,  "+>/dev/null");
    open(STDOUT, "+>&STDIN");
    open(STDERR, "+>&STDIN");
}

# little trimming sub
sub trim {
    my $res = shift;
    $res =~ s/^\s+//;
    $res =~ s/[\r\n\s]+$//;
    return $res;
}