#!/usr/bin/perl -w
use strict;
use warnings;
# use diagnostics;
use bytes;
use Getopt::Long qw(:config gnu_getopt);
use IO::Socket::INET6;
use IO::Poll     qw(POLLRDNORM POLLWRNORM POLLIN POLLHUP);
use POSIX        qw(strftime);
use Time::HiRes  qw(time);

# maximum bytes read at one time:
my $MAXREAD = 1024;

my $usage = join("\n",
    "USAGE: $0 [options] inport host outport",
    '  accepts connections on port <inport> and forwards them to <outport> on',
    '  <host>, while also printing useful debugging stuff on stdout.',
    'OPTIONS:',
    '  --hex               : produce a hexdump instead of assuming mostly ascii',
    '  --nospin            : remove the silly -|-|- whirlygig',
    '  --time    | -t      : add @millisecond timestamps to the end of each line',
    '  --in      | -i file : save a copy of the client->server to this file',
    '  --out     | -o file : save a copy of the server->client to this file',
    '  --inject  | -I file : tail this file, also inject it towards the server',
    '  --outject | -O file : tail this file, also inject it towards the client',
    '  --quiet             : no stdout, ONLY write to --in and --out',
    'NOTES:',
    '  file can contain "#" to create a new file per TCP connection',
    '  file can be ">>file" to append.',
    '  file can be "|program" to pipe --in or --out to another program',
    '  file can be "program|" to pipe another program to --inject or --outject',
    '  --in and --out can be THE SAME FILE if you wish to smush them together',
    '');
my %opt = ( spin => 1);

GetOptions(\%opt, 'hex!', 'spin!', 'time|t!',  'quiet|q!',
    'in|i=s', 'out|o=s', 'inject|I=s', 'outject|O=s' )
  or die $usage;
my $inport  = shift || die $usage;
my $host    = shift || die $usage;
my $outport = shift || die $usage;
die "Extra args @ARGV? See --help"  if @ARGV;

if ($opt{in})  { $opt{in}  =~ s/^([^>|])/>$1/; }
if ($opt{out}) { $opt{out} =~ s/^([^>|])/>$1/; }

my %esc = (); prep_esc();

my $sock = IO::Socket::INET6->new(
    Listen    => 5,
    ReuseAddr => 1,
    LocalPort => $inport,
    Proto     => 'tcp',
) or die "### Can't bind to $inport: $!";

my $poll = new IO::Poll;
$poll -> mask($sock => POLLIN);

my $start = time();
print strftime "!start=%Y-%m-%dT%H:%M:%SZ\n", gmtime($start);
print "### Waiting for connections on port $inport\n";

my $no='00';
my %pairs = ();
my $spin = 0;

$|=1;
while (1) {
    if ($opt{spin}) {
	$spin = !$spin;
	if ($spin) { print STDERR "-\r"; }
	else       { print STDERR "|\r"; }
    }
    $poll->poll(1);
    for my $in ($poll->handles(POLLIN|POLLHUP)) {
	if ($in eq $sock) { newconn(); }
	else              { read_dat($in); }
    }
}

sub describe {
    my $sock = shift;
    my $ret = $sock->peerhost();
    $ret =~ s/.+:.+/[$&]/;
    return $ret . ':' . $sock->peerport();
}

