#!/usr/bin/perl -w

use strict;
use Socket qw( sockaddr_in inet_ntoa inet_aton
    PF_INET SOCK_STREAM SOCK_DGRAM INADDR_ANY SOL_SOCKET SO_REUSEADDR );

my $Self= $0;
$Self =~ s#^(.*[/\\])?([^/\\]+?)([.][^.]*)?$#$2#;
my $whenStart= time();
my $Verbose= 0;
my $Prompt= 0;
my $PromptDelay= 0;

exit main();

sub Usage
{
    warn @_, ".\n"   if  @_;
    warn "Usage:  $Self ",
     "[-uvp] [LocalAddr:]LocalPort[In,Out] SourceHost[:SourcePort] DestHost:DestPort\n";
    if(  @_  ) {
        die qq<Enter "$Self -?" for full usage information.\n>;
    } else {
        die <<END;
    Allows internet traffic to a local port to be bi-directionally forwarded
    to another destination.  Run $Self then have SourceHost form a connection
    from its SourcePort to ThisHost:LocalPort.  The connection will be
    forwarded to DestHost:DestPort and return traffic will be forwarded
    back to SourceHost:SourcePort.  A one-port network address translation.
    -u  Specifies UDP instead of TCP.
    -v  Verbose mode (show a time stamp and size for each packet).
    -p  Prompt for each packet as to whether or not to forward/return it.
    LocalAddr   Optional IP address to listen on.  Defaults to "any".
    LocalPort   Port name/number to use on this computer or 0.  You can also
                specify LocalPortIn,LocalPortOut to use a different local port
                number for receiving and sending.  Either can be 0 to pick a
                free port and LocalPortOut can be -s to use the same port
                number as SourcePort.
    SourceHost  Host name or IP address of computer forming connection.
    SourcePort  Port name/number on SourceHost where connection comes from.
                If SourcePort is not given, then it is extracted from the
                next connection made from SourceHost to ThisHost:LocalPort.
    DestHost    Host name or IP address of computer to recieve connection.
    DestPort    Port name/number on DestHost to forward connection to.
END
    }
}

sub colons
{
    my( $secs )= @_;
    if(  $secs < 60  ) {
        return sprintf "0:%02d", $secs;
    }
    my( $mins )= $secs/60;
    $secs %= 60;
    if(  $mins < 60  ) {
        return sprintf "%d:%02d", $mins, $secs;
    }
    my( $hrs )= $mins/60;
    $mins %= 60;
    if(  $mins < 24  ) {
        return sprintf "%d:%02d:%02d", $hrs, $mins, $secs;
    }
    my( $days )= $hrs/60;
    $hrs %= 24;
    return sprintf "%d_%02d:%02d:%02d", $days, $hrs, $mins, $secs;
}

sub commas
{
    my( $num )= @_;
    1   while  $num =~ s#([._])(\d\d\d)(\d)#$1$2_$3#;
    1   while  $num =~ s#(\d)(\d\d\d)(\D|$)#$1,$2$3#;
    return $num;
}

{
    my $nl;
    BEGIN { $nl= "\n"; }
    sub nl
    {
        if(  @_  ) {
            return $nl= shift;
        } else {
            my $old= $nl;
            $nl= "";
            return $old.colons(time()-$whenStart)."- ";
        }
    }
}

{
    my $defaultAnswer;
    BEGIN { $defaultAnswer= 'y' };

    sub shouldDropPacket
    {
        my( $dir, $buf )= @_;
        my $time= join ":", map sprintf("%02d",$_), (localtime())[2,1,0];
        if(  ! $Prompt  ) {
            print nl,"$time- ${dir}ing ",length($buf)," bytes.\n";
        } else {
            print nl,"$time- $dir ",length($buf)," bytes? ";
            my $resp= ReadKey( $PromptDelay );
            $resp= $defaultAnswer   if  ! defined $resp;
            if(  $resp =~ /^\s*n/i
             ||  $resp =~ /^\s*f/i  &&  $dir !~ /^f/i
             ||  $resp =~ /^\s*r/i  &&  $dir !~ /^r/i
            ) {
                $resp= "Dropped";
                $buf= "";
                $defaultAnswer= 'n';
            } else {
                $resp= "${dir}ed";
                $defaultAnswer= 'y';
            }
            print " $resp\n";
        }
        return $buf;
    }
}

