#!/usr/bin/perl
# Filename:	serverizer
# Author:	David Ljung Madison <DaveSource.com>
# See License:	http://MarginalHacks.com/License
# Version:	1.04
#
# Description:	Run a program in the background as a tcp server
#		This is handy for programs that take a long time to startup.
#		You can "serverize" the program and then use serverizer to
#		connect to the server and run commands.
#
# Limitations:	Not multi-threaded!
#		Command input and output buffering is line-by-line, not by char
#		But line-by-line means that each command must return one
#		  (and only one!) prompt.  If you can figure out an efficient
#		  way to make this char-by-char, be my guest.
# Examples:
#
# Run a serverizer to telnet (must have a prompt that matches the default)
# % serverizor -watch='^login: '==$USER -watch=^Password:==some_pass
#
# Run a serverizer to an ftp client
# % serverizor -daemon="ftp localhost" -prompt="ftp> " -watch='^Name \(.*\): '==$USER -watch=^Password:==some_pass
#
# Once you've started up the server, you can then disconnect with "."
# and then other serverizors (as clients) will be able to connect.
# 
# When you're done, exit through serverizor (exit, quit, whatever...), or else:
# % serverizor -stop_server
# 
use strict;

use Expect;
use IO::Socket;
use Net::hostent;		# for OO version of gethostbyaddr
use POSIX qw(setsid);		# for daemonize
use Fcntl ':flock';		# for pid file

##################################################
# Setup the variables
##################################################
my $PROGNAME = $0;
$PROGNAME =~ s|.*/||;

my $daemon = "telnet localhost";
my $PORT = 7042;

# Options are:  'name' => [ <default_val>, <list?>, <usage> ],
# <list?> is 0->scalar, 1->list, {sub}->subroutine handled
my %OPTIONS = (
	'daemon'		=> [ $daemon, 0,
		"Set the program to serverize/daemonize" ],
	'daemon_args'		=> [ undef, 1,
		"List of args for the daemon" ],
	'prompt'		=> [ '\d+% ', 0,	# That's my prompt..
		"Daemon's prompt (can be regexp)" ],
	'exit'			=> [ "exit", 0,
		"The exit command for the daemon" ],
	'disconnect'		=> [ '.', 0,
		"Disconnect without quitting the daemon (or ^D)" ],
	'startup_timeout'	=> [ 300, 0,
		"Timeout for the daemon to reach the first prompt" ],
	'line_timeout'		=> [ 60, 0,
		"Timeout to wait for a line of output" ],
	'host'			=> [ "localhost", 0,
		"Host that the daemon server is on" ],
	'port'			=> [ $PORT, 0,
		"Port for the daemon server" ],
	'debug'			=> [ 0, 0,
		"Debug mode (server prints out I/O)" ],
	'start_server'		=> [ 0, 0,
		"Just start up the daemon program" ],
	'stop_server'		=> [ 0, 0,
		"Stop the daemon server" ],
	'pid_file'		=> [ "/tmp/pid.$PROGNAME-$PORT", 0,
		"Location of pid file" ],
	'watch'		=> [ undef, \&add_watch_prompt,
		"Create a watch prompt (-watch=prompt==send)" ],
	'watch_wait'		=> [ 0, 0,
		"Wait before sending watch prompts" ],
	'ignore_whitespace'	=> [ 0, 0,
		"Whitespace commands won't show a new prompt" ],
	'comment'	=> [ 0, 0,
		"Character that starts a comment (for -ignore_whitespace)" ],
	'reset_daemon'	=> [ undef, 1,
		"Cleanup commands that are run after each disconnect" ],
	);


##################################################
# Usage
##################################################
sub usage {
  my $msg;
  foreach $msg (@_) { print "ERROR:  $msg\n"; }
  print "\n";
  print "Usage:\t$PROGNAME [-d] <file>\n";
  print "\tDoes something to the given file\n";
  foreach my $opt ( sort keys %OPTIONS ) {
    my $default = $OPTIONS{$opt}[0];
    my $default = $OPTIONS{$opt}[1] ? "($default)" : "[$default]";
    printf "  -%-20s $OPTIONS{$opt}[2] $default\n",$opt;
  }
  print "\n";
  print "All options are set by '-opt=val'\n";
  print "\n";
  print "Watch prompts look for a given prompt and send the response string\n";
  print "You can have as many watch prompts as you want.  Useful for passwords.\n";
  print "\n";
  exit -1;
}