sub newconn {
    $no++;
    my $in = $sock->accept();
    print "!C${no}_connect=", describe($in), timestamp(), "\n"; # ... to $host:$outport
    $poll -> mask($in, POLLIN);
    
    my $out = IO::Socket::INET6->new(
	PeerAddr => $host,
	PeerPort => $outport,
	Proto     => 'tcp',
    ) or die "### Can't connect to $host:$outport: $!\n";
    print "!S${no}_connect=", describe($out), timestamp(), "\n";
    $poll -> mask($out, POLLIN);
    
    my $ilog = undef;
    if (my $fn = $opt{in}) {
	$fn =~ s/\#/$no/g;
	if (open($ilog, $fn)) {
	    print "### C$no saving to $fn\n";
	} else {
	    print "### C$no FAILED TO WRITE $fn : $!", timestamp(), "\n";
	    $ilog = undef;
	}
    }
    
    my $olog = undef;
    if (my $fn = $opt{out}) {
	if ($fn eq $opt{in}) {
	    $olog = $ilog;
	} else {
	    $fn =~ s/\#/$no/g;
	    if (open($olog, $fn)) {
		print "### S$no saving to $fn\n";
	    } else {
		print "### S$no FAILED TO WRITE $fn : $!", timestamp(), "\n";
		$olog = undef;
	    }
	}
    }

    $pairs{$in}      = [ "C$no", $out, $ilog,   $in, $out, $ilog, $olog];
    $pairs{$out}     = [ "S$no", $in , $olog,   $in, $out, $ilog, $olog];

    my $inject = undef;
    if (my $fn = $opt{inject}) {
	$fn =~ s/\#/$no/g;
	my $pid;
	if ($fn =~ /\|$/) {
	    $pid = open($inject, $fn) or $inject = undef;
	} else {
	    open (OUT, ">$fn") && close(OUT);
	    $pid = open($inject, "tail -fc +0 $fn |") or $inject = undef;
	}
	if ($inject) {
	    print "### c$no accepting injections from $fn\n";
	    $poll -> mask($inject, POLLIN);
	    $pairs{$inject}  = [ "c$no", $out , $ilog,   $pid, $inject ];
	    push @{$pairs{$in}},  $pid, $inject;
	    push @{$pairs{$out}}, $pid, $inject;
	} else {
	    print "### c$no FAILED TO INJECT FROM $fn : $!", timestamp(), "\n";
	}
    }

    my $outject = undef;
    if (my $fn = $opt{outject}) {
	$fn =~ s/\#/$no/g;
	my $pid;
	if ($fn =~ /\|$/) {
	    $pid = open($outject, $fn) or $outject = undef;
	} else {
	    open (OUT, ">$fn") && close(OUT);
	    $pid = open($outject, "tail -fc +0 $fn |") or $outject = undef;
	}
	if ($outject) {
	    print "### s$no accepting outjections from $fn\n";
	    $poll -> mask($outject, POLLIN);
	    $pairs{$outject} = [ "s$no", $in  , $olog,   $pid, $outject ];
	    push @{$pairs{$in}},  $pid, $outject;
	    push @{$pairs{$out}}, $pid, $outject;
	} else {
	    print "### s$no FAILED TO OUTJECT FROM $fn : $!", timestamp(), "\n";
	}
    }
}

sub read_dat {
    my $in = shift;
    my $stuff = $pairs{$in} or return print "### LOST $in", timestamp(), "\n";
    my ($name, $out, $log, @closeme) = @$stuff;
    
    my $op;
    if ($in->sysread($op, $MAXREAD)) {
	$out->syswrite($op);
	$log->syswrite($op) if $log;
	return if $opt{quiet};
	$op =~ s/./$esc{$&}/gs;
	print "$name $op", timestamp(), "\n";
	return;
    }
    
    # If we are here, we have an error / file close to clean up
    if (lc($name) eq $name) {
	print "### $name out/injector died", timestamp(), "\n";
    } else {
	print "!${name}_disconnect", timestamp(), "\n";
    }
    for ($in, @closeme) {
	if (ref $_) {
	    $poll -> mask($_, 0);
	    delete $pairs{$_};
	    $_->close();
	} elsif (defined $_) {
	    kill 'TERM', $_;
	}
    }
}

sub prep_esc {
    if ($opt{'hex'}) {
	for my $n (0 .. 255) {
	    $esc{chr($n)} = sprintf("%02x", $n);
	}
	return;
    }
    %esc = (
	"\a" => "\\a",
	"\b" => "\\b",
	"\e" => "\\e",
	"\f" => "\\f",
	"\n" => "\\n",
	"\r" => "\\r",
	"\t" => "\\t",
    );
    my $n;
    ### naievely assume ascii...
    for my $n (  0 .. 31 ) { $esc{chr($n)} //= sprintf('\%03o', $n); };
    for my $n ( 32 .. 126) { $esc{chr($n)} //= chr($n); }
    for my $n (127 .. 255) { $esc{chr($n)} //= sprintf('\%03o', $n); };
}

sub timestamp {
    return '' unless $opt{'time'};
    return ' @' . int((time()-$start)*1000);
}

__END__