# ipm [-u] [-p|801] src[:sp] dest:800
# src:[*|sp]-*:801 LISTEN
# src:1031-hub:801 connect
# src:1031-hub:801 => hub:801-dest:800
# dest:800-hub:801 => hub:801-src:1031

sub main
{
    $|= 1;
    # Kludge to turn C<$!> of C<"Unknown Error"> into something useful:
    if(  $^O =~ /^MSWin/  ) {
        $SIG{__DIE__}= sub {
            my( $msg )= @_;
            $msg =~ s#([:\s]\s*)(Unknown Error)((:? at line.*)?)\s*$#
                      $1."errno=".(0+$!).$3."\n"#e;
            die $msg;
        };
    }

    # Parse command line:
    my $proto= "tcp";
    my $quality= SOCK_STREAM;
    while(  @ARGV  &&  $ARGV[0] =~ /^-[^-]/  ) {
        my $arg= substr( $ARGV[0], 1 );
        while(  "" ne $arg  ) {
            if(  $arg =~ s/^u//i  ) {
                $proto= "udp";
                $quality= SOCK_DGRAM;
            } elsif(  $arg =~ s/^v//i  ) {
                $Verbose= ! $Verbose;
            } elsif(  $arg =~ s/^p[:=]?([\d.]*)//i  ) {
                $Prompt= ! $Prompt;
                $PromptDelay= "" eq $1 ? 0 : $1 ? $1 : -1;
            } elsif(  $arg =~ /^\?/  ) {
                Usage;
            } else {
                Usage "Unknown switch (-$arg of $ARGV[0])";
            }
        }
        shift @ARGV;
    }
    if(  $Prompt  ) {
        require Term::ReadKey;
        Term::ReadKey->import( "ReadKey" );
        Term::ReadKey::ReadMode( 3 );
    }
    if(  3 != @ARGV  ) {
        Usage "Wrong number of arguments";
    }
    my( $portLocalIn, $specSrc, $specDest )= @ARGV;

    # Validate data from command line:
    my $addrLocalIn;
    if(  $portLocalIn =~ /:/  ) {
        ( $addrLocalIn, $portLocalIn )= split /:/, $portLocalIn;
    }
    my $portLocalOut= $portLocalIn;
    if(  $portLocalIn =~ /,/  ) {
        ( $portLocalIn, $portLocalOut )= split /,/, $portLocalIn;
    }
    if(  $portLocalIn =~ /^\d+$/  ||  $portLocalIn =~ /^0x[\da-f]+$/i  ) {
        $portLocalIn=  $portLocalIn =~ /^0/
                         ?  oct($portLocalIn)  :  0+$portLocalIn;
    } else {
        my $nPort= getservbyname( $portLocalIn, $proto )
          or  die "Invalid local port name/number ($portLocalIn): $!\n";
        $portLocalIn= $nPort;
    }
    if(  $portLocalOut =~ /^\d+$/  ||  $portLocalOut =~ /^0x[\da-f]+$/i  ) {
        $portLocalOut=  $portLocalOut =~ /^0/
                          ?  oct($portLocalOut)  :  0+$portLocalOut;
    } elsif(  $portLocalOut =~ /^-[1i]$/i  ) {
        $portLocalOut= -1;
    } else {
        my $nPort= getservbyname( $portLocalOut, $proto )
          or  die "Invalid local out port name/number ($portLocalOut): $!\n";
        $portLocalOut= $nPort;
    }
    my $ipLocalIn= defined($addrLocalIn) ? inet_aton($addrLocalIn) : INADDR_ANY;
    my $inLocal= sockaddr_in( $portLocalIn, $ipLocalIn )
      or  die "Can't form local address struct: $!\n";

    my( $hostSrc, $portSrc )=  split /:/, $specSrc;
    $portSrc= 0   if  ! defined($portSrc)  ||  "" eq $portSrc;
    my $ipSrc= gethostbyname( $hostSrc )
      or  die "Invalid source host name/addr ($hostSrc): $!\n";
    if(  $portSrc =~ /^\d+$/  ||  $portSrc =~ /^0x[\da-f]+$/i  ) {
        $portSrc=  $portSrc =~ /^0/  ?  oct($portSrc)  :  0+$portSrc;
    } else {
        my $nPort= getservbyname( $portSrc, $proto )
          or  die "Invalid source port name/number ($portSrc): $!\n";
        $portSrc= $nPort;
    }
    my $inSrc= sockaddr_in( $portSrc, $ipSrc )
      or  die "Can't form source address struct: $!\n";

    my( $hostDest, $portDest )= split /:/, $specDest, 2;
    my $ipDest= gethostbyname( $hostDest )
      or  die "Invalid destination host name/addr ($hostDest): $!\n";
    if(  ! defined $portDest  ) {
        Usage qq<Required DestPort is missing; no colon in "$specDest">;
    }
    if(  $portDest =~ /^\d+$/  ||  $portDest =~ /^0x[\da-f]+$/i  ) {
        $portDest=  $portDest =~ /^0/  ?  oct($portDest)  :  0+$portDest;
    } else {
        my $nPort= getservbyname( $portDest, $proto )
          or  die "Invalid destination port name/number ($portDest): $!\n";
        $portDest= $nPort;
    }
    die "DestPort must not be 0"   if  0 == $portDest;
    my $inDest= sockaddr_in( $portDest, $ipDest )
      or  die "Can't form destination address struct; $!\n";

    # Listen for incoming connection:
    my $nProto= getprotobyname($proto)
      or  die qq<Can't find network protocol named "$proto": $!\n>;
    local( *HEAR );
    socket( HEAR, PF_INET, $quality, $nProto )
      or  die qq<Can't create PF_INET,$quality,"$proto" socket: $!\n>;
    bind( HEAR, $inLocal )
      or  die "Can't bind local port number $portLocalIn: $!\n";
    if(  0 == $portLocalIn  ) {
        $inLocal= getsockname( HEAR )
          or  die "Can't get local port number: $!\n";
        my $ipLocal;
        ( $portLocalIn, $ipLocal )= sockaddr_in( $inLocal )
          or  die "Can't decode local port number: $!\n";
    }
    local( *SRC );
    my $buf= "";
    my $lenRead= 512*512;
    if(  "tcp" eq $proto  ) {

        listen( HEAR, 5 )
          or  die "Can't listen for conections: $!\n";
        warn nl,"Listening on port $portLocalIn for ",
          "connection from $hostSrc:$portSrc...\n";
        while(  1  ) {
            $inSrc= accept( SRC, HEAR )
              or  die "Can't failed accpeting a connection: $!\n";
            my( $portTry, $ipTry )= sockaddr_in( $inSrc )
              or  die "Can't decode incoming port/address: $!\n";
            my $addrTry = inet_ntoa($ipTry);
            # my $hostTry = gethostbyaddr( $ipTry, AF_INET );
            if(  $ipTry ne $ipSrc  ) {
                warn nl,"Refusing connection from $addrTry:$portTry.\n";
            } elsif(  0 != $portSrc  &&  $portTry != $portSrc  ) {
                warn nl,"Wrong source port; ",
                  "refusing connection from $addrTry:$portTry.\n";
            } else {
                $portSrc= $portTry   if  0 == $portSrc;
                last;
            }
            close( SRC );
        }
        close( HEAR );
    } elsif(  0 == $portSrc  ) {
        warn nl,"Listening on port $portLocalIn for ",
          "first packet from $hostSrc:$portSrc...\n";
        while(  1  ) {
            my $inTry= recv( HEAR, $buf, $lenRead, 0 )
              or  die "Can't receive first packet: $!\n";
            my( $portTry, $ipTry )= sockaddr_in( $inTry )
              or  die "Can't decode first packet's port/address: $!\n";
            my $addrTry = inet_ntoa($ipTry);
            if(  $ipTry ne $ipSrc  ) {
                warn nl,"Ignoring packet from $addrTry:$portTry.\n";
            } else {
                $portSrc= $portTry;
                last;
            }
        }
        socket( SRC, PF_INET, $quality, $nProto )
          or  die "Can't create new socket: $!\n";
        close( HEAR );
        # The above C<close> means we don't need to set SO_REUSEADDR but
        # might introduce a race condition where we could lose packets
        # between above the C<close> and following the C<bind>/C<connect>.
        # However, delaying the C<close> doesn't really solve the race
        # condition because C<HEAR> could receive the next packet instead
        # of the new C<SRC>.
        #setsockopt( SRC, SOL_SOCKET, SO_REUSEADDR, 1 )
        #  or  die "Can't turn on SO_REUSEADDR for new socket: $!\n";
        bind( SRC, $inLocal )
          or  die "Can't bind new socket to port $portLocalIn: $!\n";
        $inSrc= sockaddr_in( $portSrc, $ipSrc )
          or  die "Can't form new source address struct: $!\n";
        connect( SRC, $inSrc )
          or  die "Can't connect new socket to $hostSrc:$portSrc: $!\n";
    } else {
        connect( HEAR, $inSrc )
          or  die "Can't connect socket to $hostSrc:$portSrc: $!\n";
        warn nl,"Listening on port $portLocalIn for ",
          "packets from $hostSrc:$portSrc...\n";
        *SRC= *HEAR;
        {   local( *UNDEF );
            *HEAR= *UNDEF;
        }
    }

    # Connect to destination:
    $portLocalOut= $portSrc   if  -1 == $portLocalOut;
    warn nl,"Forwarding connection from $hostSrc:$portSrc ",
      "to $hostDest:$portDest...\n";
    if(  $Verbose  ||  $Prompt  and  "" ne $buf  ) {
        $buf= shouldDropPacket( "Forward", $buf );
    }
    local( *DEST );
    socket( DEST, PF_INET, $quality, $nProto )
      or  die qq<Can't create PF_INET,$quality,"$proto" socket: $!\n>;
    if(  0 != $portLocalOut  ) {
        $inLocal= sockaddr_in( $portLocalOut, INADDR_ANY )
          or  die "Can't form local out address struct: $!\n";
        # Since we are about to C<connect>, it doesn't
        # really mattter if we reuse a local port number:
        setsockopt( DEST, SOL_SOCKET, SO_REUSEADDR, 1 )
          or  die "Can't turn on SO_REUSEADDR for out socket: $!\n";
        bind( DEST, $inLocal )
          or  die "Can't bind local out port number $portLocalOut: $!\n";
    }
    connect( DEST, $inDest )
      or  die "Can't connect to $hostDest:$portDest: $!\n";

    # Pump data between end points:
    my $fdsRead= '';
    vec( $fdsRead, fileno(SRC), 1 )= 1;
    vec( $fdsRead, fileno(DEST), 1 )= 1;
    my( $fdsInput, $cntFds );
    my $whenLog= time()-$whenStart;
    $whenLog += $whenStart+10 - $whenLog%10;
    my( $packsIn, $packsOut, $bytesIn, $bytesOut )= ( 0, 0, 0, 0 );
    if(  "" ne $buf  ) {
        $packsIn++;   $bytesIn += length($buf);
        my $wrote= 0;
        while(  $wrote < length($buf)  ) {
            my $ret=
              syswrite( DEST, $buf, length($buf)-$wrote, $wrote )
              or  die "Can't write ", length($buf)-$wrote,
                "-byte fragment of ", length($buf)," -byte packet.\n";
            if(  $wrote+$ret < length($buf)  ) {
                warn nl,"DEST only took ", $ret, " bytes of ",
                  length($buf)-$wrote, "-byte fragment of ",
                  length($buf), "-byte packet.\n";
            }
            $wrote += $ret;
        }
    }
    while(  1  ) {
        $cntFds= select( $fdsInput= $fdsRead, undef, undef, $whenLog-time() );
        if(  ! $cntFds  ) {
            $whenLog += 10   while  $whenLog <= time();
            if(  ! $Verbose  &&  ! $Prompt  ) {
                nl("\r");
                print nl,"Forwarded ",commas($bytesIn)," bytes in ",
                  commas($packsIn)," packets; returned ",commas($bytesOut)," in ",
                  commas($packsOut),".";
                nl("\n");
            }
            next;
        }
        if(  vec( $fdsInput, fileno(SRC), 1 )  ) {
            if(  sysread( SRC, $buf, $lenRead )  ) {
                $packsIn++;   $bytesIn += length($buf);
                if(  $Verbose  ||  $Prompt  ) {
                    $buf= shouldDropPacket( "Forward", $buf );
                }
                my $wrote= 0;
                while(  $wrote < length($buf)  ) {
                    my $ret=
                      syswrite( DEST, $buf, length($buf)-$wrote, $wrote )
                      or  die "Can't write ", length($buf)-$wrote,
                        "-byte fragment of ", length($buf)," -byte packet.\n";
                    if(  $wrote+$ret < length($buf)  ) {
                        warn nl,"DEST only took ", $ret, " bytes of ",
                          length($buf)-$wrote, "-byte fragment of ",
                          length($buf), "-byte packet.\n";
                    }
                    $wrote += $ret;
                }
            } elsif(  ! vec( $fdsRead, fileno(DEST), 1 )  ) {
                close( SRC )
                  or  warn nl,"Couldn't close SRC: $!\n";
                close( DEST )
                  or  warn nl,"Couldn't close DEST: $!\n";
                last;
            } elsif(  "udp" eq $proto  ) {
                #warn nl,"Ignoring EOF from SRC (",0+$!,").\n";
            } else {
                warn nl,"Read EOF from SRC (",0+$!,").\n";
                vec( $fdsRead, fileno(SRC), 1 )= 0;
                shutdown( SRC, 0 )
                  or  warn nl,"Can't shutdown read from SRC: $!\n";
                shutdown( DEST, 1 )
                  or  warn nl,"Can't shutdown writes to DEST: $!\n";
            }
        }
        if(  vec( $fdsInput, fileno(DEST), 1 )  ) {
            if(  sysread( DEST, $buf, 512*512 )  ) {
                $packsOut++;   $bytesOut += length($buf);
                if(  $Verbose  ||  $Prompt  ) {
                    $buf= shouldDropPacket( "Return", $buf );
                }
                my $wrote= 0;
                while(  $wrote < length($buf)  ) {
                    my $ret=
                      syswrite( SRC, $buf, length($buf)-$wrote, $wrote )
                      or  die "Can't write ", length($buf)-$wrote,
                        "-byte fragment of ", length($buf)," -byte packet.\n";
                    if(  $wrote+$ret < length($buf)  ) {
                        warn nl,"SRC only took ", $ret, " bytes of ",
                          length($buf)-$wrote, "-byte fragment of ",
                          length($buf), "-byte packet.\n";
                    }
                    $wrote += $ret;
                }
            } elsif(  ! vec( $fdsRead, fileno(SRC), 1 )  ) {
                close( DEST )
                  or  warn nl,"Couldn't close DEST: $!\n";
                close( SRC )
                  or  warn nl,"Couldn't close SRC: $!\n";
                last;
            } else {
                warn nl,"Read EOF from DEST.\n";
                vec( $fdsRead, fileno(DEST), 1 )= 0;
                shutdown( DEST, 0 )
                  or  warn nl,"Can't shutdown read from DEST: $!\n";
                shutdown( SRC, 1 )
                  or  warn nl,"Can't shutdown writes to SRC: $!\n";
            }
        }
    }
    warn nl,"Both sides of connection successfully shut down.\n";
    print nl,"Forwarded ",commas($bytesIn),
      " bytes in ",commas($packsIn)," packets; returned ",commas($bytesOut),
      " in ",commas($packsOut),".\n";

    return 0;
}