ljr/wcmtools/spud/bin/replicator

#!/usr/bin/perl
#
# Danga's Statistics Server Replicator
# Very lightweight program that replicates data from one statistic server (subscribe *)
# to another one.  Very, very simple and lightweight.  Uses open3() to use the ssh command
# to connect to a remote statistics server.
#
# Command line options:
#
#    -c STRING  set what config file to use for options
#    -D         if present, tell the server to daemonize
#
# Configuration file format:
#
#    server = STRING          location of statistics server
#    port = INT               port number server listens on
#    ssh_host = STRING        host of remote SSH to tunnel through
#    ssh_port = INT           port SSH daemon is on
#    ssh_key = STRING         filename to use for our private key
#    ssh_user = STRING        username to identify as to SSH server
#
# Copyright 2004, Danga Interactive
#
# Authors:
#   Mark Smith <marksmith@danga.com>
#
# License:
#   undecided.
#

# uses
use strict;
use IO::Socket;
use IO::Select;
use IPC::Open3;
use Getopt::Long;
use Carp;
use POSIX ":sys_wait_h";
use Socket qw(PF_INET IPPROTO_TCP SOCK_STREAM);
use Danga::Socket;

# Command-line options will override
my ($daemonize, $conf_file);
Getopt::Long::Configure( "bundling" );
Getopt::Long::GetOptions(
    'D|daemon'   => \$daemonize,
    'c|config=s' => \$conf_file,
);
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
my %config;
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') { $config{to}{host} = $r; }
        elsif ($l eq 'port') { $config{to}{port} = $r+0; }
        elsif ($l eq 'ssh_key') { $config{from}{key} = $r; }
        elsif ($l eq 'ssh_host') { $config{from}{host} = $r; }
        elsif ($l eq 'ssh_port') { $config{from}{port} = $r+0; }
        elsif ($l eq 'ssh_user') { $config{from}{username} = $r; };
    }
}
close FILE;

# daemonize?
daemonize() if $daemonize;

# connect to stats server we're replicating TO.  messy.
my $dsock = new_connection()
    or die "Can't get initial socket... exiting.\n";

# now setup our ssh open3 sockets
my ($reader, $writer);
my $sshpid = open3($writer, $reader, $reader, 'ssh', '-C', '-tt', '-i', $config{from}{key},
                   '-p', $config{from}{port}, '-l', $config{from}{username}, $config{from}{host});
print $writer "sub *\n";

# always kill off our SSH connection
$SIG{TERM} = sub { kill 15, $sshpid; exit 0; };
$SIG{INT} = sub { kill 15, $sshpid; exit 0; };

# variables used later
my $readbuf;

# this is our main reading loop
my $sobj = new IO::Select;
$sobj->add($reader);

# post event loop
my $postloop = sub {
    # check if somehow ssh died on us? :-/
    my $pid = waitpid -1, WNOHANG;
    if ($pid == $sshpid) {
        # sleep a few seconds and try to spawn a new one
        $sshpid = 0;
        while (!$sshpid) {
            print "Lost SSH connection... sleeping 5 seconds.\n";
            sleep 5;
            $sshpid = open3($writer, $reader, $reader, 'ssh', '-C', '-tt', '-i', $config{from}{key},
                            '-p', $config{from}{port}, '-l', $config{from}{username}, $config{from}{host});
            print $writer "sub *\n";
        }
    }

    # if our parent socket is closed...
    unless ($dsock && !$dsock->{closed}) {
        # create a new one if we can
        print "Lost SPUD connection... reconnecting...\n";
        $dsock = new_connection();
        unless ($dsock) {
            print "\tUnable to connect... pausing a second.\n";
            sleep 1;
            return 1;
        }
    }

    # see if we can read from our socket yet
    my @ready = $sobj->can_read(0.1);
    return 1 unless @ready;

    # must be ready to read
    my $bytes = sysread $reader, $readbuf, 1024, length $readbuf;
    while ($readbuf =~ s/(.+?)\r?\n//) {
        my $line = $1;
        next unless $line =~ /^set/i;
        $dsock->write("Q$line\r\n");
    }
    return 1;
};

# now configure the client
Client->SetLoopTimeout(100); # 100 milliseconds timeout
Client->SetPostLoopCallback($postloop); # have it call us

# now run the event loop
Client->EventLoop();

# kill off our child too
kill 15, $sshpid;
print "replicator terminating\n";

# 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/\s+$//;
    return $res;
}

# connect anew to the SPUD server we're replicating to
# NOTE: can return undef if we can't get to the server!
sub new_connection {
    my $sock;
    socket $sock, PF_INET, SOCK_STREAM, IPPROTO_TCP;
    die "error: can't make socket\n"
        unless $sock && defined fileno($sock);
    IO::Handle::blocking($sock, 0);
    connect $sock, Socket::sockaddr_in($config{to}{port}, Socket::inet_aton($config{to}{host}));
    my $dsock = Client->new($sock)
        or return undef;
    $dsock->watch_write(1);
    $dsock->watch_read(1);
    return $dsock;
}

###########################################################################
### Client class for use in processing input/output
package Client;

use base "Danga::Socket";

sub event_read {
    # read and toss, we don't care about input from the user here
    my Client $self = $_[0];
    my $bref = $self->read;
}

sub event_err {
    # connection died?
    my Client $self = $_[0];
    $self->close('event_err');
}

sub event_hup {
    # connection to server died...
    my Client $self = $_[0];
    $self->close('event_hup');
}