sub add_watch_prompt {
  my ($opt,$arg,$val) = @_;
  my ($prompt,$send) = split('==',$val,2);
  usage("watch prompts are of the form 'prompt==send'") unless defined $send;
  push(@{$opt->{watch_prompts}},[$prompt , sub { watch_prompt($opt,$send); exp_continue; } ]);
}

my @saved_argv = @ARGV;
sub parse_args {
  my %opt;

  # Set defaults
  map($opt{$_}=$OPTIONS{$_}[0], keys %OPTIONS);
  undef $opt{pid_file};

  while (my $arg=shift(@ARGV)) {
    if ($arg =~ /^-h$/) { usage(); }
    if ($arg =~ /^-d$/) { $opt{debug}=1; next; }
    if ($arg =~ /^-([^=]+)(=(.+))?$/) {
      my ($var,$val) = ($1, $2 ? $3 : 1);
      usage("Unknown option [$var]") unless exists $opt{$var};
      my $type = $OPTIONS{$var}[1];
      if (ref $type eq "CODE") {
        &$type(\%opt,$var,$val);
      } elsif ($type == 1) {
        push(@{$opt{$var}},$val);
      } else {
        $opt{$var}=$val;
      }
      next;
    }
    usage("Unknown argument [$arg]");
  }

  # Pid file is generated from other options (unless set)
  $opt{pid_file} = $opt{pid_file} || "/tmp/pid.$PROGNAME-$opt{port}";

  # For -ignore_whitespace
  $opt{not_command} = $opt{comment} ? "^\s*($opt{comment}.+)?\$" : "^\s*\$";
  study $opt{not_command};

  \%opt;
}

##################################################
# Daemon handling code
##################################################
# Start up the command that we want to daemonize
sub start_daemon {
  my ($opt) = @_;

  my $servex = Expect->spawn($opt->{daemon},@{$opt->{daemon_args}});
 	# We only need multiline matching if the prompt has newlines - turn
 	# it off otherwise because it's probably slower
	#$Expect::Multiline_Matching= ($opt->{prompt}=~/[\r\n]/) ? 1 : 0;
  $Expect::Multiline_Matching= 1;
  $servex->log_group(0);
  $servex->log_stdout(0);
  $opt->{ex} = $servex;

  # Wait for the first prompt
  send_command($opt,undef);
}

# Shut it down
sub stop_daemon {
  my ($opt) = @_;
  print "[Shutting down server]\n" if ($opt->{debug});
  print {$opt->{ex}} $opt->{exit}."\n"
    unless defined($opt->{ex}->exp_exitstatus());
  $opt->{ex}->soft_close();
}

# Report to the client what we've seen
my $LAST_PROMPT;
sub just_saw {
  my ($opt) = @_;
  $LAST_PROMPT = $opt->{ex}->exp_match();	# Assume it was the prompt
  my $got = $opt->{ex}->exp_before().$opt->{ex}->exp_match();
  print STDERR "//$got//" if $opt->{debug};
  print {$opt->{client}} $got if $opt->{client};
  #$opt->{ex}->clear_accum();
}

# Send a command to the daemon
sub just_send {
  my ($opt,$cmd) = @_;
  return unless defined($cmd);
  print {$opt->{ex}} $cmd."\n" if defined $cmd;
}

# We saw a watch prompt
sub watch_prompt {
  my ($opt,$send) = @_;
  just_saw($opt);
  sleep $opt->{watch_wait};
  just_send($opt,$send);
}

# Send a command to the daemon, send results to the client
sub send_command {
  my ($opt,$cmd) = @_;

  just_send($opt,$cmd);
  return 1 if (defined $cmd && $cmd =~ /$opt->{not_command}/ && $opt->{ignore_whitespace});

  my $ex = $opt->{ex};

  my ($prompt,$saw_eof,$saw_timeout) = (0,0,0);

  my $timeout = defined $cmd ? $opt->{line_timeout} : $opt->{startup_timeout};

  $opt->{ex}->expect($timeout,
#		[ "\n" => sub {
#							#my $m = $opt->{ex}->exp_before().$opt->{ex}->match();
#							#print STDERR "{$m}\n";
#							just_saw($opt);
#					}],
	      [ $opt->{prompt}, sub { just_saw($opt); $prompt=1; } ],
#	      @{$opt->{watch_prompts}},
#	      [ 'eof', sub { just_saw($opt); $saw_eof=1; } ],
#	      [ 'timeout', sub { just_saw($opt); $saw_timeout=1; } ],
	      [ "\n", sub { just_saw($opt); exp_continue_timeout; } ],
	);

  if (defined($opt->{ex}->exp_exitstatus())) {
    $opt->{server_done} = 1;
  } elsif ($saw_timeout) {
    if ($opt->{saw_timeout}) {
      print STDERR "[$PROGNAME] Saw timeout [$timeout]\n";
      $opt->{server_done} = 1;
    } else {
      print STDERR "[$PROGNAME] Saw timeout [$timeout], attempting reset\n";
      $opt->{saw_timeout} = 1;
      reset_command_state($opt);
    }
  } else {
    $opt->{saw_timeout} = 0;
  }

  return 0 if $saw_eof || $opt->{server_done};
  return 1;
}

sub reset_command_state {
  my ($opt) = @_;

  foreach my $rc ( @{$opt->{reset_daemon}} ) {
    send_command($opt,$rc);
  }
}

##################################################
# Server code
##################################################
# Doesn't work with Expect??
sub daemonize {
  # Escape to the background
  chdir '/'                 or die "[$PROGNAME] Can't chdir to /: $!";
  umask 0;
  open STDIN, '/dev/null'   or die "[$PROGNAME] Can't read /dev/null: $!";
  open STDOUT, '>/dev/null' or die "[$PROGNAME] Can't write to /dev/null: $!";
  open STDERR, '>/dev/null' or die "[$PROGNAME] Can't write to /dev/null: $!";
  defined(my $pid = fork)   or die "[$PROGNAME] Can't fork: $!";
  exit if $pid;
  setsid                    or die "[$PROGNAME] Can't start a new session: $!";
}

my $CLIENT;
sub broken_pipe { close $CLIENT if $CLIENT; }
$SIG{'PIPE'} = 'broken_pipe';

sub tcp_server {
  my ($opt) = @_;

  my $server = IO::Socket::INET->new( Proto     => 'tcp',
                                      LocalPort => $opt->{port},
                                      Listen    => SOMAXCONN,
                                      Reuse     => 1);

  die "Can't setup server: $!\n" unless $server;
  print "[Server $0 accepting clients at $opt->{port}]\n" if $opt->{debug};

  while (my $client = $server->accept()) {
    $CLIENT = $client;	# Pipe
    $opt->{client} = $client;

    $client->autoflush(1);
    print $client "[$PROGNAME: connected to $opt->{daemon} server]\n" if $opt->{debug};
    my $hostinfo = gethostbyaddr($client->peeraddr);
    my $client_name = $hostinfo->name || $client->peerhost;
    print "[Connect from $client_name]\n" if $opt->{debug};

    print $client "$LAST_PROMPT" if $LAST_PROMPT;
    while ( my $cmd = <$client> ) {
      chomp($cmd);
      print "[$client_name] $cmd\n" if $opt->{debug};

      last if ($cmd =~ /^\Q$opt->{disconnect}\E$/);
      last unless send_command($opt,$cmd);
    }
    close $client;
    undef $CLIENT;
    print "[Disconnect from $client_name]\n" if $opt->{debug};
    last if $opt->{server_done};

    # Reset command state
    reset_command_state($opt);
  }
}

# Check if a pid is active
use Errno;
sub pid_active {
  my ($pid) = @_;
  return 0 unless $pid =~ /^\d+$/;
  return ((kill 0 => $pid) || $!{EPERM}) ? 1 : 0;
}

# Is a server really running?
sub server_up {
  my ($opt) = @_;

  return 0 unless (-f $opt->{pid_file});

  # Supposedly a server is running, give them a few seconds to reveal their pid
  my $attempts=0;
  my $pid;
  while ($attempts++<3) {
    if (open(PID,"<$opt->{pid_file}")) {
      $pid = scalar <PID>;  chomp $pid;
      close(PID);
      last;
    }
    sleep 1;
  }
  return $pid if pid_active($pid);
  unlink($opt->{pid_file});	# Okay - little bit of a race here..  :(
  print STDERR "[$PROGNAME] Server pid [$pid] has died...\n";
  return 0;
}

# Make sure a server is running - if not, start one
sub check_server {
  my ($opt) = @_;

  # We can only easily check localhost
  return if ($opt->{host} ne "localhost");

  return if (server_up($opt));

  print STDERR "[$PROGNAME] Starting server, please wait\n";
  $opt->{starting_server} = 1;

  # Start one up, same args
  defined(my $pid = fork)   or die("[$PROGNAME] Can't fork: $!");
  unless ($pid) {
    exec($0,"-start_server=1",@saved_argv);
  }
}

# Start one!
sub start_server {
  my ($opt) = @_;

  # Try to get the flock
  my $save_umask = umask 022;
  open(SEM,">>$opt->{pid_file}") || die("[$PROGNAME] Can't create pid file: $opt->{pid_file}\n");
  umask $save_umask;
  # If we can't get the flock, someone else has started up.
  unless (flock(SEM, LOCK_EX|LOCK_NB)) {
    my $pid = server_up($opt);
    die("[$PROGNAME] Can't flock $opt->{pid_file}, but no server is running\n") unless $pid;
    print STDERR "[$PROGNAME] Server already running at pid $pid\n";
    return;
  }
  # Is this a race?  I don't feel like creating a second semaphore just so I
  # can write this damn file.
  open(SEM,">$opt->{pid_file}") || die("[$PROGNAME] Can't create pid file: $opt->{pid_file}\n");
  die("[$PROGNAME] Lost my lock!  Oops!  Bad race condition!\n")
    unless (flock(SEM, LOCK_EX|LOCK_NB));
  select((select(SEM), $| = 1) [0]);	# Unbuffer
  seek(SEM,0,0);
  print SEM "$$\n";

  exit unless start_daemon($opt);
  tcp_server($opt);
  stop_daemon($opt);

  unlink($opt->{pid_file});
  close(SEM);
}

##################################################
# Client code
##################################################
sub tcp_client {
  my ($opt) = @_;

  # Are we already getting feedback, or should we show the dots?
  my $getting_feedback = $opt->{debug} && $opt->{starting_server} ? 1 : 0;

  my $handle;
  my $attempts;
  while (1) {
    $handle = IO::Socket::INET->new(Proto    => "tcp",
                                    PeerAddr => $opt->{host},
                                    PeerPort => $opt->{port});
    last if $handle;
    print STDERR "[Waiting for server to finish booting (is -prompt correct?)]\n" if !$attempts && !$getting_feedback;
    sleep 1;
    print STDERR "." unless $getting_feedback;
    die("[$PROGNAME] Cannot connect to $opt->{host}:$opt->{port}\n")
      if $attempts++ > $opt->{startup_timeout};
  }
  print STDERR "\n[Connected]\n" if $attempts && !$getting_feedback;

  $handle->autoflush(1);

  my $pid = fork;
  die("[$PROGNAME] Couldn't fork: $!\n") if (!defined $pid);

  if ($pid) {
    select((select(STDOUT), $|=1)[0]);  # Don't buffer stdout
    my $c;
    while(sysread($handle,$c,1)) { print $c; }
    print "\n";
    kill("TERM",$pid);
  } else {
    select((select(STDIN), $|=1)[0]);
    my $c;
    while(sysread(STDIN,$c,1)) { print $handle $c; }
    print $handle "$opt->{disconnect}\n" if $handle;
    exit;
  }

  close ($handle)            || die "close: $!";
}

##################################################
# Main code
##################################################
sub main {
  my $opt = parse_args();

  if ($opt->{stop_server}) {
    my $pid = server_up($opt);
    print STDERR "[$PROGNAME] No server running\n" unless $pid;
    exit unless $pid;
    kill HUP => $pid;
    die("[$PROGNAME] Can't stop server started by another user [$pid]\n")
      if ($!{EPERM});
    if (pid_active($pid)) {
      sleep 1;
      if (pid_active($pid)) {
        kill TERM => $pid;
        unlink($opt->{pid_file});
      }
    }
    print STDERR "[$PROGNAME] Killed server [$pid]\n";
    unlink($opt->{pid_file});
  } elsif ($opt->{start_server}) {
    start_server($opt);
  } else {
    check_server($opt);
    tcp_client($opt);
  }
}
main();